/*
 *  ******************************************************************************
 *  *
 *  *
 *  * 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.dl4jcore.nn.layers;

import lombok.extern.slf4j.Slf4j;
import org.deeplearning4j.BaseDL4JTest;
import org.eclipse.deeplearning4j.dl4jcore.TestUtils;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.*;
import org.deeplearning4j.nn.conf.distribution.NormalDistribution;
import org.deeplearning4j.nn.conf.layers.*;
import org.deeplearning4j.nn.conf.preprocessor.FeedForwardToRnnPreProcessor;
import org.deeplearning4j.nn.conf.preprocessor.RnnToFeedForwardPreProcessor;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.nn.layers.recurrent.RnnOutputLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.nn.workspace.LayerWorkspaceMgr;
import org.deeplearning4j.optimize.api.TrainingListener;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
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.buffer.DataType;
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.NoOp;
import org.nd4j.linalg.learning.config.Sgd;
import org.nd4j.linalg.lossfunctions.LossFunctions.LossFunction;
import java.util.Collections;
import java.util.Random;
import static org.junit.jupiter.api.Assertions.*;
import org.junit.jupiter.api.DisplayName;

@Slf4j
@DisplayName("Output Layer Test")
@NativeTag
@Tag(TagNames.CUSTOM_FUNCTIONALITY)
@Tag(TagNames.DL4J_OLD_API)
class OutputLayerTest extends BaseDL4JTest {

    @Test
    @DisplayName("Test Set Params")
    void testSetParams() {
        NeuralNetConfiguration conf = new NeuralNetConfiguration.Builder().optimizationAlgo(OptimizationAlgorithm.LINE_GRADIENT_DESCENT).updater(new Sgd(1e-1)).layer(new org.deeplearning4j.nn.conf.layers.OutputLayer.Builder().nIn(4).nOut(3).weightInit(WeightInit.ZERO).activation(Activation.SOFTMAX).lossFunction(LossFunction.MCXENT).build()).build();
        long numParams = conf.getLayer().initializer().numParams(conf);
        INDArray params = Nd4j.create(1, numParams);
        org.deeplearning4j.nn.layers.OutputLayer l = (org.deeplearning4j.nn.layers.OutputLayer) conf.getLayer().instantiate(conf, Collections.<TrainingListener>singletonList(new ScoreIterationListener(1)), 0, params, true, params.dataType());
        params = l.params();
        l.setParams(params);
        assertEquals(params, l.params());
    }

    @Test
    @DisplayName("Test Output Layers Rnn Forward Pass")
    void testOutputLayersRnnForwardPass() {
        // Test output layer with RNNs (
        // Expect all outputs etc. to be 2d
        int nIn = 2;
        int nOut = 5;
        int layerSize = 4;
        int timeSeriesLength = 6;
        int miniBatchSize = 3;
        Random r = new Random(12345L);
        INDArray input = Nd4j.zeros(miniBatchSize, nIn, timeSeriesLength);
        for (int i = 0; i < miniBatchSize; i++) {
            for (int j = 0; j < nIn; j++) {
                for (int k = 0; k < timeSeriesLength; k++) {
                    input.putScalar(new int[] { i, j, k }, r.nextDouble() - 0.5);
                }
            }
        }
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().seed(12345L).list().layer(0, new GravesLSTM.Builder().nIn(nIn).nOut(layerSize).dist(new NormalDistribution(0, 1)).activation(Activation.TANH).updater(new NoOp()).build()).layer(1, new org.deeplearning4j.nn.conf.layers.OutputLayer.Builder(LossFunction.MCXENT).activation(Activation.SOFTMAX).nIn(layerSize).nOut(nOut).dist(new NormalDistribution(0, 1)).updater(new NoOp()).build()).inputPreProcessor(1, new RnnToFeedForwardPreProcessor()).build();
        MultiLayerNetwork mln = new MultiLayerNetwork(conf);
        mln.init();
        INDArray out2d = mln.feedForward(input).get(2);
        assertArrayEquals(out2d.shape(), new long[] { miniBatchSize * timeSeriesLength, nOut });
        INDArray out = mln.output(input);
        assertArrayEquals(out.shape(), new long[] { miniBatchSize * timeSeriesLength, nOut });
        INDArray preout = mln.output(input);
        assertArrayEquals(preout.shape(), new long[] { miniBatchSize * timeSeriesLength, nOut });
        // As above, but for RnnOutputLayer. Expect all activations etc. to be 3d
        MultiLayerConfiguration confRnn = new NeuralNetConfiguration.Builder().seed(12345L).list().layer(0, new GravesLSTM.Builder().nIn(nIn).nOut(layerSize).dist(new NormalDistribution(0, 1)).activation(Activation.TANH).updater(new NoOp()).build()).layer(1, new org.deeplearning4j.nn.conf.layers.RnnOutputLayer.Builder(LossFunction.MCXENT).activation(Activation.SOFTMAX).nIn(layerSize).nOut(nOut).dist(new NormalDistribution(0, 1)).updater(new NoOp()).build()).build();
        MultiLayerNetwork mlnRnn = new MultiLayerNetwork(confRnn);
        mln.init();
        INDArray out3d = mlnRnn.feedForward(input).get(2);
        assertArrayEquals(out3d.shape(), new long[] { miniBatchSize, nOut, timeSeriesLength });
        INDArray outRnn = mlnRnn.output(input);
        assertArrayEquals(outRnn.shape(), new long[] { miniBatchSize, nOut, timeSeriesLength });
        INDArray preoutRnn = mlnRnn.output(input);
        assertArrayEquals(preoutRnn.shape(), new long[] { miniBatchSize, nOut, timeSeriesLength });
    }

    @Test
    @DisplayName("Test Rnn Output Layer Inc Edge Cases")
    void testRnnOutputLayerIncEdgeCases() {
        // Basic test + test edge cases: timeSeriesLength==1, miniBatchSize==1, both
        int[] tsLength = { 5, 1, 5, 1 };
        int[] miniBatch = { 7, 7, 1, 1 };
        int nIn = 3;
        int nOut = 6;
        int layerSize = 4;
        FeedForwardToRnnPreProcessor proc = new FeedForwardToRnnPreProcessor();
        for (int t = 0; t < tsLength.length; t++) {
            Nd4j.getRandom().setSeed(12345);
            int timeSeriesLength = tsLength[t];
            int miniBatchSize = miniBatch[t];
            Random r = new Random(12345L);
            INDArray input = Nd4j.zeros(miniBatchSize, nIn, timeSeriesLength);
            for (int i = 0; i < miniBatchSize; i++) {
                for (int j = 0; j < nIn; j++) {
                    for (int k = 0; k < timeSeriesLength; k++) {
                        input.putScalar(new int[] { i, j, k }, r.nextDouble() - 0.5);
                    }
                }
            }
            INDArray labels3d = Nd4j.zeros(miniBatchSize, nOut, timeSeriesLength);
            for (int i = 0; i < miniBatchSize; i++) {
                for (int j = 0; j < timeSeriesLength; j++) {
                    int idx = r.nextInt(nOut);
                    labels3d.putScalar(new int[] { i, idx, j }, 1.0f);
                }
            }
            INDArray labels2d = proc.backprop(labels3d, miniBatchSize, LayerWorkspaceMgr.noWorkspaces());
            MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().seed(12345L).list().layer(0, new GravesLSTM.Builder().nIn(nIn).nOut(layerSize).dist(new NormalDistribution(0, 1)).activation(Activation.TANH).updater(new NoOp()).build()).layer(1, new org.deeplearning4j.nn.conf.layers.OutputLayer.Builder(LossFunction.MCXENT).activation(Activation.SOFTMAX).nIn(layerSize).nOut(nOut).dist(new NormalDistribution(0, 1)).updater(new NoOp()).build()).inputPreProcessor(1, new RnnToFeedForwardPreProcessor()).build();
            MultiLayerNetwork mln = new MultiLayerNetwork(conf);
            mln.init();
            INDArray out2d = mln.feedForward(input).get(2);
            INDArray out3d = proc.preProcess(out2d, miniBatchSize, LayerWorkspaceMgr.noWorkspaces());
            MultiLayerConfiguration confRnn = new NeuralNetConfiguration.Builder().seed(12345L).list().layer(0, new GravesLSTM.Builder().nIn(nIn).nOut(layerSize).dist(new NormalDistribution(0, 1)).activation(Activation.TANH).updater(new NoOp()).build()).layer(1, new org.deeplearning4j.nn.conf.layers.RnnOutputLayer.Builder(LossFunction.MCXENT).activation(Activation.SOFTMAX).nIn(layerSize).nOut(nOut).dist(new NormalDistribution(0, 1)).updater(new NoOp()).build()).build();
            MultiLayerNetwork mlnRnn = new MultiLayerNetwork(confRnn);
            mlnRnn.init();
            INDArray outRnn = mlnRnn.feedForward(input).get(2);
            mln.setLabels(labels2d);
            mlnRnn.setLabels(labels3d);
            mln.computeGradientAndScore();
            mlnRnn.computeGradientAndScore();
            // score is average over all examples.
            // However: OutputLayer version has miniBatch*timeSeriesLength "examples" (after reshaping)
            // RnnOutputLayer has miniBatch examples
            // Hence: expect difference in scores by factor of timeSeriesLength
            double score = mln.score() * timeSeriesLength;
            double scoreRNN = mlnRnn.score();
            assertTrue(!Double.isNaN(score));
            assertTrue(!Double.isNaN(scoreRNN));
            double relError = Math.abs(score - scoreRNN) / (Math.abs(score) + Math.abs(scoreRNN));
            System.out.println(relError);
            assertTrue(relError < 1e-6);
            // Check labels and inputs for output layer:
            org.deeplearning4j.nn.layers.OutputLayer ol = (org.deeplearning4j.nn.layers.OutputLayer) mln.getOutputLayer();
            assertArrayEquals(ol.getInput().shape(), new long[] { miniBatchSize * timeSeriesLength, layerSize });
            assertArrayEquals(ol.getLabels().shape(), new long[] { miniBatchSize * timeSeriesLength, nOut });
            RnnOutputLayer rnnol = (RnnOutputLayer) mlnRnn.getOutputLayer();
            // assertArrayEquals(rnnol.getInput().shape(),new int[]{miniBatchSize,layerSize,timeSeriesLength});
            // Input may be set by BaseLayer methods. Thus input may end up as reshaped 2d version instead of original 3d version.
            // Not ideal, but everything else works.
            assertArrayEquals(rnnol.getLabels().shape(), new long[] { miniBatchSize, nOut, timeSeriesLength });
            // Check shapes of output for both:
            assertArrayEquals(out2d.shape(), new long[] { miniBatchSize * timeSeriesLength, nOut });
            INDArray out = mln.output(input);
            assertArrayEquals(out.shape(), new long[] { miniBatchSize * timeSeriesLength, nOut });
            INDArray preout = mln.output(input);
            assertArrayEquals(preout.shape(), new long[] { miniBatchSize * timeSeriesLength, nOut });
            INDArray outFFRnn = mlnRnn.feedForward(input).get(2);
            assertArrayEquals(outFFRnn.shape(), new long[] { miniBatchSize, nOut, timeSeriesLength });
            INDArray outRnn2 = mlnRnn.output(input);
            assertArrayEquals(outRnn2.shape(), new long[] { miniBatchSize, nOut, timeSeriesLength });
            INDArray preoutRnn = mlnRnn.output(input);
            assertArrayEquals(preoutRnn.shape(), new long[] { miniBatchSize, nOut, timeSeriesLength });
        }
    }

    @Test
    @DisplayName("Test Compare Rnn Output Rnn Loss")
    void testCompareRnnOutputRnnLoss() {
        Nd4j.getRandom().setSeed(12345);
        int timeSeriesLength = 4;
        int nIn = 5;
        int layerSize = 6;
        int nOut = 6;
        int miniBatchSize = 3;
        MultiLayerConfiguration conf1 = new NeuralNetConfiguration.Builder().seed(12345L).updater(new NoOp()).list().layer(new LSTM.Builder().nIn(nIn).nOut(layerSize).activation(Activation.TANH).dist(new NormalDistribution(0, 1.0)).updater(new NoOp()).build()).layer(new DenseLayer.Builder().nIn(layerSize).nOut(nOut).activation(Activation.IDENTITY).build()).layer(new RnnLossLayer.Builder(LossFunction.MCXENT).activation(Activation.SOFTMAX).build()).build();
        MultiLayerNetwork mln = new MultiLayerNetwork(conf1);
        mln.init();
        MultiLayerConfiguration conf2 = new NeuralNetConfiguration.Builder().seed(12345L).updater(new NoOp()).list().layer(new LSTM.Builder().nIn(nIn).nOut(layerSize).activation(Activation.TANH).dist(new NormalDistribution(0, 1.0)).updater(new NoOp()).build()).layer(new org.deeplearning4j.nn.conf.layers.RnnOutputLayer.Builder(LossFunction.MCXENT).activation(Activation.SOFTMAX).nIn(layerSize).nOut(nOut).build()).build();
        MultiLayerNetwork mln2 = new MultiLayerNetwork(conf2);
        mln2.init();
        mln2.setParams(mln.params());
        INDArray in = Nd4j.rand(new int[] { miniBatchSize, nIn, timeSeriesLength });
        INDArray out1 = mln.output(in);
        INDArray out2 = mln.output(in);
        assertEquals(out1, out2);
        Random r = new Random(12345);
        INDArray labels = Nd4j.create(miniBatchSize, nOut, timeSeriesLength);
        for (int i = 0; i < miniBatchSize; i++) {
            for (int j = 0; j < timeSeriesLength; j++) {
                labels.putScalar(i, r.nextInt(nOut), j, 1.0);
            }
        }
        mln.setInput(in);
        mln.setLabels(labels);
        mln2.setInput(in);
        mln2.setLabels(labels);
        mln.computeGradientAndScore();
        mln2.computeGradientAndScore();
        assertEquals(mln.gradient().gradient(), mln2.gradient().gradient());
        assertEquals(mln.score(), mln2.score(), 1e-6);
        TestUtils.testModelSerialization(mln);
    }

    @Test
    @DisplayName("Test Cnn Loss Layer")
    void testCnnLossLayer() {
        for (WorkspaceMode ws : WorkspaceMode.values()) {
            log.info("*** Testing workspace: " + ws);
            for (Activation a : new Activation[] { Activation.TANH, Activation.SELU }) {
                // Check that (A+identity) is equal to (identity+A), for activation A
                // i.e., should get same output and weight gradients for both
                MultiLayerConfiguration conf1 = new NeuralNetConfiguration.Builder().seed(12345L).updater(new NoOp()).convolutionMode(ConvolutionMode.Same).inferenceWorkspaceMode(ws).trainingWorkspaceMode(ws).list().layer(new ConvolutionLayer.Builder().nIn(3).nOut(4).activation(Activation.IDENTITY).kernelSize(2, 2).stride(1, 1).dist(new NormalDistribution(0, 1.0)).updater(new NoOp()).build()).layer(new CnnLossLayer.Builder(LossFunction.MSE).activation(a).build()).build();
                MultiLayerConfiguration conf2 = new NeuralNetConfiguration.Builder().seed(12345L).updater(new NoOp()).convolutionMode(ConvolutionMode.Same).inferenceWorkspaceMode(ws).trainingWorkspaceMode(ws).list().layer(new ConvolutionLayer.Builder().nIn(3).nOut(4).activation(a).kernelSize(2, 2).stride(1, 1).dist(new NormalDistribution(0, 1.0)).updater(new NoOp()).build()).layer(new CnnLossLayer.Builder(LossFunction.MSE).activation(Activation.IDENTITY).build()).build();
                MultiLayerNetwork mln = new MultiLayerNetwork(conf1);
                mln.init();
                MultiLayerNetwork mln2 = new MultiLayerNetwork(conf2);
                mln2.init();
                mln2.setParams(mln.params());
                INDArray in = Nd4j.rand(new int[] { 3, 3, 5, 5 });
                INDArray out1 = mln.output(in);
                INDArray out2 = mln2.output(in);
                assertEquals(out1, out2);
                INDArray labels = Nd4j.rand(out1.shape());
                mln.setInput(in);
                mln.setLabels(labels);
                mln2.setInput(in);
                mln2.setLabels(labels);
                mln.computeGradientAndScore();
                mln2.computeGradientAndScore();
                assertEquals(mln.score(), mln2.score(), 1e-6);
                assertEquals(mln.gradient().gradient(), mln2.gradient().gradient());
                // Also check computeScoreForExamples
                INDArray in2a = Nd4j.rand(new int[] { 1, 3, 5, 5 });
                INDArray labels2a = Nd4j.rand(new int[] { 1, 4, 5, 5 });
                INDArray in2 = Nd4j.concat(0, in2a, in2a);
                INDArray labels2 = Nd4j.concat(0, labels2a, labels2a);
                INDArray s = mln.scoreExamples(new DataSet(in2, labels2), false);
                assertArrayEquals(new long[] { 2, 1 }, s.shape());
                assertEquals(s.getDouble(0), s.getDouble(1), 1e-6);
                TestUtils.testModelSerialization(mln);
            }
        }
    }

    @Test
    @DisplayName("Test Cnn Loss Layer Comp Graph")
    void testCnnLossLayerCompGraph() {
        for (WorkspaceMode ws : WorkspaceMode.values()) {
            log.info("*** Testing workspace: " + ws);
            for (Activation a : new Activation[] { Activation.TANH, Activation.SELU }) {
                // Check that (A+identity) is equal to (identity+A), for activation A
                // i.e., should get same output and weight gradients for both
                ComputationGraphConfiguration conf1 = new NeuralNetConfiguration.Builder().seed(12345L).updater(new NoOp()).convolutionMode(ConvolutionMode.Same).inferenceWorkspaceMode(ws).trainingWorkspaceMode(ws).graphBuilder().addInputs("in").addLayer("0", new ConvolutionLayer.Builder().nIn(3).nOut(4).activation(Activation.IDENTITY).kernelSize(2, 2).stride(1, 1).dist(new NormalDistribution(0, 1.0)).updater(new NoOp()).build(), "in").addLayer("1", new CnnLossLayer.Builder(LossFunction.MSE).activation(a).build(), "0").setOutputs("1").build();
                ComputationGraphConfiguration conf2 = new NeuralNetConfiguration.Builder().seed(12345L).updater(new NoOp()).convolutionMode(ConvolutionMode.Same).inferenceWorkspaceMode(ws).trainingWorkspaceMode(ws).graphBuilder().addInputs("in").addLayer("0", new ConvolutionLayer.Builder().nIn(3).nOut(4).activation(a).kernelSize(2, 2).stride(1, 1).dist(new NormalDistribution(0, 1.0)).updater(new NoOp()).build(), "in").addLayer("1", new CnnLossLayer.Builder(LossFunction.MSE).activation(Activation.IDENTITY).build(), "0").setOutputs("1").build();
                ComputationGraph graph = new ComputationGraph(conf1);
                graph.init();
                ComputationGraph graph2 = new ComputationGraph(conf2);
                graph2.init();
                graph2.setParams(graph.params());
                INDArray in = Nd4j.rand(new int[] { 3, 3, 5, 5 });
                INDArray out1 = graph.outputSingle(in);
                INDArray out2 = graph2.outputSingle(in);
                assertEquals(out1, out2);
                INDArray labels = Nd4j.rand(out1.shape());
                graph.setInput(0, in);
                graph.setLabels(labels);
                graph2.setInput(0, in);
                graph2.setLabels(labels);
                graph.computeGradientAndScore();
                graph2.computeGradientAndScore();
                assertEquals(graph.score(), graph2.score(), 1e-6);
                assertEquals(graph.gradient().gradient(), graph2.gradient().gradient());
                // Also check computeScoreForExamples
                INDArray in2a = Nd4j.rand(new int[] { 1, 3, 5, 5 });
                INDArray labels2a = Nd4j.rand(new int[] { 1, 4, 5, 5 });
                INDArray in2 = Nd4j.concat(0, in2a, in2a);
                INDArray labels2 = Nd4j.concat(0, labels2a, labels2a);
                INDArray s = graph.scoreExamples(new DataSet(in2, labels2), false);
                assertArrayEquals(new long[] { 2, 1 }, s.shape());
                assertEquals(s.getDouble(0), s.getDouble(1), 1e-6);
                TestUtils.testModelSerialization(graph);
            }
        }
    }

    @Test
    @DisplayName("Test Cnn Output Layer Softmax")
    void testCnnOutputLayerSoftmax() {
        // Check that softmax is applied channels-wise
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().seed(12345L).updater(new NoOp()).convolutionMode(ConvolutionMode.Same).list().layer(new ConvolutionLayer.Builder().nIn(3).nOut(4).activation(Activation.IDENTITY).dist(new NormalDistribution(0, 1.0)).updater(new NoOp()).build()).layer(new CnnLossLayer.Builder(LossFunction.MSE).activation(Activation.SOFTMAX).build()).build();
        MultiLayerNetwork net = new MultiLayerNetwork(conf);
        net.init();
        INDArray in = Nd4j.rand(new int[] { 2, 3, 4, 5 });
        INDArray out = net.output(in);
        double min = out.minNumber().doubleValue();
        double max = out.maxNumber().doubleValue();
        assertTrue(min >= 0 && max <= 1.0);
        INDArray sum = out.sum(1);
        assertEquals(Nd4j.ones(DataType.FLOAT, 2, 4, 5), sum);
    }

    @Test
    @DisplayName("Test Output Layer Defaults")
    void testOutputLayerDefaults() {
        new NeuralNetConfiguration.Builder().list().layer(new org.deeplearning4j.nn.conf.layers.OutputLayer.Builder().nIn(10).nOut(10).build()).build();
        new NeuralNetConfiguration.Builder().list().layer(new org.deeplearning4j.nn.conf.layers.LossLayer.Builder().build()).build();
        new NeuralNetConfiguration.Builder().list().layer(new CnnLossLayer.Builder().build()).build();
        new NeuralNetConfiguration.Builder().list().layer(new CenterLossOutputLayer.Builder().build()).build();
    }
}
