package fr.perceptron;

import java.io.BufferedWriter;
import java.io.FileWriter;

public class Main {

    public static void main(String[] args) {

        int[] layerSizes = { 2, 50, 10, 1 };
        if (args.length > 0) {
            layerSizes = new int[args.length];
            for (int i = 0; i < args.length; ++i) {
                layerSizes[i] = Integer.parseInt(args[i]);
            }
        }

        String file = "data/twocircles_data_";
        DataSet data = new DataSet(20);
        DataSet eval = new DataSet(1);
        data.loadDataSet(file + "train.csv");
        eval.loadDataSet(file + "eval.csv");

        NeuralNetwork nn = new NeuralNetwork(layerSizes);
        System.out.println(nn + "\n");

        Window window = new Window();
        window.init(nn, data);

        long startTime = System.currentTimeMillis();
        long endTime, epochStartTime, elapsedTime;
        double totalCost = 0.0;

        double[] learnRate;
        double[] NNCost;
        learnRate = new double[500];
        NNCost = new double[500];
        double oldCost;

        // Train the Neural Network
        for (int i = 0; i < 500; i++) {
            System.out.println("Epoch " + i);
            epochStartTime = System.currentTimeMillis();
            nn.train(data);
            elapsedTime = System.currentTimeMillis();
            System.out.println("Time elapsed: " + (elapsedTime - epochStartTime) + " ms");
            oldCost = totalCost;
            totalCost = nn.evaluateTotalCost(data);
            System.out.println("Total cost: " + totalCost + " Learning rate: " + nn.learningRate);
            window.update();
            learnRate[i] = nn.learningRate;
            NNCost[i] = totalCost;
            if (totalCost < NeuralNetwork.COST_TOLERANCE) break;
            if (Math.abs(oldCost - totalCost) < NeuralNetwork.COST_VARIATION_TOLERANCE) break;
            data.resetBatches();
        }

        endTime = System.currentTimeMillis() - startTime;
        System.out.println("Total training time: " + endTime + " ms (" + (endTime / 1000.0) + " s)");
        window.update();

        int correct = 0;
        DataSet d = eval;
        for (int i = 0; i < d.points.length; i++) {
            double res = nn.calculateOutputs(d.points[i].inputs)[0];
            correct += (res > 0.5 ? 1.0 : 0.0) == d.points[i].expectedOutputs[0] ? 1 : 0;
        }
        System.out.println("Correct points on eval set: " + correct + " / " + d.points.length);
        System.out.println("Accuracy: " + (correct * 100.0 / d.points.length) + "%");

        // Save the Neural Network
        saveStats(learnRate, NNCost, null, "stats.txt");
    }

    public static void saveStats(double[] learnRate, double[] NNCost, double[] accuracies, String filename) {
        // Save learn rate and cost in a text file, where each line contains:
        // learnRate[i], NNCost[i]
        try (BufferedWriter writer = new BufferedWriter(new FileWriter("stats/" + filename))) {
            for (int i = 0; i < learnRate.length; i++) {
                writer.write(learnRate[i] + "," + NNCost[i]);
                if (accuracies != null) {
                    writer.write("," + accuracies[i]);
                }
                writer.newLine();
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
