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

import org.deeplearning4j.BaseDL4JTest;
import org.deeplearning4j.nn.conf.GradientNormalization;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.dropout.Dropout;
import org.deeplearning4j.nn.conf.layers.*;
import org.deeplearning4j.nn.weights.WeightInit;
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.activations.IActivation;
import org.nd4j.linalg.activations.impl.ActivationSoftmax;
import org.nd4j.linalg.activations.impl.ActivationTanH;
import org.nd4j.linalg.convolution.Convolution;
import org.nd4j.linalg.learning.config.AdaGrad;
import org.nd4j.linalg.learning.config.IUpdater;
import org.nd4j.linalg.lossfunctions.LossFunctions.LossFunction;
import java.io.*;
import static org.junit.jupiter.api.Assertions.*;
import org.junit.jupiter.api.DisplayName;

/**
 * @author Jeffrey Tang.
 */
@DisplayName("Layer Builder Test")
@NativeTag
@Tag(TagNames.DL4J_OLD_API)
class LayerBuilderTest extends BaseDL4JTest {

    final double DELTA = 1e-15;

    int numIn = 10;

    int numOut = 5;

    double drop = 0.3;

    IActivation act = new ActivationSoftmax();

    PoolingType poolType = PoolingType.MAX;

    int[] kernelSize = new int[] { 2, 2 };

    int[] stride = new int[] { 2, 2 };

    int[] padding = new int[] { 1, 1 };

    int k = 1;

    Convolution.Type convType = Convolution.Type.VALID;

    LossFunction loss = LossFunction.MCXENT;

    WeightInit weight = WeightInit.XAVIER;

    double corrupt = 0.4;

    double sparsity = 0.3;

    double corruptionLevel = 0.5;

    double dropOut = 0.1;

    IUpdater updater = new AdaGrad();

    GradientNormalization gradNorm = GradientNormalization.ClipL2PerParamType;

    double gradNormThreshold = 8;

    @Test
    @DisplayName("Test Layer")
    void testLayer() throws Exception {
        DenseLayer layer = new DenseLayer.Builder().activation(act).weightInit(weight).dropOut(dropOut).updater(updater).gradientNormalization(gradNorm).gradientNormalizationThreshold(gradNormThreshold).build();
        checkSerialization(layer);
        assertEquals(act, layer.getActivationFn());
        assertEquals(weight.getWeightInitFunction(), layer.getWeightInitFn());
        assertEquals(new Dropout(dropOut), layer.getIDropout());
        assertEquals(updater, layer.getIUpdater());
        assertEquals(gradNorm, layer.getGradientNormalization());
        assertEquals(gradNormThreshold, layer.getGradientNormalizationThreshold(), 0.0);
    }

    @Test
    @DisplayName("Test Feed Forward Layer")
    void testFeedForwardLayer() throws Exception {
        DenseLayer ff = new DenseLayer.Builder().nIn(numIn).nOut(numOut).build();
        checkSerialization(ff);
        assertEquals(numIn, ff.getNIn());
        assertEquals(numOut, ff.getNOut());
    }

    @Test
    @DisplayName("Test Convolution Layer")
    void testConvolutionLayer() throws Exception {
        ConvolutionLayer conv = new ConvolutionLayer.Builder(kernelSize, stride, padding).build();
        checkSerialization(conv);
        // assertEquals(convType, conv.getConvolutionType());
        assertArrayEquals(kernelSize, conv.getKernelSize());
        assertArrayEquals(stride, conv.getStride());
        assertArrayEquals(padding, conv.getPadding());
    }

    @Test
    @DisplayName("Test Subsampling Layer")
    void testSubsamplingLayer() throws Exception {
        SubsamplingLayer sample = new SubsamplingLayer.Builder(poolType, stride).kernelSize(kernelSize).padding(padding).build();
        checkSerialization(sample);
        assertArrayEquals(padding, sample.getPadding());
        assertArrayEquals(kernelSize, sample.getKernelSize());
        assertEquals(poolType, sample.getPoolingType());
        assertArrayEquals(stride, sample.getStride());
    }

    @Test
    @DisplayName("Test Output Layer")
    void testOutputLayer() throws Exception {
        OutputLayer out = new OutputLayer.Builder(loss).build();
        checkSerialization(out);
    }

    @Test
    @DisplayName("Test Rnn Output Layer")
    void testRnnOutputLayer() throws Exception {
        RnnOutputLayer out = new RnnOutputLayer.Builder(loss).build();
        checkSerialization(out);
    }

    @Test
    @DisplayName("Test Auto Encoder")
    void testAutoEncoder() throws Exception {
        AutoEncoder enc = new AutoEncoder.Builder().corruptionLevel(corruptionLevel).sparsity(sparsity).build();
        checkSerialization(enc);
        assertEquals(corruptionLevel, enc.getCorruptionLevel(), DELTA);
        assertEquals(sparsity, enc.getSparsity(), DELTA);
    }

    @Test
    @DisplayName("Test Graves LSTM")
    void testGravesLSTM() throws Exception {
        GravesLSTM glstm = new GravesLSTM.Builder().forgetGateBiasInit(1.5).activation(Activation.TANH).nIn(numIn).nOut(numOut).build();
        checkSerialization(glstm);
        assertEquals(glstm.getForgetGateBiasInit(), 1.5, 0.0);
        assertEquals(glstm.getNIn(), numIn);
        assertEquals(glstm.getNOut(), numOut);
        assertTrue(glstm.getActivationFn() instanceof ActivationTanH);
    }

    @Test
    @DisplayName("Test Graves Bidirectional LSTM")
    void testGravesBidirectionalLSTM() throws Exception {
        final GravesBidirectionalLSTM glstm = new GravesBidirectionalLSTM.Builder().forgetGateBiasInit(1.5).activation(Activation.TANH).nIn(numIn).nOut(numOut).build();
        checkSerialization(glstm);
        assertEquals(1.5, glstm.getForgetGateBiasInit(), 0.0);
        assertEquals(glstm.getNIn(), numIn);
        assertEquals(glstm.getNOut(), numOut);
        assertTrue(glstm.getActivationFn() instanceof ActivationTanH);
    }

    @Test
    @DisplayName("Test Embedding Layer")
    void testEmbeddingLayer() throws Exception {
        EmbeddingLayer el = new EmbeddingLayer.Builder().nIn(10).nOut(5).build();
        checkSerialization(el);
        assertEquals(10, el.getNIn());
        assertEquals(5, el.getNOut());
    }

    @Test
    @DisplayName("Test Batch Norm Layer")
    void testBatchNormLayer() throws Exception {
        BatchNormalization bN = new BatchNormalization.Builder().nIn(numIn).nOut(numOut).gamma(2).beta(1).decay(0.5).lockGammaBeta(true).build();
        checkSerialization(bN);
        assertEquals(numIn, bN.getNIn());
        assertEquals(numOut, bN.getNOut());
        assertEquals(true, bN.isLockGammaBeta());
        assertEquals(0.5, bN.getDecay(), 1e-4);
        assertEquals(2, bN.getGamma(), 1e-4);
        assertEquals(1, bN.getBeta(), 1e-4);
    }

    @Test
    @DisplayName("Test Activation Layer")
    void testActivationLayer() throws Exception {
        ActivationLayer activationLayer = new ActivationLayer.Builder().activation(act).build();
        checkSerialization(activationLayer);
        assertEquals(act, activationLayer.getActivationFn());
    }

    private void checkSerialization(Layer layer) throws Exception {
        NeuralNetConfiguration confExpected = new NeuralNetConfiguration.Builder().layer(layer).build();
        NeuralNetConfiguration confActual;
        // check Java serialization
        byte[] data;
        try (ByteArrayOutputStream bos = new ByteArrayOutputStream();
            ObjectOutput out = new ObjectOutputStream(bos)) {
            out.writeObject(confExpected);
            data = bos.toByteArray();
        }
        try (ByteArrayInputStream bis = new ByteArrayInputStream(data);
            ObjectInput in = new ObjectInputStream(bis)) {
            confActual = (NeuralNetConfiguration) in.readObject();
        }
        assertEquals(confExpected.getLayer(), confActual.getLayer(), "unequal Java serialization");
        // check JSON
        String json = confExpected.toJson();
        confActual = NeuralNetConfiguration.fromJson(json);
        assertEquals(confExpected.getLayer(), confActual.getLayer(), "unequal JSON serialization");
        // check YAML
        String yaml = confExpected.toYaml();
        confActual = NeuralNetConfiguration.fromYaml(yaml);
        assertEquals(confExpected.getLayer(), confActual.getLayer(), "unequal YAML serialization");
        // check the layer's use of callSuper on equals method
        confActual.getLayer().setIDropout(new Dropout(new java.util.Random().nextDouble()));
        assertNotEquals(confExpected.getLayer(), confActual.getLayer(), "broken equals method (missing callSuper?)");
    }
}
