package com.ty.test;

import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.nd4j.evaluation.classification.Evaluation;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.learning.config.Nesterovs;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.io.File;
import java.io.FileNotFoundException;
import java.util.Scanner;

public class Test {
    public static void main(String[] args) throws FileNotFoundException {

        final int numRows = 38;
        final int numColumns = 50;
        int outputNum = 2; // number of output classes
        int rngSeed = 123; // random number seed for reproducibility
        int numEpochs = 100; // number of epochs to perform

        Scanner scanner = new Scanner(new File("out.txt"));
        double[][] vectors = new double[numRows][numColumns];
        double[][] outputs = new double[numRows][2];
        int index = 0;
        while (scanner.hasNextLine()) {
            scanner.nextLine();
            vectors[index][index] = 1;
            index++;
        }

        for(int i=0;i<38;i++){
            if(i==7||i==14){
                outputs[i][1] = 1;
            }else{
                outputs[i][0] = 1;
            }
        }

        INDArray inputs = Nd4j.create(vectors);
        INDArray desiredOutputs = Nd4j.create(outputs);
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
                .seed(rngSeed) //include a random seed for reproducibility
                // use stochastic gradient descent as an optimization algorithm
                .updater(new Nesterovs(0.9, 0.9))
                .l2(1e-4)
                .list()
                .layer(0, new DenseLayer.Builder() //create the first, input layer with xavier initialization
                        .nIn(numColumns)
                        .nOut(2)
                        .activation(Activation.RELU)
                        .weightInit(WeightInit.XAVIER)
                        .build())
                .layer(1, new OutputLayer.Builder(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD) //create hidden layer
                        .nIn(2)
                        .nOut(outputNum)
                        .activation(Activation.SOFTMAX)
                        .weightInit(WeightInit.XAVIER)
                        .build())
                .build();

        MultiLayerNetwork model = new MultiLayerNetwork(conf);
        model.init();
        //print the score with every 1 iteration
        model.setListeners(new ScoreIterationListener(1));

        System.out.println("train model");
        for (int j = 0; j < numEpochs; j++) {
            model.fit(inputs, desiredOutputs);
        }

        Evaluation eval = new Evaluation(outputNum); //create an evaluation object with 10 possible classes

        for (int i = 0; i < numRows; i++) {
            INDArray inputs2 = Nd4j.create(vectors[i]);
            INDArray outputs2 = Nd4j.create(new double[]{outputs[i][0], outputs[i][1]});
            DataSet ds = new DataSet(inputs2, outputs2);
            INDArray output = model.output(ds.getFeatures());
            System.out.println(output);
            eval.eval(ds.getLabels(), output);
        }


        System.out.println(eval.stats());

    }
}