package org.deeplearning4j.nn.layers.convolution;

import org.deeplearning4j.datasets.iterator.impl.MnistDataSetIterator;
import org.deeplearning4j.nn.api.Layer;
import org.deeplearning4j.nn.api.LayerFactory;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.*;
import org.deeplearning4j.nn.conf.layers.ConvolutionLayer;
import org.deeplearning4j.nn.conf.override.ConfOverride;
import org.deeplearning4j.nn.layers.OutputLayer;
import org.deeplearning4j.nn.layers.factory.LayerFactories;
import org.deeplearning4j.nn.layers.feedforward.rbm.RBM;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.junit.Test;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.convolution.Convolution;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.lossfunctions.LossFunctions;


import static org.junit.Assert.*;

/**
 * @author Adam Gibson
 */
public class TestConvolutionLayer {

    @Test
    public void testNumExamplesMatch() throws Exception  {
        DataSetIterator mnist = new MnistDataSetIterator(10,10);
        DataSet next = mnist.next();

        Layer layer = getCNNConfig(1, 2, 9, 9, new int[] {3,3}, new int[] {1,1});

        INDArray input = next.getFeatureMatrix().reshape(next.numExamples(),1,28,28);
        INDArray conv = layer.activate(input);

        assertEquals(input.slices(), conv.slices());

    }


    @Test
    public void testFeatureMapSpaceSize() throws Exception  {
        int inputWidth = 28;
        int inputHeight = 28;
        int kernelWidth = 9;
        int kernelHeight = 9;
        int[] stride = new int[] {3,3};
        int[] padding = new int[] {1,1};
        int nIn = 1;
        int nOut = 2;

        DataSetIterator mnist = new MnistDataSetIterator(10, 10);
        DataSet next = mnist.next();

        Layer layer = getCNNConfig(nIn, nOut, kernelWidth, kernelHeight, stride, padding);

        INDArray input = next.getFeatureMatrix().reshape(next.numExamples(),1, inputWidth, inputHeight);
        INDArray conv = layer.activate(input);

        int featureMapWidth = (inputWidth - kernelWidth) + (2*padding[0]) / stride[0] + 1;
        assertEquals(featureMapWidth, conv.shape()[0]);
        assertEquals(nOut, conv.shape()[2]);

    }



    private static Layer getCNNConfig(int nIn, int nOut, int kernelWidth, int kernelHeight, int[] stride, int[] padding){
        NeuralNetConfiguration conf = new NeuralNetConfiguration.Builder()
                .activationFunction("relu")
                .iterations(1)
                .stride(stride)
                .padding(padding)
                .layer(new ConvolutionLayer.Builder(new int[]{kernelWidth, kernelHeight}, Convolution.Type.SAME)
                        .nIn(nIn)
                        .nOut(nOut)
                        .build())
                .build();
        return LayerFactories.getFactory(new ConvolutionLayer()).create(conf);

    }

}

@@ -1905,19 +1905,19 @@
  int%5B%5D %7B
-3,3
+1,1
 %7D;%0A     
@@ -1938,35 +1938,35 @@
 ng = new int%5B%5D %7B
-1,1
+0,0
 %7D;%0A        int n
@@ -1992,17 +1992,18 @@
  nOut = 
-2
+10
 ;%0A%0A     
@@ -2350,24 +2350,26 @@
 );%0A%0A        
+//
 int featureM
@@ -2412,10 +2412,9 @@
  + (
-2*
+ 
 padd
@@ -2434,24 +2434,58 @@
 ide%5B0%5D + 1;%0A
+        int featureMapWidth = 10;%0A
         asse
@@ -2520,17 +2520,14 @@
 nv.s
-hape()%5B0%5D
+ize(0)
 );%0A 
@@ -2562,17 +2562,14 @@
 nv.s
-hape()%5B2%5D
+ize(0)
 );%0A%0A

update test