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

#include "aidge/backend/cpu.hpp"

using namespace Aidge;

TEST_CASE("[cpu/operator] Transpose(forward)") {
    SECTION("3D Tensor") {
        std::shared_ptr<Tensor> input = std::make_shared<Tensor>(Array3D<float,2,3,4> {
            {
                {{0.42507452, 0.11244237, 0.43243718, 0.62354952},
                {0.90250170, 0.48719984, 0.45781207, 0.92536664},
                {0.06348717, 0.91678733, 0.64452291, 0.00484818}},

                {{0.66873497, 0.99508536, 0.55714869, 0.84887981},
                {0.41666120, 0.92365038, 0.80034822, 0.38721532},
                {0.52037925, 0.53937608, 0.66380072, 0.36330253}}
            }
        });
        std::shared_ptr<Tensor> output = std::make_shared<Tensor>(Array3D<float,2,4,3> { 
            {
                {{0.42507452, 0.90250170, 0.06348717},
                {0.11244237, 0.48719984, 0.91678733},
                {0.43243718, 0.45781207, 0.64452291},
                {0.62354952, 0.92536664, 0.00484818}},

                {{0.66873497, 0.41666120, 0.52037925},
                {0.99508536, 0.92365038, 0.53937608},
                {0.55714869, 0.80034822, 0.66380072},
                {0.84887981, 0.38721532, 0.36330253}}
            }
        });
        std::shared_ptr<Node> myTranspose = Transpose<3>(std::array<DimSize_t,3>{{0,2,1}});
        auto op = std::static_pointer_cast<OperatorTensor>(myTranspose -> getOperator());
        op->associateInput(0,input);
        op->setDataType(DataType::Float32);
        op->setBackend("cpu");
        op->computeOutputDims();
        myTranspose->forward();

        REQUIRE(*(op->getOutput(0)) == *output);
    }
    SECTION("4D Tensor") {
        std::shared_ptr<Tensor> input = std::make_shared<Tensor>(Array4D<int,2,3,1,4> {
            {
                {
                    {
                        {1, 2, 3, 4}
                    },
                    {
                        {5, 6, 7, 8}
                    },
                    {
                        {9, 10, 11, 12}
                    }
                },
                {
                    {
                        {13, 14, 15, 16}
                    },
                    {
                        {17, 18, 19, 20}
                    },
                    {
                        {21, 22, 23, 24}
                    }
                }
            }
        });
        std::shared_ptr<Tensor> output = std::make_shared<Tensor>(Array4D<int,2,4,1,3> { 
            {
                {
                    {
                        {1, 5, 9}
                    },
                    {
                        {2, 6, 10}
                    },
                    {
                        {3, 7, 11}
                    },
                    {
                        {4, 8, 12}
                    }
                },
                {
                    {
                        {13, 17, 21}
                    },
                    {
                        {14, 18, 22}
                    },
                    {
                        {15, 19, 23}
                    },
                    {
                        {16, 20, 24}
                    }
                }
            }
        });
        std::shared_ptr<Node> myTranspose = Transpose<4>(std::array<DimSize_t,4>{{0,3,2,1}});
        auto op = std::static_pointer_cast<OperatorTensor>(myTranspose -> getOperator());
        op->associateInput(0,input);
        op->setDataType(DataType::Int32);
        op->setBackend("cpu");
        op->computeOutputDims();
        myTranspose->forward();

        REQUIRE(*(op->getOutput(0)) == *output);
    }
}