package cz.fidentis.analyst.visitors.face;

import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.function.ToDoubleFunction;
import java.util.stream.Collector;

/**
 * A collector for calculation of the weighted average, suitable for use in Java 8 streams.
 *
 * <p>
 * A mutable reduction operation that accumulates input elements into a mutable
 * result container, transforming the accumulated result into a final
 * representation after all input elements have been processed.
 * </p>
 * 
 * @author Daniel Schramm
 * @param <T> Data type of the stream elements
 */
public class WeightedAverageCollector<T> implements Collector<T, IntemediateResults, Double> {

    private final ToDoubleFunction<? super T> valueFunction;
    private final ToDoubleFunction<? super T> weightFunction;
    
    /**
     * Constructor.
     * 
     * @param valueFunction Function returning the value for a given stream element
     * @param weightFunction Function returning the weight for a given stream element
     */
    public WeightedAverageCollector(
            ToDoubleFunction<? super T> valueFunction,
            ToDoubleFunction<? super T> weightFunction) {
        this.valueFunction = valueFunction;
        this.weightFunction = weightFunction;
    }
    
    /**
     * Returns a {@link Collector} interface object used to compute the weighted average.
     * 
     * @param <T> Data type of the stream elements
     * @param valueFunction Function returning the value for a given stream element
     * @param weightFunction Function returning the weight for a given stream element
     * @return A Collector interface object used to compute the weighted average
     */
    public static <T> Collector<T, ?, Double> toWeightedAverage(
            ToDoubleFunction<? super T> valueFunction,
            ToDoubleFunction<? super T> weightFunction) {
        return new WeightedAverageCollector<>(valueFunction, weightFunction);
    }
    
    /**
     * A function that creates and returns a new mutable result container.
     *
     * @return A function which returns a new, mutable result container
     */
    @Override
    public Supplier<IntemediateResults> supplier() {
        return IntemediateResults::new;
    }
    
    /**
     * A function that folds a value into a mutable result container.
     *
     * @return A function which folds a value into a mutable result container
     */
    @Override
    public BiConsumer<IntemediateResults, T> accumulator() {
        return (iResult, streamElement) -> {
            iResult.addWeightedValSum(valueFunction.applyAsDouble(streamElement)
                    * weightFunction.applyAsDouble(streamElement));
            iResult.addWeightSum(weightFunction.applyAsDouble(streamElement));
        };
    }
    
    /**
     * A function that accepts two partial results and merges them. The
     * combiner function may fold state from one argument into the other and
     * return that, or may return a new result container.
     *
     * @return A function which combines two partial results into a combined
     *         result
     */
    @Override
    public BinaryOperator<IntemediateResults> combiner() {
        return (iResult1, iResult2) -> {
            iResult1.addWeightedValSum(iResult2.getWeightedValSum());
            iResult1.addWeightSum(iResult2.getWeightSum());
            
            return iResult1;
        };
    }
    
    /**
     * Perform the final transformation from the intermediate accumulation type
     * {@link IntemediateResults} to the final result type {@link Double}.
     *
     * <p>
     * If the characteristic {@code IDENTITY_FINISH} is
     * set, this function may be presumed to be an identity transform with an
     * unchecked cast from {@link IntemediateResults} to {@link Double}.
     * </p>
     *
     * @return A function which transforms the intermediate result to the final
     *         result
     */
    @Override
    public Function<IntemediateResults, Double> finisher() {
        return iResult -> iResult.getWeightedValSum() / iResult.getWeightSum();
    }
    
    /**
     * Returns a {@code Set} of {@code Collector.Characteristics} indicating
     * the characteristics of this Collector.
     *
     * @return An immutable set of collector characteristics
     */
    @Override
    public Set<Characteristics> characteristics() {
        return Set.of(Characteristics.UNORDERED);
    }
}

/**
 * A helper class which stores intermediate results
 * for the {@link WeightedAverageCollector} class.
 * 
 * @author Daniel Schramm
 */
class IntemediateResults {
    private double weightedValSum = 0;
    private double weightSum = 0;

    public double getWeightedValSum() {
        return weightedValSum;
    }

    public double getWeightSum() {
        return weightSum;
    }

    public void addWeightedValSum(double weightedValSum) {
        this.weightedValSum += weightedValSum;
    }

    public void addWeightSum(double weightSum) {
        this.weightSum += weightSum;
    }
}
