// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-License-Identifier: MPL-2.0

#include "main.h"

#include <Eigen/Tensor>

using Eigen::RowMajor;
using Eigen::Tensor;

static void test_0d() {
  Tensor<int, 0> scalar1;
  Tensor<int, 0, RowMajor> scalar2;

  TensorMap<const Tensor<int, 0>> scalar3(scalar1.data());
  TensorMap<const Tensor<int, 0, RowMajor>> scalar4(scalar2.data());

  scalar1() = 7;
  scalar2() = 13;

  VERIFY_IS_EQUAL(scalar1.rank(), 0);
  VERIFY_IS_EQUAL(scalar1.size(), 1);

  VERIFY_IS_EQUAL(scalar3(), 7);
  VERIFY_IS_EQUAL(scalar4(), 13);
}

static void test_1d() {
  Tensor<int, 1> vec1(6);
  Tensor<int, 1, RowMajor> vec2(6);

  TensorMap<const Tensor<int, 1>> vec3(vec1.data(), 6);
  TensorMap<const Tensor<int, 1, RowMajor>> vec4(vec2.data(), 6);

  vec1(0) = 4;
  vec2(0) = 0;
  vec1(1) = 8;
  vec2(1) = 1;
  vec1(2) = 15;
  vec2(2) = 2;
  vec1(3) = 16;
  vec2(3) = 3;
  vec1(4) = 23;
  vec2(4) = 4;
  vec1(5) = 42;
  vec2(5) = 5;

  VERIFY_IS_EQUAL(vec1.rank(), 1);
  VERIFY_IS_EQUAL(vec1.size(), 6);
  VERIFY_IS_EQUAL(vec1.dimension(0), 6);

  VERIFY_IS_EQUAL(vec3(0), 4);
  VERIFY_IS_EQUAL(vec3(1), 8);
  VERIFY_IS_EQUAL(vec3(2), 15);
  VERIFY_IS_EQUAL(vec3(3), 16);
  VERIFY_IS_EQUAL(vec3(4), 23);
  VERIFY_IS_EQUAL(vec3(5), 42);

  VERIFY_IS_EQUAL(vec4(0), 0);
  VERIFY_IS_EQUAL(vec4(1), 1);
  VERIFY_IS_EQUAL(vec4(2), 2);
  VERIFY_IS_EQUAL(vec4(3), 3);
  VERIFY_IS_EQUAL(vec4(4), 4);
  VERIFY_IS_EQUAL(vec4(5), 5);
}

static void test_2d() {
  Tensor<int, 2> mat1(2, 3);
  Tensor<int, 2, RowMajor> mat2(2, 3);

  mat1(0, 0) = 0;
  mat1(0, 1) = 1;
  mat1(0, 2) = 2;
  mat1(1, 0) = 3;
  mat1(1, 1) = 4;
  mat1(1, 2) = 5;

  mat2(0, 0) = 0;
  mat2(0, 1) = 1;
  mat2(0, 2) = 2;
  mat2(1, 0) = 3;
  mat2(1, 1) = 4;
  mat2(1, 2) = 5;

  TensorMap<const Tensor<int, 2>> mat3(mat1.data(), 2, 3);
  TensorMap<const Tensor<int, 2, RowMajor>> mat4(mat2.data(), 2, 3);

  VERIFY_IS_EQUAL(mat3.rank(), 2);
  VERIFY_IS_EQUAL(mat3.size(), 6);
  VERIFY_IS_EQUAL(mat3.dimension(0), 2);
  VERIFY_IS_EQUAL(mat3.dimension(1), 3);

  VERIFY_IS_EQUAL(mat4.rank(), 2);
  VERIFY_IS_EQUAL(mat4.size(), 6);
  VERIFY_IS_EQUAL(mat4.dimension(0), 2);
  VERIFY_IS_EQUAL(mat4.dimension(1), 3);

  VERIFY_IS_EQUAL(mat3(0, 0), 0);
  VERIFY_IS_EQUAL(mat3(0, 1), 1);
  VERIFY_IS_EQUAL(mat3(0, 2), 2);
  VERIFY_IS_EQUAL(mat3(1, 0), 3);
  VERIFY_IS_EQUAL(mat3(1, 1), 4);
  VERIFY_IS_EQUAL(mat3(1, 2), 5);

  VERIFY_IS_EQUAL(mat4(0, 0), 0);
  VERIFY_IS_EQUAL(mat4(0, 1), 1);
  VERIFY_IS_EQUAL(mat4(0, 2), 2);
  VERIFY_IS_EQUAL(mat4(1, 0), 3);
  VERIFY_IS_EQUAL(mat4(1, 1), 4);
  VERIFY_IS_EQUAL(mat4(1, 2), 5);
}

static void test_3d() {
  Tensor<int, 3> mat1(2, 3, 7);
  Tensor<int, 3, RowMajor> mat2(2, 3, 7);

  int val = 0;
  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      for (int k = 0; k < 7; ++k) {
        mat1(i, j, k) = val;
        mat2(i, j, k) = val;
        val++;
      }
    }
  }

  TensorMap<const Tensor<int, 3>> mat3(mat1.data(), 2, 3, 7);
  TensorMap<const Tensor<int, 3, RowMajor>> mat4(mat2.data(), 2, 3, 7);

  VERIFY_IS_EQUAL(mat3.rank(), 3);
  VERIFY_IS_EQUAL(mat3.size(), 2 * 3 * 7);
  VERIFY_IS_EQUAL(mat3.dimension(0), 2);
  VERIFY_IS_EQUAL(mat3.dimension(1), 3);
  VERIFY_IS_EQUAL(mat3.dimension(2), 7);

  VERIFY_IS_EQUAL(mat4.rank(), 3);
  VERIFY_IS_EQUAL(mat4.size(), 2 * 3 * 7);
  VERIFY_IS_EQUAL(mat4.dimension(0), 2);
  VERIFY_IS_EQUAL(mat4.dimension(1), 3);
  VERIFY_IS_EQUAL(mat4.dimension(2), 7);

  val = 0;
  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      for (int k = 0; k < 7; ++k) {
        VERIFY_IS_EQUAL(mat3(i, j, k), val);
        VERIFY_IS_EQUAL(mat4(i, j, k), val);
        val++;
      }
    }
  }
}

static void test_from_tensor() {
  Tensor<int, 3> mat1(2, 3, 7);
  Tensor<int, 3, RowMajor> mat2(2, 3, 7);

  int val = 0;
  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      for (int k = 0; k < 7; ++k) {
        mat1(i, j, k) = val;
        mat2(i, j, k) = val;
        val++;
      }
    }
  }

  TensorMap<Tensor<int, 3>> mat3(mat1);
  TensorMap<Tensor<int, 3, RowMajor>> mat4(mat2);

  VERIFY_IS_EQUAL(mat3.rank(), 3);
  VERIFY_IS_EQUAL(mat3.size(), 2 * 3 * 7);
  VERIFY_IS_EQUAL(mat3.dimension(0), 2);
  VERIFY_IS_EQUAL(mat3.dimension(1), 3);
  VERIFY_IS_EQUAL(mat3.dimension(2), 7);

  VERIFY_IS_EQUAL(mat4.rank(), 3);
  VERIFY_IS_EQUAL(mat4.size(), 2 * 3 * 7);
  VERIFY_IS_EQUAL(mat4.dimension(0), 2);
  VERIFY_IS_EQUAL(mat4.dimension(1), 3);
  VERIFY_IS_EQUAL(mat4.dimension(2), 7);

  val = 0;
  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      for (int k = 0; k < 7; ++k) {
        VERIFY_IS_EQUAL(mat3(i, j, k), val);
        VERIFY_IS_EQUAL(mat4(i, j, k), val);
        val++;
      }
    }
  }

  TensorFixedSize<int, Sizes<2, 3, 7>> mat5;

  val = 0;
  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      for (int k = 0; k < 7; ++k) {
        array<ptrdiff_t, 3> coords;
        coords[0] = i;
        coords[1] = j;
        coords[2] = k;
        mat5(coords) = val;
        val++;
      }
    }
  }

  TensorMap<TensorFixedSize<int, Sizes<2, 3, 7>>> mat6(mat5);

  VERIFY_IS_EQUAL(mat6.rank(), 3);
  VERIFY_IS_EQUAL(mat6.size(), 2 * 3 * 7);
  VERIFY_IS_EQUAL(mat6.dimension(0), 2);
  VERIFY_IS_EQUAL(mat6.dimension(1), 3);
  VERIFY_IS_EQUAL(mat6.dimension(2), 7);

  val = 0;
  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      for (int k = 0; k < 7; ++k) {
        VERIFY_IS_EQUAL(mat6(i, j, k), val);
        val++;
      }
    }
  }
}

static int f(const TensorMap<Tensor<int, 3>>& tensor) {
  //  Size<0> empty;
  EIGEN_STATIC_ASSERT((internal::array_size<Sizes<>>::value == 0), YOU_MADE_A_PROGRAMMING_MISTAKE);
  EIGEN_STATIC_ASSERT((internal::array_size<DSizes<int, 0>>::value == 0), YOU_MADE_A_PROGRAMMING_MISTAKE);
  Tensor<int, 0> result = tensor.sum();
  return result();
}

static void test_casting() {
  Tensor<int, 3> tensor(2, 3, 7);

  int val = 0;
  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      for (int k = 0; k < 7; ++k) {
        tensor(i, j, k) = val;
        val++;
      }
    }
  }

  TensorMap<Tensor<int, 3>> map(tensor);
  int sum1 = f(map);
  int sum2 = f(tensor);

  VERIFY_IS_EQUAL(sum1, sum2);
  VERIFY_IS_EQUAL(sum1, 861);
}

template <typename T>
static const T& add_const(T& value) {
  return value;
}

static void test_0d_const_tensor() {
  Tensor<int, 0> scalar1;
  Tensor<int, 0, RowMajor> scalar2;

  TensorMap<const Tensor<int, 0>> scalar3(add_const(scalar1).data());
  TensorMap<const Tensor<int, 0, RowMajor>> scalar4(add_const(scalar2).data());

  scalar1() = 7;
  scalar2() = 13;

  VERIFY_IS_EQUAL(scalar1.rank(), 0);
  VERIFY_IS_EQUAL(scalar1.size(), 1);

  VERIFY_IS_EQUAL(scalar3(), 7);
  VERIFY_IS_EQUAL(scalar4(), 13);
}

static void test_0d_const_tensor_map() {
  Tensor<int, 0> scalar1;
  Tensor<int, 0, RowMajor> scalar2;

  const TensorMap<Tensor<int, 0>> scalar3(scalar1.data());
  const TensorMap<Tensor<int, 0, RowMajor>> scalar4(scalar2.data());

  // Although TensorMap is constant, we still can write to the underlying
  // storage, because we map over non-constant Tensor.
  scalar3() = 7;
  scalar4() = 13;

  VERIFY_IS_EQUAL(scalar1(), 7);
  VERIFY_IS_EQUAL(scalar2(), 13);

  // Pointer to the underlying storage is also non-const.
  scalar3.data()[0] = 8;
  scalar4.data()[0] = 14;

  VERIFY_IS_EQUAL(scalar1(), 8);
  VERIFY_IS_EQUAL(scalar2(), 14);
}

// Regression test for issue #1616: unsigned index and dimension types must not
// trigger narrowing errors in the variadic constructor and accessors.
static void test_unsigned_indices() {
  Tensor<int, 3> tensor(2, 3, 7);
  tensor.setZero();

  TensorMap<Tensor<int, 3>> map(tensor.data(), 2u, 3ul, std::size_t(7));
  map(1u, std::size_t(2), 6ul) = 5;
  VERIFY_IS_EQUAL(map(1ul, 2u, std::size_t(6)), 5);
  VERIFY_IS_EQUAL(tensor(1, 2, 6), 5);

  TensorMap<const Tensor<int, 3>> const_map(tensor.data(), 2u, 3ul, std::size_t(7));
  VERIFY_IS_EQUAL(const_map(1u, 2ul, std::size_t(6)), 5);
}

// An index or dimension the tensor's index type cannot represent must assert rather than silently truncate.
static void test_narrowing_indices() {
  // Fixed width: on ILP32 targets shifting std::size_t by 40 is undefined, and GCC folds it to an in-range 0.
  const std::uint64_t too_large = std::uint64_t(1) << 40;

  Tensor<int, 3, ColMajor, int> tensor(2, 3, 7);
  tensor.setZero();

  VERIFY_RAISES_ASSERT((TensorMap<Tensor<int, 3, ColMajor, int>>(tensor.data(), 2, 3, too_large)));

  TensorMap<Tensor<int, 3, ColMajor, int>> map(tensor.data(), 2, 3, 7);
  VERIFY_RAISES_ASSERT(map(0, 0, too_large) = 1);

  TensorMap<const Tensor<int, 3, ColMajor, int>> const_map(tensor.data(), 2, 3, 7);
  VERIFY_RAISES_ASSERT(const_map(0, 0, too_large));
}

EIGEN_DECLARE_TEST(tensor_map) {
  CALL_SUBTEST(test_0d());
  CALL_SUBTEST(test_1d());
  CALL_SUBTEST(test_2d());
  CALL_SUBTEST(test_3d());

  CALL_SUBTEST(test_from_tensor());
  CALL_SUBTEST(test_casting());

  CALL_SUBTEST(test_0d_const_tensor());
  CALL_SUBTEST(test_0d_const_tensor_map());

  CALL_SUBTEST(test_unsigned_indices());
  CALL_SUBTEST(test_narrowing_indices());
}