/********************************************************************************
 * 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 <array>

#include <catch2/catch_test_macros.hpp>

#include "aidge/data/Tensor.hpp"
#include "aidge/utils/TensorUtils.hpp"
#include "aidge/backend/cpu/data/TensorImpl.hpp"

using namespace Aidge;

TEST_CASE("Tensor creation") {
  SECTION("from const array") {
    Tensor x = Array3D<int, 2, 2, 2>{{{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}}};

    Tensor xCopy = Array3D<int, 2, 2, 2>{{{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}}};

    Tensor xFloat =
        Array3D<float, 2, 2, 2>{{{{1., 2.}, {3., 4.}}, {{5., 6.}, {7., 8.}}}};

    SECTION("Tensor features") {
      REQUIRE(x.nbDims() == 3);
      REQUIRE(x.dims()[0] == 2);
      REQUIRE(x.dims()[1] == 2);
      REQUIRE(x.dims()[2] == 2);
      REQUIRE(x.size() == 8);
    }

    SECTION("Access to array") {
      REQUIRE(static_cast<int *>(x.getImpl()->rawPtr())[0] == 1);
      REQUIRE(static_cast<int *>(x.getImpl()->rawPtr())[7] == 8);
    }

    SECTION("get function") {
      REQUIRE(x.get<int>({0, 0, 0}) == 1);
      REQUIRE(x.get<int>({0, 0, 1}) == 2);
      REQUIRE(x.get<int>({0, 1, 1}) == 4);
      REQUIRE(x.get<int>({1, 1, 0}) == 7);
      x.set<int>({1, 1, 1}, 36);
      REQUIRE(x.get<int>({1, 1, 1}) == 36);
    }

    SECTION("Pretty printing for debug") { REQUIRE_NOTHROW(x.print()); }

    SECTION("Tensor (in)equality") {
      REQUIRE(x == xCopy);
      REQUIRE_FALSE(x == xFloat);
    }
  }
}

TEST_CASE("Tensor methods") {
  Tensor x = Array3D<int, 2, 2, 2>{{
    {{1, 2},
     {3, 4}},
    {{5, 6},
     {7, 8}}
  }};

  Tensor xCopy = Array3D<int, 2, 2, 2>{{{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}}};

  Tensor xFloat =
      Array3D<float, 2, 2, 2>{{{{1., 2.}, {3., 4.}}, {{5., 6.}, {7., 8.}}}};

  SECTION("Tensor sharing") {
    Tensor xCopyCtor(x);
    REQUIRE(xCopyCtor.getImpl() == x.getImpl());

    Tensor xEqOp = x;
    REQUIRE(xEqOp.getImpl() == x.getImpl());

    Tensor xCloned = x.clone();
    REQUIRE(xCloned.getImpl() != x.getImpl());
    REQUIRE(xCloned == x);
  }

  SECTION("Tensor extract") {
    Tensor y = x.extract({0, 1});
    REQUIRE(y.getImpl() == x.getImpl());
    REQUIRE(approxEq<int>(y, Array1D<int, 2>{{3, 4}}));
    REQUIRE(y.isContiguous());
    
    Tensor y2 = x.extract({0, 1, 1}, {2, 1, 1});
    REQUIRE(y2.getImpl() == x.getImpl());
    REQUIRE(!y2.isContiguous());
    Tensor y3 = y2.clone();
    REQUIRE(y3.isContiguous());
    REQUIRE(approxEq<int>(y3, Array3D<int, 2, 1, 1>{{{{4}}, {{8}}}}));
  }
}