/*
 *  ******************************************************************************
 *  *
 *  *
 *  * 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.deeplearning4j.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.weights.WeightInit;
import org.junit.jupiter.api.Test;
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.*;

/**
 * @author Jeffrey Tang.
 */
public 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
    public 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
    public 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
    public 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
    public 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
    public void testOutputLayer() throws Exception {
        OutputLayer out = new OutputLayer.Builder(loss).build();

        checkSerialization(out);
    }

    @Test
    public void testRnnOutputLayer() throws Exception {
        RnnOutputLayer out = new RnnOutputLayer.Builder(loss).build();

        checkSerialization(out);
    }

    @Test
    public 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
    public 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.nIn, numIn);
        assertEquals(glstm.nOut, numOut);
        assertTrue(glstm.getActivationFn() instanceof ActivationTanH);
    }

    @Test
    public 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.nIn, numIn);
        assertEquals(glstm.nOut, numOut);
        assertTrue(glstm.getActivationFn() instanceof ActivationTanH);
    }

    @Test
    public 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
    public 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.nIn);
        assertEquals(numOut, bN.nOut);
        assertEquals(true, bN.isLockGammaBeta());
        assertEquals(0.5, bN.decay, 1e-4);
        assertEquals(2, bN.gamma, 1e-4);
        assertEquals(1, bN.beta, 1e-4);
    }

    @Test
    public void testActivationLayer() throws Exception {
        ActivationLayer activationLayer = new ActivationLayer.Builder().activation(act).build();

        checkSerialization(activationLayer);

        assertEquals(act, activationLayer.activationFn);
    }

    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?)");
    }

}
