using namespace mlpack;
using namespace mlpack::ann;

BOOST_AUTO_TEST_SUITE(ConvolutionalNetworkTest);


template<typename Net>
void predict_digits(Net &net, arma::cube const &input,
                    size_t begin, size_t end,
                    size_t activate_node)
{
    arma::cube predict_data(input.n_rows, input.n_cols, 1);
    arma::mat label;
    double sum = 0;
    for(size_t i = begin; i != end; ++i)
    {
        predict_data.slice(0) = input.slice(i);
        net.Predict(predict_data, label);
        if(label(activate_node) == 1){
            ++sum;
        }
    }

    std::cout<<"predict accuracy of "<<activate_node<<" : "
            <<(sum/(end - begin))<<"\n";
}

/**
 * Train and evaluate a vanilla network with the specified structure.
 */
template<
        typename PerformanceFunction
        >
void BuildVanillaNetwork()
{
    arma::mat X;
    X.load("mnist_first250_training_4s_and_9s.arm");

    // Normalize each point since these are images.
    arma::uword nPoints = X.n_cols;
    for (arma::uword i = 0; i < nPoints; i++)
    {
        X.col(i) /= norm(X.col(i), 2);
    }

    // Build the target matrix.
    arma::mat Y = arma::zeros<arma::mat>(10, nPoints);
    for (size_t i = 0; i < nPoints; i++)
    {
        if (i < nPoints / 2)
        {
            Y.col(i)(5) = 1;
        }
        else
        {
            Y.col(i)(8) = 1;
        }
    }

    arma::cube input = arma::cube(28, 28, nPoints);
    for (size_t i = 0; i < nPoints; i++)
        input.slice(i) = arma::mat(X.colptr(i), 28, 28);

    /*
   * Construct a convolutional neural network with a 28x28x1 input layer,
   * 24x24x8 convolution layer, 12x12x8 pooling layer, 8x8x12 convolution layer
   * and a 4x4x12 pooling layer which is fully connected with the output layer.
   * The network structure looks like:
   *
   * Input    Convolution  Pooling      Convolution  Pooling      Output
   * Layer    Layer        Layer        Layer        Layer        Layer
   *
   *          +---+        +---+        +---+        +---+
   *          | +---+      | +---+      | +---+      | +---+
   * +---+    | | +---+    | | +---+    | | +---+    | | +---+    +---+
   * |   |    | | |   |    | | |   |    | | |   |    | | |   |    |   |
   * |   +--> +-+ |   +--> +-+ |   +--> +-+ |   +--> +-+ |   +--> |   |
   * |   |      +-+   |      +-+   |      +-+   |      +-+   |    |   |
   * +---+        +---+        +---+        +---+        +---+    +---+
   */

    ConvLayer<RMSPROP> convLayer0(1, 8, 5, 5);
    BiasLayer2D<RMSPROP, ZeroInitialization> biasLayer0(8);
    BaseLayer2D<PerformanceFunction> baseLayer0;
    PoolingLayer<> poolingLayer0(2);




    ConvLayer<RMSPROP> convLayer1(8, 12, 5, 5);
    BiasLayer2D<RMSPROP, ZeroInitialization> biasLayer1(12);
    BaseLayer2D<PerformanceFunction> baseLayer1;
    PoolingLayer<> poolingLayer1(2);

    LinearMappingLayer<RMSPROP> linearLayer0(192, 10);
    BiasLayer<RMSPROP> biasLayer2(10);
    SoftmaxLayer<> softmaxLayer0;

    OneHotLayer outputLayer;

    auto modules = std::tie(convLayer0, biasLayer0, baseLayer0, poolingLayer0,
                            convLayer1, biasLayer1, baseLayer1, poolingLayer1,
                            linearLayer0, biasLayer2, softmaxLayer0);

    CNN<decltype(modules), decltype(outputLayer)>
            net(modules, outputLayer);

    Trainer<decltype(net)> trainer(net, 50, 1, 0.7);
    trainer.Train(input, Y, input, Y);

    predict_digits(net, input, 0, nPoints/2, 5);
    predict_digits(net, input, nPoints/2, nPoints, 8);

    BOOST_REQUIRE_LE(trainer.ValidationError(), 0.7);
}

/**
 * Train the vanilla network on a larger dataset.
 */
BOOST_AUTO_TEST_CASE(VanillaNetworkTest)
{
    BuildVanillaNetwork<LogisticFunction>();
}

/**
 * Train and evaluate a vanilla network with the specified structure.
 */
template<
        typename PerformanceFunction
        >
void BuildVanillaDropoutNetwork()
{
    arma::mat X;
    X.load("mnist_first250_training_4s_and_9s.arm");

    // Normalize each point since these are images.
    arma::uword nPoints = X.n_cols;
    for (arma::uword i = 0; i < nPoints; i++)
    {
        X.col(i) /= norm(X.col(i), 2);
    }

    // Build the target matrix.
    arma::mat Y = arma::zeros<arma::mat>(10, nPoints);
    for (size_t i = 0; i < nPoints; i++)
    {
        if (i < nPoints / 2)
        {
            Y.col(i)(5) = 1;
        }
        else
        {
            Y.col(i)(8) = 1;
        }
    }

    arma::cube input = arma::cube(28, 28, nPoints);
    for (size_t i = 0; i < nPoints; i++)
        input.slice(i) = arma::mat(X.colptr(i), 28, 28);

    /*
   * Construct a convolutional neural network with a 28x28x1 input layer,
   * 24x24x4 convolution layer, 24x24x4 dropout layer, 12x12x4 pooling layer,
   * 8x8x8 convolution layer,8x8x8 Dropout Layer and a 4x4x12 pooling layer
   * which is fully connected with the output layer. The network structure looks
   * like:
   *
   * Input    Convolution  Dropout      Pooling     Convolution,     Output
   * Layer    Layer        Layer        Layer       Dropout,         Layer
   *                                                Pooling Layer
   *          +---+        +---+        +---+
   *          | +---+      | +---+      | +---+
   * +---+    | | +---+    | | +---+    | | +---+                    +---+
   * |   |    | | |   |    | | |   |    | | |   |                    |   |
   * |   +--> +-+ |   +--> +-+ |   +--> +-+ |   +--> ............--> |   |
   * |   |      +-+   |      +-+   |      +-+   |                    |   |
   * +---+        +---+        +---+        +---+                    +---+
   */

    ConvLayer<AdaDelta> convLayer0(1, 4, 5, 5);
    BiasLayer2D<AdaDelta, ZeroInitialization> biasLayer0(4);
    DropoutLayer2D<> dropoutLayer0;
    BaseLayer2D<PerformanceFunction> baseLayer0;
    PoolingLayer<> poolingLayer0(2);

    ConvLayer<AdaDelta> convLayer1(4, 8, 5, 5);
    BiasLayer2D<AdaDelta, ZeroInitialization> biasLayer1(8);
    DropoutLayer2D<> dropoutLayer1;
    BaseLayer2D<PerformanceFunction> baseLayer1;
    PoolingLayer<> poolingLayer1(2);

    LinearMappingLayer<AdaDelta> linearLayer0(128, 10);
    BiasLayer<AdaDelta> biasLayer2(10);
    SoftmaxLayer<> softmaxLayer0;

    OneHotLayer outputLayer;

    auto modules =
            std::tie(convLayer0, biasLayer0, dropoutLayer0, baseLayer0, poolingLayer0,
                     convLayer1, biasLayer1, dropoutLayer1, baseLayer1, poolingLayer1,
                     linearLayer0, biasLayer2, softmaxLayer0);

    CNN<decltype(modules), decltype(outputLayer)>
            net(modules, outputLayer);

    Trainer<decltype(net)> trainer(net, 50, 1, 0.7);
    trainer.Train(input, Y, input, Y);

    predict_digits(net, input, 0, nPoints/2, 5);
    predict_digits(net, input, nPoints/2, nPoints, 8);

    BOOST_REQUIRE_LE(trainer.ValidationError(), 0.7);
}

/**
 * Train the network on a larger dataset using dropout.
 */
BOOST_AUTO_TEST_CASE(VanillaNetworkDropoutTest)
{
    BuildVanillaDropoutNetwork<RectifierFunction>();
}

BOOST_AUTO_TEST_SUITE_END();