/*
 *  ******************************************************************************
 *  *
 *  *
 *  * 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.longrunning.downloads;

import org.datavec.api.records.reader.impl.csv.CSVRecordReader;
import org.datavec.api.split.FileSplit;
import org.datavec.image.loader.CifarLoader;
import org.datavec.image.loader.LFWLoader;
import org.deeplearning4j.BaseDL4JTest;
import org.deeplearning4j.datasets.datavec.RecordReaderDataSetIterator;
import org.deeplearning4j.datasets.fetchers.DataSetType;
import org.deeplearning4j.datasets.iterator.IteratorDataSetIterator;
import org.deeplearning4j.datasets.iterator.impl.*;
import org.deeplearning4j.eval.Evaluation;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.GradientNormalization;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.inputs.InputType;
import org.deeplearning4j.nn.conf.layers.ConvolutionLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.conf.layers.SubsamplingLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.optimize.listeners.CollectScoresIterationListener;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
import org.nd4j.common.tests.tags.NativeTag;
import org.nd4j.common.tests.tags.TagNames;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.common.io.ClassPathResource;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import static org.junit.jupiter.api.Assertions.*;
import org.junit.jupiter.api.DisplayName;

@DisplayName("Data Set Iterator Test")
@Tag(TagNames.FILE_IO)
@NativeTag
@Tag(TagNames.LARGE_RESOURCES)
@Tag(TagNames.LONG_TEST)
@Tag(TagNames.DOWNLOADS)
@Disabled("Disabling till we migrate to a new storage area for datasets")
class DataSetIteratorTest extends BaseDL4JTest {




    @Override
    public long getTimeoutMilliseconds() {
        // Should run quickly; increased to large timeout due to occasional slow CI downloads
        return 360000;
    }

    @Test
    @DisplayName("Test Batch Size Of One Iris")
    void testBatchSizeOfOneIris() throws Exception {
        // Test for (a) iterators returning correct number of examples, and
        // (b) Labels are a proper one-hot vector (i.e., sum is 1.0)
        // Iris:
        DataSetIterator iris = new IrisDataSetIterator(1, 5);
        int irisC = 0;
        while (iris.hasNext()) {
            irisC++;
            DataSet ds = iris.next();
            assertTrue(ds.getLabels().sum(Integer.MAX_VALUE).getDouble(0) == 1.0);
        }
        assertEquals(5, irisC);
    }

    @Test
    @DisplayName("Test Batch Size Of One Mnist")
    void testBatchSizeOfOneMnist() throws Exception {
        // MNIST:
        DataSetIterator mnist = new MnistDataSetIterator(1, 5);
        int mnistC = 0;
        while (mnist.hasNext()) {
            mnistC++;
            DataSet ds = mnist.next();
            assertTrue(ds.getLabels().sum(Integer.MAX_VALUE).getDouble(0) == 1.0);
        }
        assertEquals(5, mnistC);
    }

    @Test
    @DisplayName("Test Mnist")
    void testMnist(@TempDir Path tempDir) throws Exception {
        ClassPathResource cpr = new ClassPathResource("mnist_first_200.txt");
        CSVRecordReader rr = new CSVRecordReader(0, ',');
        rr.initialize(new FileSplit(cpr.getTempFileFromArchive()));
        RecordReaderDataSetIterator dsi = new RecordReaderDataSetIterator(rr, 10, 0, 10);
        MnistDataSetIterator iter = new MnistDataSetIterator(10, 200, false, true, false, 0,tempDir.toFile());
        while (dsi.hasNext()) {
            DataSet dsExp = dsi.next();
            DataSet dsAct = iter.next();
            INDArray fExp = dsExp.getFeatures();
            fExp.divi(255);
            INDArray lExp = dsExp.getLabels();
            INDArray fAct = dsAct.getFeatures();
            INDArray lAct = dsAct.getLabels();
            assertEquals(fExp, fAct.castTo(fExp.dataType()));
            assertEquals(lExp, lAct.castTo(lExp.dataType()));
        }
        assertFalse(iter.hasNext());
    }

    @Test
    @DisplayName("Test Lfw Iterator")
    void testLfwIterator() throws Exception {
        int numExamples = 1;
        int row = 28;
        int col = 28;
        int channels = 1;
        LFWDataSetIterator iter = new LFWDataSetIterator(numExamples, new int[] { row, col, channels }, true);
        assertTrue(iter.hasNext());
        DataSet data = iter.next();
        assertEquals(numExamples, data.getLabels().size(0));
        assertEquals(row, data.getFeatures().size(2));
    }

    @Test
    @DisplayName("Test Tiny Image Net Iterator")
    void testTinyImageNetIterator() throws Exception {
        int numClasses = 200;
        int row = 64;
        int col = 64;
        int channels = 3;
        TinyImageNetDataSetIterator iter = new TinyImageNetDataSetIterator(1, DataSetType.TEST);
        assertTrue(iter.hasNext());
        DataSet data = iter.next();
        assertEquals(numClasses, data.getLabels().size(1));
        assertArrayEquals(new long[] { 1, channels, row, col }, data.getFeatures().shape());
    }

    @Test
    @DisplayName("Test Tiny Image Net Iterator 2")
    void testTinyImageNetIterator2() throws Exception {
        int numClasses = 200;
        int row = 224;
        int col = 224;
        int channels = 3;
        TinyImageNetDataSetIterator iter = new TinyImageNetDataSetIterator(1, new int[] { row, col }, DataSetType.TEST);
        assertTrue(iter.hasNext());
        DataSet data = iter.next();
        assertEquals(numClasses, data.getLabels().size(1));
        assertArrayEquals(new long[] { 1, channels, row, col }, data.getFeatures().shape());
    }

    @Test
    @DisplayName("Test Lfw Model")
    void testLfwModel() throws Exception {
        final int numRows = 28;
        final int numColumns = 28;
        int numChannels = 3;
        int outputNum = LFWLoader.NUM_LABELS;
        int numSamples = LFWLoader.NUM_IMAGES;
        int batchSize = 2;
        int seed = 123;
        int listenerFreq = 1;
        LFWDataSetIterator lfw = new LFWDataSetIterator(batchSize, numSamples, new int[] { numRows, numColumns, numChannels }, outputNum, false, true, 1.0, new Random(seed));
        MultiLayerConfiguration.Builder builder = new NeuralNetConfiguration.Builder().seed(seed).gradientNormalization(GradientNormalization.RenormalizeL2PerLayer).optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT).list().layer(0, new ConvolutionLayer.Builder(5, 5).nIn(numChannels).nOut(6).weightInit(WeightInit.XAVIER).activation(Activation.RELU).build()).layer(1, new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX, new int[] { 2, 2 }).stride(1, 1).build()).layer(2, new OutputLayer.Builder(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD).nOut(outputNum).weightInit(WeightInit.XAVIER).activation(Activation.SOFTMAX).build()).setInputType(InputType.convolutionalFlat(numRows, numColumns, numChannels));
        MultiLayerNetwork model = new MultiLayerNetwork(builder.build());
        model.init();
        model.setListeners(new ScoreIterationListener(listenerFreq));
        model.fit(lfw.next());
        DataSet dataTest = lfw.next();
        INDArray output = model.output(dataTest.getFeatures());
        Evaluation eval = new Evaluation(outputNum);
        eval.eval(dataTest.getLabels(), output);
        // System.out.println(eval.stats());
    }

    @Test
    @DisplayName("Test Cifar 10 Iterator")
    void testCifar10Iterator() throws Exception {
        int numExamples = 1;
        int row = 32;
        int col = 32;
        int channels = 3;
        Cifar10DataSetIterator iter = new Cifar10DataSetIterator(numExamples);
        assertTrue(iter.hasNext());
        DataSet data = iter.next();
        assertEquals(numExamples, data.getLabels().size(0));
        assertEquals(channels * row * col, data.getFeatures().ravel().length());
    }

    @Test
    @DisplayName("Test Cifar Model")
    void testCifarModel() throws Exception {
        // Streaming
        runCifar(false);
        // Preprocess
        runCifar(true);
    }

    public void runCifar(boolean preProcessCifar) throws Exception {
        final int height = 32;
        final int width = 32;
        int channels = 3;
        int outputNum = CifarLoader.NUM_LABELS;
        int batchSize = 5;
        int seed = 123;
        int listenerFreq = 1;
        Cifar10DataSetIterator cifar = new Cifar10DataSetIterator(batchSize);
        MultiLayerConfiguration.Builder builder = new NeuralNetConfiguration.Builder().seed(seed).gradientNormalization(GradientNormalization.RenormalizeL2PerLayer).optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT).list().layer(0, new ConvolutionLayer.Builder(5, 5).nIn(channels).nOut(6).weightInit(WeightInit.XAVIER).activation(Activation.RELU).build()).layer(1, new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX, new int[] { 2, 2 }).build()).layer(2, new OutputLayer.Builder(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD).nOut(outputNum).weightInit(WeightInit.XAVIER).activation(Activation.SOFTMAX).build()).setInputType(InputType.convolutionalFlat(height, width, channels));
        MultiLayerNetwork model = new MultiLayerNetwork(builder.build());
        model.init();
        // model.setListeners(Arrays.asList((TrainingListener) new ScoreIterationListener(listenerFreq)));
        CollectScoresIterationListener listener = new CollectScoresIterationListener(listenerFreq);
        model.setListeners(listener);
        model.fit(cifar);
        cifar = new Cifar10DataSetIterator(batchSize);
        Evaluation eval = new Evaluation(cifar.getLabels());
        while (cifar.hasNext()) {
            DataSet testDS = cifar.next(batchSize);
            INDArray output = model.output(testDS.getFeatures());
            eval.eval(testDS.getLabels(), output);
        }
        // System.out.println(eval.stats(true));
        listener.exportScores(System.out);
    }

    @Test
    @DisplayName("Test Iterator Data Set Iterator Combining")
    void testIteratorDataSetIteratorCombining() {
        // Test combining of a bunch of small (size 1) data sets together
        int batchSize = 3;
        int numBatches = 4;
        int featureSize = 5;
        int labelSize = 6;
        Nd4j.getRandom().setSeed(12345);
        List<DataSet> orig = new ArrayList<>();
        for (int i = 0; i < batchSize * numBatches; i++) {
            INDArray features = Nd4j.rand(1, featureSize);
            INDArray labels = Nd4j.rand(1, labelSize);
            orig.add(new DataSet(features, labels));
        }
        DataSetIterator iter = new IteratorDataSetIterator(orig.iterator(), batchSize);
        int count = 0;
        while (iter.hasNext()) {
            DataSet ds = iter.next();
            assertArrayEquals(new long[] { batchSize, featureSize }, ds.getFeatures().shape());
            assertArrayEquals(new long[] { batchSize, labelSize }, ds.getLabels().shape());
            List<INDArray> fList = new ArrayList<>();
            List<INDArray> lList = new ArrayList<>();
            for (int i = 0; i < batchSize; i++) {
                DataSet dsOrig = orig.get(count * batchSize + i);
                fList.add(dsOrig.getFeatures());
                lList.add(dsOrig.getLabels());
            }
            INDArray fExp = Nd4j.vstack(fList);
            INDArray lExp = Nd4j.vstack(lList);
            assertEquals(fExp, ds.getFeatures());
            assertEquals(lExp, ds.getLabels());
            count++;
        }
        assertEquals(count, numBatches);
    }

    @Test
    @DisplayName("Test Iterator Data Set Iterator Splitting")
    void testIteratorDataSetIteratorSplitting() {
        // Test splitting large data sets into smaller ones
        int origBatchSize = 4;
        int origNumDSs = 3;
        int batchSize = 3;
        int numBatches = 4;
        int featureSize = 5;
        int labelSize = 6;
        Nd4j.getRandom().setSeed(12345);
        List<DataSet> orig = new ArrayList<>();
        for (int i = 0; i < origNumDSs; i++) {
            INDArray features = Nd4j.rand(origBatchSize, featureSize);
            INDArray labels = Nd4j.rand(origBatchSize, labelSize);
            orig.add(new DataSet(features, labels));
        }
        List<DataSet> expected = new ArrayList<>();
        expected.add(new DataSet(orig.get(0).getFeatures().getRows(0, 1, 2), orig.get(0).getLabels().getRows(0, 1, 2)));
        expected.add(new DataSet(Nd4j.vstack(orig.get(0).getFeatures().getRows(3), orig.get(1).getFeatures().getRows(0, 1)), Nd4j.vstack(orig.get(0).getLabels().getRows(3), orig.get(1).getLabels().getRows(0, 1))));
        expected.add(new DataSet(Nd4j.vstack(orig.get(1).getFeatures().getRows(2, 3), orig.get(2).getFeatures().getRows(0)), Nd4j.vstack(orig.get(1).getLabels().getRows(2, 3), orig.get(2).getLabels().getRows(0))));
        expected.add(new DataSet(orig.get(2).getFeatures().getRows(1, 2, 3), orig.get(2).getLabels().getRows(1, 2, 3)));
        DataSetIterator iter = new IteratorDataSetIterator(orig.iterator(), batchSize);
        int count = 0;
        while (iter.hasNext()) {
            DataSet ds = iter.next();
            assertEquals(expected.get(count), ds);
            count++;
        }
        assertEquals(count, numBatches);
    }
}
