package org.studiorailgun.conversation.categorization;

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

import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.BackpropType;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.optimize.api.InvocationType;
import org.deeplearning4j.optimize.listeners.EvaluativeListener;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.nd4j.evaluation.classification.Evaluation;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.learning.config.Sgd;

/**
 * Computational model for categorizing conversations
 */
public class ConvCategorizationNet {

    /**
     * The configuration for the network
     */
    MultiLayerConfiguration conf;

    /**
     * The actual model created from the configuration
     */
    MultiLayerNetwork model;

    /**
     * The number of epochs to try to train for
     */
    int targetEpochs;
    
    /**
     * Initializes the model
     */
    public void init(){
        this.conf = new NeuralNetConfiguration.Builder()
        //setup base config
            .weightInit(WeightInit.XAVIER)
            .activation(Activation.RELU)
            .optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT)
            .updater(new Sgd(0.05))
            .list()
        //add layers
            .layer(0, new DenseLayer.Builder()
                .nIn(4096)
                .nOut(4096)
                .build()
            )
        //add back propagation
            .backpropType(BackpropType.Standard)
        //finalize
        .build();

        this.model = new MultiLayerNetwork(conf);
        model.init();
    }

    public void apply(String input){
        // this.model.output();
    }

    /**
     * Trains the model on a set of data
     * @param trainData The training data
     * @param testData The test data
     */
    public void train(DataSetIterator trainData, DataSetIterator testData){
        System.out.println("Train model..");
        model.setListeners(new ScoreIterationListener(1), new EvaluativeListener(testData, 1, InvocationType.EPOCH_END));
        model.fit(trainData, targetEpochs);
    }

    /**
     * Evaluates the model against the test data
     * @param testData The test data
     */
    public void evaluate(DataSetIterator testData){
        System.out.println("Evaluate model..");
        Evaluation eval = model.evaluate(testData);
        System.out.println(eval.stats());
    }

    /**
     * Saves the model to disk
     * @param location The file to save the model into
     */
    public void save(File location){
        try {
            model.save(location, true);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * Loads the model from disk
     * @param location The location of the model file
     */
    public void load(File location){
        try {
            this.model = MultiLayerNetwork.load(location, true);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

}
