/********************************************************************************
 * 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 "aidge/data/Tensor.hpp"
#include "aidge/data/TensorImpl.hpp"
#include "aidge_cpu.hpp"
#include "aidge/operator/Add.hpp"

using namespace Aidge;

TEST_CASE("[cpu/operator] Add(forward)") {
    std::shared_ptr<Tensor> input1 = std::make_shared<Tensor>(Array4D<int,3,3,3,2> { 
        {
            {
                {{20, 47},{21, 48},{22, 49}},
                {{23, 50},{24, 51},{25, 52}},
                {{26, 53},{27, 54},{28, 55}}
            },
            {
                {{29, 56},{30, 57},{31, 58}},
                {{32, 59},{33, 60},{34, 61}},
                {{35, 62},{36, 63},{37, 64}}
            },
            {
                {{38, 65},{39, 66},{40, 67}},
                {{41, 68},{42, 69},{43, 70}},
                {{44, 71},{45, 72},{46, 73}}
            }
        }
    });

    SECTION("One input") {
        std::shared_ptr<Node> myAdd = Add<1>();
        myAdd->getOperator()->setBackend("cpu");
        myAdd->getOperator()->setDatatype(DataType::Int32);
        myAdd->getOperator()->associateInput(0, input1);
        myAdd->getOperator()->computeOutputDims();
        myAdd->forward();

        REQUIRE(*std::static_pointer_cast<Tensor>(myAdd->getOperator()->getOutput(0)) == *input1);
    }

    SECTION("Two inputs") {
        std::shared_ptr<Tensor> expectedOutput = std::make_shared<Tensor>(Array4D<int,3,3,3,2> { 
            {
                {
                    {{40,  94},{42,  96},{44,  98}},
                    {{46, 100},{48, 102},{50, 104}},
                    {{52, 106},{54, 108},{56, 110}}
                },
                {
                    {{58, 112},{60, 114},{62, 116}},
                    {{64, 118},{66, 120},{68, 122}},
                    {{70, 124},{72, 126},{74, 128}}
                },
                {
                    {{76, 130},{78, 132},{80, 134}},
                    {{82, 136},{84, 138},{86, 140}},
                    {{88, 142},{90, 144},{92, 146}}
                }
            }
        });

        std::shared_ptr<Node> myAdd = Add<2>();
        myAdd->getOperator()->setDatatype(DataType::Int32);
        myAdd->getOperator()->setBackend("cpu");
        myAdd->getOperator()->associateInput(0, input1);
        myAdd->getOperator()->associateInput(1, input1);
        myAdd->getOperator()->computeOutputDims();
        myAdd->forward();

        REQUIRE(*std::static_pointer_cast<Tensor>(myAdd->getOperator()->getOutput(0)) == *expectedOutput);
    }
    
    SECTION("Three inputs") {
        std::shared_ptr<Tensor> expectedOutput = std::make_shared<Tensor>(Array4D<int,3,3,3,2> { 
            {
                {
                    {{ 60, 141},{ 63, 144},{ 66, 147}},
                    {{ 69, 150},{ 72, 153},{ 75, 156}},
                    {{ 78, 159},{ 81, 162},{ 84, 165}}
                },
                {
                    {{ 87, 168},{ 90, 171},{ 93, 174}},
                    {{ 96, 177},{ 99, 180},{102, 183}},
                    {{105, 186},{108, 189},{111, 192}}
                },
                {
                    {{114, 195},{117, 198},{120, 201}},
                    {{123, 204},{126, 207},{129, 210}},
                    {{132, 213},{135, 216},{138, 219}}
                }
            }
        });

        std::shared_ptr<Node> myAdd = Add<3>();
        myAdd->getOperator()->setDatatype(DataType::Int32);
        myAdd->getOperator()->setBackend("cpu");
        myAdd->getOperator()->associateInput(0, input1);
        myAdd->getOperator()->associateInput(1, input1);
        myAdd->getOperator()->associateInput(2, input1);
        myAdd->getOperator()->computeOutputDims();
        myAdd->forward();

        REQUIRE(*std::static_pointer_cast<Tensor>(myAdd->getOperator()->getOutput(0)) == *expectedOutput);
    }
}