/********************************************************************************
 * Copyright (c) 2023 CEA-List
 *
 * This program and the accompanying materials are made available under the
 * terms of the Eclipse Public License 2.0 which is available at
 * http://www.eclipse.org/legal/epl-2.0.
 *
 * SPDX-License-Identifier: EPL-2.0
 *
 ********************************************************************************/

#include <catch2/catch_test_macros.hpp>
#include <memory>
#include <cstdlib>

#include "aidge/data/Tensor.hpp"
#include "aidge/operator/MaxPooling.hpp"

#include "aidge/backend/cpu.hpp"

using namespace Aidge;


TEST_CASE("[cpu/operator] MaxPooling(forward)", "[MaxPooling][CPU]") {
    std::shared_ptr<Tensor> myInput = std::make_shared<Tensor>(Array4D<float,2,2,5,5> { //NCHW
        {
            {
                {{-0.3848,  0.2166, -0.4373,  0.6142,  0.5277},
                 {0.7995,  0.3638, -1.4589, -1.0843,  1.0918},
            	 {0.7147,  0.0936, -1.2902,  1.2037,  0.4874},
                 {-0.5981,  2.1184, -0.9175,  1.3859,  0.3305},
                 {-1.7700,  0.0563, -0.3914,  0.0538, -0.3955}},

                {{-3.1409, -0.4554,  0.0524,  2.2291,  0.4859},
                 {-0.7465, -0.6567, -2.3703, -0.6386, -1.4152},
                 { 2.2329, -0.5850,  0.0700,  1.2838, -1.7363},
                 { 0.2139,  0.0624, -1.0689, -0.8221, -0.8038},
                 { 0.1886, -0.7840, -0.2313,  0.2651, -1.6244}}
            },
            {
                {{ 0.4371,  1.6417,  0.9129,  0.6325,  0.5438},
                 {-2.3552, -0.8850, -0.0232, -0.5462, -1.2011},
                 {1.7653, -1.6668, -1.0814,  0.6182,  1.2071},
                 {0.9541, -0.5133,  0.8664, -0.8892,  1.4585},
                 {1.0220, -0.5107,  0.1829, -0.2301, -0.4268}},

                {{ 1.0429,  0.6279, -0.2875,  0.7187, -0.1500},
                 {1.6041,  2.9635,  1.4172, -0.7517,  0.5441},
                 {-0.2276,  0.0857,  0.6776, -0.1389, -0.0614},
                 {-0.1547, -0.3435,  0.0650, -0.5095, -1.8073},
                 {1.7217,  0.3999, -0.5953,  1.0604, -0.4126}}
            }
        }
    });
    SECTION("Stride") {
        std::shared_ptr<Node> myMaxPool = MaxPooling({2,2}, "mycdw", {2,2});
        auto op = std::static_pointer_cast<OperatorTensor>(myMaxPool -> getOperator());

        std::shared_ptr<Tensor> myOutput = std::make_shared<Tensor>(Array4D<float,2,2,2,2> {
            {
                {
                    {{  0.7995,  0.6142},
                     { 2.1184,  1.3859}},
                    {{ -0.4554,  2.2291},
                     {  2.2329,  1.2838}}
                },
                {
                    {{1.6417,  0.9129},
                     {1.7653,  0.8664}},
                    {{2.9635,  1.4172},
                     {0.0857,  0.6776}}
                }
            }
        });
        myMaxPool->getOperator()->associateInput(0,myInput);
        myMaxPool->getOperator()->setDataType(DataType::Float32);
        myMaxPool->getOperator()->setBackend("cpu");
        op->computeOutputDims();
        myMaxPool->forward();
        op->getOutput(0)->print();
        REQUIRE(*(op->getOutput(0)) == *myOutput);
    }
}