package fr.perceptron;

import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.util.zip.GZIPInputStream;

public class MNIST {

    private static double[] labels_train;
    private static double[][] images_train;
    private static double[] labels_eval;
    private static double[][] images_eval;

    public static double[] getEvalLabels() {
        return labels_eval;
    }

    public static double[][] getEvalImages() {
        return images_eval;
    }

    public static double[] getTrainLabels() {
        return labels_train;
    }

    public static double[][] getTrainImages() {
        return images_train;
    }

    public static void loadEval() {
        final String trainImagesPath = "./mnist/t10k-images-idx3-ubyte.gz";
        final String trainLabelsPath = "./mnist/t10k-labels-idx1-ubyte.gz";

        try {
            InputStream imgIn = new GZIPInputStream(new FileInputStream(trainImagesPath));
            InputStream lblIn = new GZIPInputStream(new FileInputStream(trainLabelsPath));

            byte[] tempBuffer = new byte[16];
            imgIn.read(tempBuffer, 0, 16);
            lblIn.read(tempBuffer, 0, 8); // Les labels n'ont que 8 octets d'en-tête

            byte[] dataBuffer = new byte[1];
            labels_eval = new double[10000];
            images_eval = new double[10000][784];

            System.out.println("Start loading MNIST set.");

            for (int i = 0; i < 10000; i++) {
                lblIn.read(dataBuffer, 0, 1);
                labels_eval[i] = dataBuffer[0] & 0xFF;

                for (int j = 0; j < 784; j++) {
                    imgIn.read(dataBuffer, 0, 1);
                    double pixelVal = (dataBuffer[0] & 0xFF) / 255.0;
                    images_eval[i][j] = pixelVal;
                }
            }
            System.out.println("MNIST set loaded.");
            imgIn.close();
            lblIn.close();

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public static void loadTrain() {
        final String trainImagesPath = "./mnist/train-images-idx3-ubyte.gz";
        final String trainLabelsPath = "./mnist/train-labels-idx1-ubyte.gz";

        try {
            InputStream imgIn = new GZIPInputStream(new FileInputStream(trainImagesPath));
            InputStream lblIn = new GZIPInputStream(new FileInputStream(trainLabelsPath));

            byte[] tempBuffer = new byte[16];
            imgIn.read(tempBuffer, 0, 16);
            lblIn.read(tempBuffer, 0, 8); // Les labels n'ont que 8 octets d'en-tête

            byte[] dataBuffer = new byte[1];
            labels_train = new double[60000];
            images_train = new double[60000][784];

            System.out.println("Start loading MNIST set.");

            for (int i = 0; i < 60000; i++) {
                lblIn.read(dataBuffer, 0, 1);
                labels_train[i] = dataBuffer[0] & 0xFF;

                for (int j = 0; j < 784; j++) {
                    imgIn.read(dataBuffer, 0, 1);
                    double pixelVal = (dataBuffer[0] & 0xFF) / 255.0;
                    images_train[i][j] = pixelVal;
                }
            }
            System.out.println("MNIST set loaded.");
            imgIn.close();
            lblIn.close();

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public static void save() {
        try {
            System.out.println("Saving set to cache...");
            ObjectOutputStream oos = new ObjectOutputStream(new FileOutputStream("data.dat"));
            oos.writeObject(labels_train);
            oos.writeObject(images_train);
            oos.writeObject(labels_eval);
            oos.writeObject(images_eval);

            oos.close();
            System.err.println("Set saved.");

        } catch (Exception e) {
        }
    }

    public static void load() {
        try {
            System.out.println("Loading cache...");
            ObjectInputStream ois = new ObjectInputStream(new FileInputStream("data.dat"));       
            labels_train = (double[]) ois.readObject();
            images_train = (double[][]) ois.readObject();
            labels_eval = (double[]) ois.readObject();
            images_eval = (double[][]) ois.readObject();
            ois.close();
            System.out.println("Cache loaded.");
        } catch (Exception e) {
        }
    }
}
