package com.naaturel.ANN.domain.model.training;

import com.naaturel.ANN.domain.abstraction.AlgorithmStep;
import com.naaturel.ANN.domain.abstraction.TrainingContext;
import com.naaturel.ANN.infrastructure.dataset.DataSetEntry;
import com.naaturel.ANN.infrastructure.visualization.GraphVisualizer;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.function.Consumer;
import java.util.function.Predicate;

public class TrainingPipeline {

    private final List<AlgorithmStep> steps;
    private Consumer<TrainingContext> beforeEpoch;
    private Consumer<TrainingContext> afterEpoch;
    private Predicate<TrainingContext> stopCondition;

    private boolean verbose;
    private boolean visualization;
    private boolean timeMeasurement;

    private GraphVisualizer visualizer;
    private int verboseDelay;

    public TrainingPipeline(List<AlgorithmStep> steps) {
        this.steps = new ArrayList<>(steps);
        this.stopCondition = (ctx) -> false;
        this.beforeEpoch = (context -> {});
        this.afterEpoch = (context -> {});
    }

    public TrainingPipeline stopCondition(Predicate<TrainingContext> predicate) {
        this.stopCondition = predicate;
        return this;
    }

    public TrainingPipeline beforeEpoch(Consumer<TrainingContext> consumer) {
        this.beforeEpoch = consumer;
        return this;
    }

    public TrainingPipeline afterEpoch(Consumer<TrainingContext> consumer) {
        this.afterEpoch = consumer;
        return this;
    }

    public TrainingPipeline withVerbose(boolean enabled, int epochDelay) {
        if(epochDelay <= 0) throw new IllegalArgumentException("Epoch delay cannot lower or equal to 0");
        this.verbose = enabled;
        this.verboseDelay = epochDelay;
        return this;
    }

    public TrainingPipeline withVisualization(boolean enabled, GraphVisualizer visualizer) {
        this.visualization = enabled;
        this.visualizer = visualizer;
        return this;
    }

    public TrainingPipeline withTimeMeasurement(boolean enabled) {
        this.timeMeasurement = enabled;
        return this;
    }

    public void run(TrainingContext ctx) {

        long start = this.timeMeasurement ? System.currentTimeMillis() : 0;

        do {
            this.beforeEpoch.accept(ctx);
            this.executeSteps(ctx);
            this.afterEpoch.accept(ctx);
            if(this.verbose && ctx.epoch % this.verboseDelay == 0) {
                System.out.printf("[Global error] : %f\n", ctx.globalLoss);
            }
            ctx.epoch += 1;
        } while (!this.stopCondition.test(ctx));

        if(this.timeMeasurement) {
            long end = System.currentTimeMillis();
            System.out.printf("[Training finished in %.3fs]\n", (end-start)/1000.0);
        }
        System.out.printf("[Final global error] : %f\n", ctx.globalLoss);
        System.out.println("Accuracy: " +  accuracyCheck(ctx) + "%");
    }

    private void executeSteps(TrainingContext ctx){
        for (DataSetEntry entry : ctx.dataset) {

            ctx.currentEntry = entry;
            ctx.expectations = ctx.dataset.getLabelsAsFloat(entry);

            for (AlgorithmStep step : steps) {
                step.run();
            }

            if(this.verbose && ctx.epoch % this.verboseDelay == 0) {
                System.out.printf("Epoch : %d, ", ctx.epoch);
                System.out.printf("predicted : %s, ", Arrays.toString(ctx.predictions));
                System.out.printf("expected : %s, ", Arrays.toString(ctx.expectations.toArray()));
                System.out.printf("delta : %s, ", Arrays.toString(ctx.deltas));
                System.out.printf("loss : %.5f\n", ctx.localLoss);
            }
        }
    }

    private float accuracyCheck(TrainingContext ctx){
        int correct = 0;
        for (DataSetEntry entry : ctx.dataset) {
            float[] predictions = ctx.model.predict(entry.getDataAsFloat());
            List<Float> labels = ctx.dataset.getLabelsAsFloat(entry);

            // argmax of predictions
            int predictedClass = 0;
            for (int i = 1; i < predictions.length; i++) {
                if (predictions[i] > predictions[predictedClass]) predictedClass = i;
            }

            // argmax of labels
            int expectedClass = 0;
            for (int i = 1; i < labels.size(); i++) {
                if (labels.get(i) > labels.get(expectedClass)) expectedClass = i;
            }

            if (predictedClass == expectedClass) correct++;
        }

        return 100f * correct / ctx.dataset.size();
    }

}
