package com.naaturel.ANN.implementation.training;

import com.naaturel.ANN.domain.abstraction.AlgorithmStep;
import com.naaturel.ANN.domain.abstraction.Model;
import com.naaturel.ANN.domain.abstraction.Trainer;
import com.naaturel.ANN.domain.model.helpers.ModelCreator;
import com.naaturel.ANN.domain.model.training.TrainingPipeline;
import com.naaturel.ANN.implementation.gradientDescent.SquareLossStep;
import com.naaturel.ANN.implementation.multiLayers.*;
import com.naaturel.ANN.implementation.simplePerceptron.SimplePredictionStep;
import com.naaturel.ANN.infrastructure.dataset.DataSet;
import com.naaturel.ANN.infrastructure.persistence.ModelSnapshot;

import java.util.List;

public class GradientBackpropagationTraining implements Trainer {

    private Model model;

    public GradientBackpropagationTraining(int[] neurons, int nbrInputs){
        model = ModelCreator.createModel(neurons, nbrInputs, new TanH());
    }

    public Model getModel() {
        return model;
    }

    @Override
    public void train(float learningRate, int epoch, DataSet dataset) {
        GradientBackpropagationContext context =
                new GradientBackpropagationContext(model, dataset, learningRate, 10);

        List<AlgorithmStep> steps = List.of(
                new SimplePredictionStep(context),
                new OutputLayerErrorStep(context),
                new ErrorSignalStep(context),
                new BackpropagationCorrectionStep(context),
                new SquareLossStep(context)
        );

        new TrainingPipeline(steps)
                .stopCondition(ctx -> ctx.globalLoss <= 0.00F || ctx.epoch > epoch)
                .beforeEpoch(ctx -> {
                    ctx.globalLoss = 0.0F;
                })
                .afterEpoch(ctx -> {
                    ctx.globalLoss /= dataset.size();
                })
                .withVerbose(true,epoch/10)
                .withTimeMeasurement(true)
                .run(context);
    }

    @Override
    public void saveModel(ModelSnapshot snapshot, String path) {
        try {
            snapshot.saveToFile(model, path);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
}
