
package fr.perceptron;

import java.io.File;
import java.io.IOException;

public class MnistMain {

    public static void main(String[] args) {

        int[] layerSizes = {
            784,    // 28 * 28 = 784
            512,    // 256 neurons in the hidden layer 1
            // 128,    // 128 neurons in the hidden layer 2
            10      // 10 classes (0-9 digits)
        };
        if (args.length > 0) {
            layerSizes = new int[args.length];
            for (int i = 0; i < args.length; ++i) {
                layerSizes[i] = Integer.parseInt(args[i]);
            }
        }

        if (!new File("data.dat").exists()) {
            MNIST.loadTrain();
            MNIST.loadEval();
            MNIST.save();
        } else {
            MNIST.load();
        }

        DataSet data = new DataSet(1);
        DataSet eval = new DataSet(1);
        data.loadFromArray(MNIST.getTrainImages(), MNIST.getTrainLabels());
        data.batchesOf(1000);
        eval.loadFromArray(MNIST.getEvalImages(), MNIST.getEvalLabels());

        String cachedNN = "nn.dat";
        NeuralNetwork nn;
        System.out.println("Try loading cached NN...");
        try {
            nn = NeuralNetwork.load(cachedNN);
        } catch (ClassNotFoundException | IOException e) {
            System.out.println("Loading failed, creating new NN...");
            nn = new NeuralNetwork(layerSizes);
        }
        System.out.println("Done.");
        System.out.println(nn + "\n");

        int numEpochs = 500;
        long startTime = System.currentTimeMillis();
        long endTime, epochStartTime, elapsedTime;
        double[] learnRates = new double[numEpochs];
        double[] dataAccuracies = new double[numEpochs];
        // double[] NNCost = new double[numEpochs];
        double[] evalAccuracies = new double[numEpochs];
        // double cost = 0.0;

        // Train the Neural Network
        for (int i = nn.epoch; i < numEpochs; i++) {
            System.out.println("Epoch " + i);
            epochStartTime = System.currentTimeMillis();
            nn.train(data);
            elapsedTime = System.currentTimeMillis() - epochStartTime;
            System.out.println("Time elapsed: " + elapsedTime + "ms (" + (elapsedTime / 1000.0) + "s)");
            // cost = nn.evaluateTotalCost(data);
            double dataAccuracy = nn.evaluateAccuracy(data);
            double evalAccuracy = nn.evaluateAccuracy(eval);
            System.out.println(String.format("Learning rate: %.10f", nn.learningRate));
            System.out.println(String.format("Train data accuracy: %.5f%%", dataAccuracy * 100));
            System.out.println(String.format("Eval data accuracy: %.5f%%", evalAccuracy * 100));
            learnRates[i] = nn.learningRate;
            // NNCost[i] = cost;
            dataAccuracies[i] = dataAccuracy;
            evalAccuracies[i] = evalAccuracy;
            data.resetBatches();
            // Main.saveStats(learnRates, NNCost, dataAccuracies, "stats.txt");
            Main.saveStats(learnRates, dataAccuracies, evalAccuracies, "stats.txt");
            try {
                nn.save(cachedNN);
            } catch (IOException e) {
                System.out.println(e);
                e.printStackTrace();
                return;
            }
        }

        endTime = System.currentTimeMillis() - startTime;
        System.out.println("Total training time: " + endTime + "ms (" + (endTime / 1000.0) + "s)");
        // Main.saveStats(learnRates, NNCost, dataAccuracies, "stats.txt");
        Main.saveStats(learnRates, dataAccuracies, evalAccuracies, "stats.txt");

        try {
            nn.save(cachedNN);
        } catch (IOException e) {
            System.out.println(e);
            e.printStackTrace();
            return;
        }
    }
}
