package fr.perceptron;

import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.Serializable;
import java.util.Random;
// import java.util.concurrent.ExecutorService;
// import java.util.concurrent.Executors;
// import java.util.concurrent.Future;

public class NeuralNetwork implements Serializable {
    
    public int[] sizes; // number of neurons for each layer of the neural network
    public Layer[] layers;

    private transient Random random;
    private static final double initialLearningRate = 0.1d;
    private static final double decayRate = 0.95d;
    private static final double decayFactor = 10.0d;
    public static final double COST_TOLERANCE = 10e-5;
    public static final double COST_VARIATION_TOLERANCE = 10e-7;
    public int epoch;
    public double learningRate;

    public NeuralNetwork(int[] layerSizes) {
        this.sizes = layerSizes;
        // the "input layer" doesn't truly exist so we don't add it
        this.layers = new Layer[layerSizes.length - 1];
        // init to a fixed seed to debug
        this.random = new Random(0);

        // set the epoch to 0 and the learning rate to the initial value
        this.epoch = 0;
        this.learningRate = initialLearningRate;

        // create each layer with the number of input neurons corresponding
        // to the number of output neurons of the previous index
        for (int i = 0; i < layerSizes.length - 1; ++i) {
            layers[i] = new Layer(layerSizes[i], layerSizes[i + 1], this.random);
        }
    }

    // calculates the output of the Neural Network based on given input values
    // the inputs are forwarded through each layer of the NN, using the forward()
    // method of each layer.
    public double[] calculateOutputs(double[] inputs) {
        double[] layerOutput = inputs;

        // forward the inputs through each layer
        for (int i = 0; i < this.layers.length; ++i) {
            layerOutput = this.layers[i].forward(layerOutput);
        }

        return layerOutput;
    }

    // returns an int that corresponds to the index of the highest activation
    // value of an output. (i.e. the index of the max of all values of the array)
    public int classify(double[] networkOutput) {
        int maxActivationIndex = 0;
        double maxActivation = Double.MIN_VALUE;

        for (int i = 0; i < networkOutput.length; ++i) {
            if (networkOutput[i] > maxActivation) {
                maxActivation = networkOutput[i];
                maxActivationIndex = i;
            }
        }

        return maxActivationIndex;
    }

    public int classify(double input) {
        // return 1 if input > 0 else 0
        return (input > 0.5 ? 1 : 0);
    }

    // calculates the cost of an evaluation by comparing the values predicted
    // with the expected output. The formula used for the cost is the
    // mean square error, defined as :
    // E(w) = 1/2 * sum( (expected - predicted) ^ 2)
    public double calculateCost(double[] expected, double[] outputs) {

        double cost = 0;
        for (int i = 0; i < outputs.length; ++i) {
            cost += Math.pow(expected[i] - outputs[i], 2);
        }
        return 0.5 * cost;
    }

    public void train(DataSet data) {
        // Train the Neural Network using the given data set
        // the data set is split into batches, and each batch is used to
        // update the weights and biases of the Neural Network
        // This implementation is sequential, but it could be parallelized (TODO)
        DataSet currentSet = data.generateBatch();

        while(currentSet != null) {
            System.out.print(String.format("Batch %2d/%d\r", data.currentBatch, data.numBatches));
            // for each point in the batch, calculate the outputs and update
            // the gradients of the Neural Network
            for (DataPoint currentPoint : currentSet.points) {
                double[] values = this.calculateOutputs(currentPoint.inputs);
                this.updateGradients(currentPoint, values);
            }
            currentSet = data.generateBatch();
        }
        System.out.println();
        // update the epoch and learning rate
        this.epoch++;
        this.learningRate = initialLearningRate
            * Math.pow(decayRate, this.epoch / decayFactor);
    }

    public void updateGradients(DataPoint point, double[] values) {
        // update the gradients of the Neural Network
        int finalLayer = this.layers.length - 1;
        double[][] layerErrors = new double[finalLayer + 1][];
        layerErrors[finalLayer] = new double[this.layers[finalLayer].numNeurons];

        // calculate the delta for the output layer
        for (int i = 0; i < layerErrors[finalLayer].length; i++) {
            layerErrors[finalLayer][i] = -(point.expectedOutputs[i] - values[i])
                * Layer.activationDerivativeFunction(this.layers[finalLayer].rawLayerData[i]);
        }

        // calculate the delta for the hidden layers
        for (int i = finalLayer - 1; i >= 0; i--) {
            int layerSize = this.layers[i].numNeurons;
            layerErrors[i] = new double[layerSize];
            for (int j = 0; j < layerSize; j++) {
                layerErrors[i][j] = 0;
                for (int k = 0; k < this.layers[i+1].numNeurons; k++) {
                    layerErrors[i][j] += this.layers[i + 1].getWeight(j, k) * layerErrors[i + 1][k];
                }
                layerErrors[i][j] *= Layer.activationDerivativeFunction(layers[i].rawLayerData[j]);
            }
        }

        // update the weights and biases of each layer
        for (int l = 0; l < this.layers.length; l++) { // for each layer
            for (int j = 0; j < this.layers[l].numNeurons; j++) { // for each neuron
                for (int i = 0; i < this.layers[l].numInputNeurons; i++) { // for each input neuron
                    double w = this.layers[l].getWeight(i, j);
                    // update the weight using the delta of the neuron
                    // find the raw data of the neuron
                    // if l == 0, the raw data is the input data
                    double rawData = (l == 0) ? point.inputs[i] : layers[l - 1].rawLayerData[i];
                    // update the weight
                    // w = w - learningRate * delta * f(rawData)
                    w = w - this.learningRate * layerErrors[l][j]
                        * Layer.activationFunction(rawData);
                    // set the new weight
                    this.layers[l].setWeight(i, j, w);
                }
                // update the biases
                this.layers[l].biases[j] -= this.learningRate * layerErrors[l][j];
            } 
        }
    }

    public double evaluateTotalCost(DataSet data) {
        double res = 0.0;

        for (int i = 0; i < data.numPoints; ++i) {
            DataPoint point = data.points[i];
            double[] outputs = calculateOutputs(point.inputs);
            res += calculateCost(point.expectedOutputs, outputs);
        }

        return res / data.numPoints;
    }

    public double evaluateAccuracy(DataSet data) {
        int correct = 0;
        for (int i = 0; i < data.numPoints; ++i) {
            DataPoint point = data.points[i];
            double[] outputs = calculateOutputs(point.inputs);
            int predicted = classify(outputs);
            int expected = classify(point.expectedOutputs);
            if (predicted == expected) {
                correct++;
            }
        }
        return (double) correct / data.numPoints;
    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        sb.append("NeuralNetwork{\n");
        for (int i = 0; i < this.layers.length; ++i) {
            sb.append("    " + this.layers[i].toString() + "\n");
        }
        sb.append("}");
        return sb.toString();
    }

    public void save(String filename) throws IOException {
        try (ObjectOutputStream oos = new ObjectOutputStream(new FileOutputStream(filename))) {
            oos.writeObject(this);
        }
    }

    public static NeuralNetwork load(String filename) throws IOException, ClassNotFoundException {
        try (ObjectInputStream ois = new ObjectInputStream(new FileInputStream(filename))) {
            NeuralNetwork nn = (NeuralNetwork) ois.readObject();
            nn.random = new Random(0); 
            return nn;
        }
    }
}
