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

import org.deeplearning4j.BaseDL4JTest;
import org.deeplearning4j.datasets.iterator.utilty.ListDataSetIterator;
import org.deeplearning4j.exception.DL4JInvalidConfigException;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.CNN2DFormat;
import org.deeplearning4j.nn.conf.ComputationGraphConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.inputs.InputType;
import org.deeplearning4j.nn.conf.layers.*;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.nn.weights.WeightInit;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Disabled;
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.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.factory.Nd4j;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import static org.junit.jupiter.api.Assertions.*;


//@Disabled
@NativeTag
@Tag(TagNames.DL4J_OLD_API)
public class TestCompGraphCNN extends BaseDL4JTest {

    protected ComputationGraphConfiguration conf;
    protected ComputationGraph graph;
    protected DataSetIterator dataSetIterator;
    protected DataSet ds;

    protected static ComputationGraphConfiguration getMultiInputGraphConfig() {
        ComputationGraphConfiguration conf =
                new NeuralNetConfiguration.Builder()
                        .optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT)
                        .graphBuilder().addInputs("input")
                        .setInputTypes(InputType.convolutional(32, 32, 3))
                        .addLayer("cnn1",
                                new ConvolutionLayer.Builder(4, 4).stride(2, 2).nIn(3).nOut(3)
                                        .build(),
                                "input")
                        .addLayer("cnn2",
                                new ConvolutionLayer.Builder(4, 4).stride(2, 2).nIn(3).nOut(3)
                                        .build(),
                                "input")
                        .addLayer("max1",
                                new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX)
                                        .stride(1, 1).kernelSize(2, 2).build(),
                                "cnn1", "cnn2")
                        .addLayer("dnn1", new DenseLayer.Builder().nOut(7).build(), "max1")
                        .addLayer("output", new OutputLayer.Builder().nIn(7).nOut(10).activation(Activation.SOFTMAX).build(), "dnn1")
                        .setOutputs("output").build();

        return conf;
    }

    protected static DataSetIterator getDS() {

        List<DataSet> list = new ArrayList<>(5);
        for (int i = 0; i < 5; i++) {
            INDArray f = Nd4j.create(1, 32 * 32 * 3);
            INDArray l = Nd4j.create(1, 10);
            l.putScalar(i, 1.0);
            list.add(new DataSet(f, l));
        }
        return new ListDataSetIterator(list, 5);
    }

    protected static int getNumParams() {
        return 2 * (3 * 1 * 4 * 4 * 3 + 3) + (7 * 14 * 14 * 6 + 7) + (7 * 10 + 10);
    }

    @BeforeEach
    @Disabled
    public void beforeDo() {
        conf = getMultiInputGraphConfig();
        graph = new ComputationGraph(conf);
        graph.init();

        dataSetIterator = getDS();
        ds = dataSetIterator.next();

    }

    @Test
    public void testConfigBasic() {
        //Check the order. there are 2 possible valid orders here
        int[] order = graph.topologicalSortOrder();
        int[] expOrder1 = new int[] {0, 1, 2, 4, 3, 5, 6}; //First of 2 possible valid orders
        int[] expOrder2 = new int[] {0, 2, 1, 4, 3, 5, 6}; //Second of 2 possible valid orders
        boolean orderOK = Arrays.equals(expOrder1, order) || Arrays.equals(expOrder2, order);
        assertTrue(orderOK);

        INDArray params = graph.params();
        assertNotNull(params);

        // confirm param shape is what is expected
        int nParams = getNumParams();
        assertEquals(nParams, params.length());

        INDArray arr = Nd4j.linspace(0, nParams, nParams, DataType.FLOAT).reshape(nParams);
        assertEquals(nParams, arr.length());

        // params are set
        graph.setParams(arr);
        params = graph.params();
        assertEquals(arr, params);

        //Number of inputs and outputs:
        assertEquals(1, graph.getNumInputArrays());
        assertEquals(1, graph.getNumOutputArrays());

    }

    @Test()
    public void testCNNComputationGraphKernelTooLarge() {
       assertThrows(DL4JInvalidConfigException.class,() -> {
           int imageWidth = 23;
           int imageHeight = 19;
           int nChannels = 1;
           int classes = 2;
           int numSamples = 200;

           int kernelHeight = 3;
           int kernelWidth = imageWidth;


           DataSet trainInput;

           ComputationGraphConfiguration conf =
                   new NeuralNetConfiguration.Builder()
                           .optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT)
                           .seed(123).graphBuilder().addInputs("input")
                           .setInputTypes(InputType.convolutional(nChannels, imageWidth,
                                   imageHeight))
                           .addLayer("conv1", new ConvolutionLayer.Builder()
                                   .kernelSize(kernelHeight, kernelWidth).stride(1, 1)
                                   .dataFormat(CNN2DFormat.NCHW)
                                   .nIn(nChannels).nOut(2).weightInit(WeightInit.XAVIER)
                                   .activation(Activation.RELU).build(), "input")
                           .addLayer("pool1",
                                   new SubsamplingLayer.Builder()
                                           .dataFormat(CNN2DFormat.NCHW)
                                           .poolingType(SubsamplingLayer.PoolingType.MAX)
                                           .kernelSize(imageHeight - kernelHeight + 1, 1)
                                           .stride(1, 1).build(),
                                   "conv1")
                           .addLayer("output", new OutputLayer.Builder().nOut(classes).activation(Activation.SOFTMAX).build(), "pool1")
                           .setOutputs("output").build();


           ComputationGraph model = new ComputationGraph(conf);
           model.init();


           INDArray emptyFeatures = Nd4j.zeros(numSamples, imageWidth * imageHeight * nChannels);
           INDArray emptyLables = Nd4j.zeros(numSamples, classes);

           trainInput = new DataSet(emptyFeatures, emptyLables);

           model.fit(trainInput);
       });

    }
}
