package gradientDescent;

import com.naaturel.ANN.domain.abstraction.AlgorithmStep;
import com.naaturel.ANN.domain.abstraction.Model;
import com.naaturel.ANN.domain.model.helpers.ModelCreator;
import com.naaturel.ANN.infrastructure.dataset.DataSet;
import com.naaturel.ANN.infrastructure.dataset.DatasetExtractor;
import com.naaturel.ANN.domain.model.neuron.*;
import com.naaturel.ANN.domain.model.training.TrainingPipeline;
import com.naaturel.ANN.implementation.gradientDescent.*;
import com.naaturel.ANN.implementation.simplePerceptron.*;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.util.ArrayList;
import java.util.List;

import static org.junit.jupiter.api.Assertions.*;


public class GradientDescentTest {

    private DataSet dataset;
    private GradientDescentTrainingContext context;

    private Model model;

    private TrainingPipeline pipeline;

    @BeforeEach
    public void init(){
        dataset = new DatasetExtractor()
                .extract("C:/Users/Laurent/Desktop/ANN-framework/src/main/resources/assets/and-gradient.csv", 1);

        model = ModelCreator.createModel(new int[]{1}, 2, new Linear(1, 0));

        context = new GradientDescentTrainingContext(model, dataset);
        context.correctorTerms = new ArrayList<>();

        List<AlgorithmStep> steps = List.of(
                new SimplePredictionStep(context),
                new SimpleDeltaStep(context),
                new SquareLossStep(context),
                new GradientDescentErrorStrategy(context)
        );

        pipeline = new TrainingPipeline(steps)
                .stopCondition(ctx -> ctx.globalLoss <= 0.125F || ctx.epoch > 100)
                .beforeEpoch(ctx -> {
                    GradientDescentTrainingContext gdCtx = (GradientDescentTrainingContext) ctx;
                    gdCtx.globalLoss = 0.0F;
                    gdCtx.correctorTerms.clear();
                    for (int i = 0; i < ctx.model.synCount(); i++){
                        gdCtx.correctorTerms.add(0F);
                    }
                });
    }

    @Test
    public void test_the_whole_algorithm(){

        List<Float> expectedGlobalLosses = List.of(
                0.5F,
                0.38F,
                0.3176F,
                0.272096F,
                0.237469F
        );

        context.learningRate = 0.2F;
        pipeline.afterEpoch(ctx -> {
            context.globalLoss /= context.dataset.size();
            new GradientDescentCorrectionStrategy(context).run();

            int index = ctx.epoch-1;
            if(index >= expectedGlobalLosses.size()) return;

            assertEquals(expectedGlobalLosses.get(index), context.globalLoss, 0.00001f);
        });

        pipeline
                .withVerbose(true, 1)
                .run(context);
        assertEquals(67, context.epoch);
    }
}
