/*
 *  ******************************************************************************
 *  *
 *  *
 *  * This program and the accompanying materials are made available under the
 *  * terms of the Apache License, Version 2.0 which is available at
 *  * https://www.apache.org/licenses/LICENSE-2.0.
 *  *
 *  *  See the NOTICE file distributed with this work for additional
 *  *  information regarding copyright ownership.
 *  * Unless required by applicable law or agreed to in writing, software
 *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 *  * License for the specific language governing permissions and limitations
 *  * under the License.
 *  *
 *  * SPDX-License-Identifier: Apache-2.0
 *  *****************************************************************************
 */

package org.eclipse.deeplearning4j.integration.testcases.dl4j;

import org.eclipse.deeplearning4j.integration.ModelType;
import org.eclipse.deeplearning4j.integration.TestCase;
import org.eclipse.deeplearning4j.integration.testcases.dl4j.misc.CharacterIterator;
import org.nd4j.evaluation.IEvaluation;
import org.nd4j.evaluation.classification.Evaluation;
import org.nd4j.evaluation.classification.EvaluationCalibration;
import org.nd4j.evaluation.classification.ROCMultiClass;
import org.nd4j.linalg.api.buffer.DataType;
import org.nd4j.linalg.dataset.adapter.MultiDataSetIteratorAdapter;
import org.nd4j.linalg.dataset.api.preprocessor.CompositeMultiDataSetPreProcessor;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.shade.guava.io.Files;
import org.datavec.api.records.reader.SequenceRecordReader;
import org.datavec.api.records.reader.impl.csv.CSVSequenceRecordReader;
import org.datavec.api.split.NumberedFileInputSplit;
import org.deeplearning4j.datasets.datavec.SequenceRecordReaderDataSetIterator;
import org.deeplearning4j.datasets.iterator.EarlyTerminationDataSetIterator;
import org.deeplearning4j.nn.conf.BackpropType;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.inputs.InputType;
import org.deeplearning4j.nn.conf.layers.*;
import org.deeplearning4j.nn.conf.layers.recurrent.Bidirectional;
import org.deeplearning4j.nn.weights.WeightInit;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.api.MultiDataSet;
import org.nd4j.linalg.dataset.api.MultiDataSetPreProcessor;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.dataset.api.iterator.MultiDataSetIterator;
import org.nd4j.linalg.dataset.api.preprocessor.MultiDataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.MultiNormalizerMinMaxScaler;
import org.nd4j.linalg.dataset.api.preprocessor.MultiNormalizerStandardize;
import org.nd4j.linalg.indexing.NDArrayIndex;
import org.nd4j.common.io.ClassPathResource;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.lossfunctions.LossFunctions;
import org.nd4j.common.primitives.Pair;

import java.io.File;
import java.util.Collections;
import java.util.List;

public class RNNTestCases {

    /**
     * RNN + global pooling + CSV + normalizer
     */
    public static TestCase getRnnCsvSequenceClassificationTestCase1(){
        return new RnnCsvSequenceClassificationTestCase1();
    }

    public static TestCase getRnnCsvSequenceClassificationTestCase2(){
        return new RnnCsvSequenceClassificationTestCase2();
    }

    public static TestCase getRnnCharacterTestCase(){
        return new TestCase() {
            {
                testName = "RnnCharacterTestCase";
                testType = TestType.RANDOM_INIT;
                testPredictions = true;
                testTrainingCurves = true;
                testGradients = true;
                testParamsPostTraining = true;
                testEvaluation = true;
                testOverfitting = false;            //Not much point on this one - it already fits very well...
                //Gradients depend on a lot of chained steps, numerical differences can accumulate
                maxRelativeErrorGradients = 5e-4;
                minAbsErrorGradients = 2e-4;
                maxRelativeErrorParamsPostTraining = 1e-4;
                minAbsErrorParamsPostTraining = 2e-3;
            }

            private int miniBatchSize = 32;
            private int exampleLength = 200;


            @Override
            public ModelType modelType() {
                return ModelType.MLN;
            }

            @Override
            public Object getConfiguration() throws Exception {
                Nd4j.getRandom().setSeed(12345);

                CharacterIterator iter = CharacterIterator.getShakespeareIterator(miniBatchSize,exampleLength);
                int nOut = iter.totalOutcomes();

                int lstmLayerSize = 200;					//Number of units in each GravesLSTM layer
                int tbpttLength = 50;                       //Length for truncated backpropagation through time. i.e., do parameter updates ever 50 characters

                return new NeuralNetConfiguration.Builder()
                        .dataType(DataType.FLOAT)
                        .seed(12345)
                        .l2(0.001)
                        .weightInit(WeightInit.XAVIER)
                        .updater(new Adam(1e-3))
                        .list()
                        .layer(0, new LSTM.Builder().nIn(iter.inputColumns()).nOut(lstmLayerSize)
                                .activation(Activation.TANH).build())
                        .layer(1, new LSTM.Builder().nIn(lstmLayerSize).nOut(lstmLayerSize)
                                .activation(Activation.TANH).build())
                        .layer(2, new RnnOutputLayer.Builder(LossFunctions.LossFunction.MCXENT).activation(Activation.SOFTMAX)        //MCXENT + softmax for classification
                                .nIn(lstmLayerSize).nOut(nOut).build())
                        .backpropType(BackpropType.TruncatedBPTT).tBPTTForwardLength(tbpttLength).tBPTTBackwardLength(tbpttLength)

                        .build();
            }

            @Override
            public List<Pair<INDArray[], INDArray[]>> getPredictionsTestData() throws Exception {
                MultiDataSet mds = getTrainingData().next();
                return Collections.singletonList(new Pair<>(mds.getFeatures(), mds.getFeaturesMaskArrays()));
            }

            @Override
            public MultiDataSet getGradientsTestData() throws Exception {
                return getTrainingData().next();
            }

            @Override
            public MultiDataSetIterator getTrainingData() throws Exception {
                DataSetIterator iter = CharacterIterator.getShakespeareIterator(miniBatchSize,exampleLength);
                iter = new EarlyTerminationDataSetIterator(iter, 2);    //2 minibatches, 200/50 = 4 updates per minibatch
                return new MultiDataSetIteratorAdapter(iter);
            }

            @Override
            public IEvaluation[] getNewEvaluations(){
                return new IEvaluation[]{
                        new Evaluation(),
                        new ROCMultiClass()};
            }

            @Override
            public MultiDataSetIterator getEvaluationTestData() throws Exception {
                return getTrainingData();
            }
        };
    }

    protected static class RnnCsvSequenceClassificationTestCase1 extends TestCase {
        protected   RnnCsvSequenceClassificationTestCase1(){
            testName = "RnnCsvSequenceClassification1";
            testType = TestType.RANDOM_INIT;
            testPredictions = true;
            testTrainingCurves = true;
            testGradients = true;
            testParamsPostTraining = true;
            testEvaluation = true;
            testOverfitting = false;            //Not much point on this one - it already fits very well...
        }


        protected MultiDataNormalization normalizer;

        protected MultiDataNormalization getNormalizer() throws Exception {
            if(normalizer != null){
                return normalizer;
            }

            normalizer = new MultiNormalizerStandardize();
            normalizer.fit(getTrainingDataUnnormalized());

            return normalizer;
        }

        @Override
        public ModelType modelType() {
            return ModelType.MLN;
        }

        @Override
        public Object getConfiguration() throws Exception {
            return new NeuralNetConfiguration.Builder()
                    .dataType(DataType.FLOAT)
                    .seed(12345)
                    .updater(new Adam(5e-2))
                    .l1(1e-3).l2(1e-3)
                    .list()
                    .layer(0, new LSTM.Builder().activation(Activation.TANH).nOut(10).build())
                    .layer(new GlobalPoolingLayer.Builder().poolingType(PoolingType.AVG).build())
                    .layer(new OutputLayer.Builder().nOut(6)
                            .lossFunction(LossFunctions.LossFunction.MCXENT)
                            .activation(Activation.SOFTMAX)
                            .build())
                    .setInputType(InputType.recurrent(1))
                    .build();
        }

        @Override
        public List<Pair<INDArray[], INDArray[]>> getPredictionsTestData() throws Exception {
            MultiDataSet mds = getTrainingData().next();
            return Collections.singletonList(new Pair<>(mds.getFeatures(), mds.getFeaturesMaskArrays()));
        }

        @Override
        public MultiDataSet getGradientsTestData() throws Exception {
            return getTrainingData().next();
        }

        @Override
        public MultiDataSetIterator getTrainingData() throws Exception {
            MultiDataSetIterator iter = getTrainingDataUnnormalized();

            MultiDataSetPreProcessor pp = multiDataSet -> {
                INDArray l = multiDataSet.getLabels(0);
                l = l.get(NDArrayIndex.all(), NDArrayIndex.all(), NDArrayIndex.point(l.size(2)-1));
                multiDataSet.setLabels(0, l);
                multiDataSet.setLabelsMaskArray(0, null);
            };


            iter.setPreProcessor(new CompositeMultiDataSetPreProcessor(getNormalizer(),pp));

            return iter;
        }

        protected MultiDataSetIterator getTrainingDataUnnormalized() throws Exception {
            int miniBatchSize = 10;
            int numLabelClasses = 6;

            File featuresDirTrain = Files.createTempDir();
            File labelsDirTrain = Files.createTempDir();
            new ClassPathResource("dl4j-integration-tests/data/uci_seq/train/features/").copyDirectory(featuresDirTrain);
            new ClassPathResource("dl4j-integration-tests/data/uci_seq/train/labels/").copyDirectory(labelsDirTrain);

            SequenceRecordReader trainFeatures = new CSVSequenceRecordReader();
            trainFeatures.initialize(new NumberedFileInputSplit(featuresDirTrain.getAbsolutePath() + "/%d.csv", 0, 449));
            SequenceRecordReader trainLabels = new CSVSequenceRecordReader();
            trainLabels.initialize(new NumberedFileInputSplit(labelsDirTrain.getAbsolutePath() + "/%d.csv", 0, 449));

            DataSetIterator trainData = new SequenceRecordReaderDataSetIterator(trainFeatures, trainLabels, miniBatchSize, numLabelClasses,
                    false, SequenceRecordReaderDataSetIterator.AlignmentMode.ALIGN_END);

            MultiDataSetIterator iter = new MultiDataSetIteratorAdapter(trainData);
            return iter;
        }

        @Override
        public IEvaluation[] getNewEvaluations(){
            return new IEvaluation[]{
                    new Evaluation(),
                    new ROCMultiClass(),
                    new EvaluationCalibration()
            };
        }

        @Override
        public MultiDataSetIterator getEvaluationTestData() throws Exception {
            int miniBatchSize = 10;
            int numLabelClasses = 6;

//            File featuresDirTest = new ClassPathResource("/RnnCsvSequenceClassification/uci_seq/test/features/").getFile();
//            File labelsDirTest = new ClassPathResource("/RnnCsvSequenceClassification/uci_seq/test/labels/").getFile();
            File featuresDirTest = Files.createTempDir();
            File labelsDirTest = Files.createTempDir();
            new ClassPathResource("dl4j-integration-tests/data/uci_seq/test/features/").copyDirectory(featuresDirTest);
            new ClassPathResource("dl4j-integration-tests/data/uci_seq/test/labels/").copyDirectory(labelsDirTest);

            SequenceRecordReader trainFeatures = new CSVSequenceRecordReader();
            trainFeatures.initialize(new NumberedFileInputSplit(featuresDirTest.getAbsolutePath() + "/%d.csv", 0, 149));
            SequenceRecordReader trainLabels = new CSVSequenceRecordReader();
            trainLabels.initialize(new NumberedFileInputSplit(labelsDirTest.getAbsolutePath() + "/%d.csv", 0, 149));

            DataSetIterator testData = new SequenceRecordReaderDataSetIterator(trainFeatures, trainLabels, miniBatchSize, numLabelClasses,
                    false, SequenceRecordReaderDataSetIterator.AlignmentMode.ALIGN_END);

            MultiDataSetIterator iter = new MultiDataSetIteratorAdapter(testData);

            MultiDataSetPreProcessor pp = multiDataSet -> {
                INDArray l = multiDataSet.getLabels(0);
                l = l.get(NDArrayIndex.all(), NDArrayIndex.all(), NDArrayIndex.point(l.size(2)-1));
                multiDataSet.setLabels(0, l);
                multiDataSet.setLabelsMaskArray(0, null);
            };


            iter.setPreProcessor(new CompositeMultiDataSetPreProcessor(getNormalizer(),pp));

            return iter;
        }
    }

    /**
     * Similar to test case 1 - but using GravesLSTM + bidirectional wrapper + min/max scaler normalizer
     */
    protected static class RnnCsvSequenceClassificationTestCase2 extends RnnCsvSequenceClassificationTestCase1 {
        protected RnnCsvSequenceClassificationTestCase2() {
            super();
            testName = "RnnCsvSequenceClassification2";
        }

        @Override
        public Object getConfiguration() throws Exception {
            return new NeuralNetConfiguration.Builder()
                    .dataType(DataType.FLOAT)
                    .seed(12345)
                    .updater(new Adam(5e-2))
                    .l1(1e-3).l2(1e-3)
                    .list()
                    .layer(0, new Bidirectional(new LSTM.Builder().activation(Activation.TANH).nOut(10).build()))
                    .layer(new GlobalPoolingLayer.Builder().poolingType(PoolingType.AVG).build())
                    .layer(new OutputLayer.Builder().nOut(6)
                            .lossFunction(LossFunctions.LossFunction.MCXENT)
                            .activation(Activation.SOFTMAX)
                            .build())
                    .setInputType(InputType.recurrent(1))
                    .build();
        }

        protected MultiDataNormalization getNormalizer() throws Exception {
            if(normalizer != null){
                return normalizer;
            }

            normalizer = new MultiNormalizerMinMaxScaler();
            normalizer.fit(getTrainingDataUnnormalized());

            return normalizer;
        }
    }


}
