/********************************************************************************
 * 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/operator/Reshape.hpp"

#include "aidge/backend/cpu.hpp"

#include <memory>

using namespace Aidge;

TEST_CASE("[cpu/operator] Reshape(forward)") {
    SECTION("1D Tensor") {
        std::shared_ptr<Tensor> input = std::make_shared<Tensor>(Array1D<float,6> {
            {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}
        });
        std::shared_ptr<Tensor> expectedOutput = std::make_shared<Tensor>(Array2D<float,2,3> {
            {
                {1.0, 2.0, 3.0},
                {4.0, 5.0, 6.0}
            }
        });

        std::shared_ptr<Node> myReshape = Reshape({2, 3});
        auto op = std::static_pointer_cast<OperatorTensor>(myReshape -> getOperator());
        op->associateInput(0, input);
        op->setDataType(DataType::Float32);
        op->setBackend("cpu");
        op->computeOutputDims();
        myReshape->forward();

        REQUIRE(*(op->getOutput(0)) == *expectedOutput);
    }
    SECTION("2D Tensor") {
        std::shared_ptr<Tensor> input = std::make_shared<Tensor>(Array2D<float,2,3> {
            {
                {1.0, 2.0, 3.0},
                {4.0, 5.0, 6.0}
            }

        });
        std::shared_ptr<Tensor> expectedOutput = std::make_shared<Tensor>(Array2D<float,3,2> {
            {
                {1.0, 2.0},
                {3.0, 4.0},
                {5.0, 6.0}
            }
        });

        std::shared_ptr<Node> myReshape = Reshape({3, 2});
        auto op = std::static_pointer_cast<OperatorTensor>(myReshape -> getOperator());
        op->associateInput(0, input);
        op->setDataType(DataType::Float32);
        op->setBackend("cpu");
        op->computeOutputDims();
        myReshape->forward();

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