// 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 <limits>
#include <numeric>
#include <Eigen/Tensor>

using Eigen::Tensor;

template <int DataLayout>
static void test_trivial_reductions() {
  {
    Tensor<float, 0, DataLayout> tensor;
    tensor.setRandom();
    array<ptrdiff_t, 0> reduction_axis;

    Tensor<float, 0, DataLayout> result = tensor.sum(reduction_axis);
    VERIFY_IS_EQUAL(result(), tensor());
  }

  {
    Tensor<float, 1, DataLayout> tensor(7);
    tensor.setRandom();
    array<ptrdiff_t, 0> reduction_axis;

    Tensor<float, 1, DataLayout> result = tensor.sum(reduction_axis);
    VERIFY_IS_EQUAL(result.dimension(0), 7);
    for (int i = 0; i < 7; ++i) {
      VERIFY_IS_EQUAL(result(i), tensor(i));
    }
  }

  {
    Tensor<float, 2, DataLayout> tensor(2, 3);
    tensor.setRandom();
    array<ptrdiff_t, 0> reduction_axis;

    Tensor<float, 2, DataLayout> result = tensor.sum(reduction_axis);
    VERIFY_IS_EQUAL(result.dimension(0), 2);
    VERIFY_IS_EQUAL(result.dimension(1), 3);
    for (int i = 0; i < 2; ++i) {
      for (int j = 0; j < 3; ++j) {
        VERIFY_IS_EQUAL(result(i, j), tensor(i, j));
      }
    }
  }
}

template <typename Scalar, int DataLayout>
static void test_scalar_reduction_conversion() {
  using TensorType = Tensor<Scalar, 4, DataLayout>;
  using PartialReductionDims = array<ptrdiff_t, 2>;
  using FullReduction = decltype(std::declval<const TensorType&>().sum());
  using PartialReduction = decltype(std::declval<const TensorType&>().sum(std::declval<const PartialReductionDims&>()));

  static_assert(std::is_convertible<FullReduction, Scalar>::value, "full reductions should convert to scalars");
  static_assert(!std::is_convertible<PartialReduction, Scalar>::value,
                "partial reductions should not convert to scalars");

  TensorType tensor(2, 2, 2, 2);
  Scalar expected_sum(0);
  Scalar expected_prod(1);
  Scalar expected_min = NumTraits<Scalar>::highest();
  Scalar expected_max = NumTraits<Scalar>::lowest();
  for (int i = 0; i < tensor.size(); ++i) {
    // Bound the product so the half test does not overflow.
    const Scalar value = static_cast<Scalar>((i % 3) + 1);
    tensor(i) = value;
    expected_sum += value;
    expected_prod *= value;
    expected_min = numext::mini(expected_min, value);
    expected_max = numext::maxi(expected_max, value);
  }

  const Scalar sum = tensor.sum();
  array<ptrdiff_t, 4> reduction_axis4 = {0, 1, 2, 3};
  const Scalar sum_from_dims = tensor.sum(reduction_axis4);
  const Scalar prod = tensor.prod();
  const Scalar smallest = tensor.minimum();
  const Scalar largest = tensor.maximum();
  const Scalar mean = tensor.mean();
  const Scalar expected_mean = expected_sum / static_cast<Scalar>(tensor.size());

  if (NumTraits<Scalar>::IsInteger) {
    VERIFY_IS_EQUAL(sum, expected_sum);
    VERIFY_IS_EQUAL(sum_from_dims, expected_sum);
    VERIFY_IS_EQUAL(prod, expected_prod);
    VERIFY_IS_EQUAL(smallest, expected_min);
    VERIFY_IS_EQUAL(largest, expected_max);
    VERIFY_IS_EQUAL(mean, expected_mean);
  } else {
    VERIFY_IS_APPROX(sum, expected_sum);
    VERIFY_IS_APPROX(sum_from_dims, expected_sum);
    VERIFY_IS_APPROX(prod, expected_prod);
    VERIFY_IS_APPROX(smallest, expected_min);
    VERIFY_IS_APPROX(largest, expected_max);
    VERIFY_IS_APPROX(mean, expected_mean);
  }
}

template <typename Scalar, int DataLayout>
static void test_scalar_reduction_mixed_operators() {
  // Mixed scalar operations must stay lazy despite the reduction's implicit scalar conversion.
  Tensor<Scalar, 2, DataLayout> tensor(3, 4);
  tensor.setConstant(Scalar(2));

  // Scalar-left overloads must not pollute ADL for NumTraits enums.
  static_assert(NumTraits<Scalar>::MulCost + 1 > 0, "NumTraits enum arithmetic should not be ambiguous");

  static_assert(!std::is_same<decltype(tensor.sum() > 0), bool>::value,
                "scalar comparisons on a full reduction should stay lazy expressions");
  static_assert(!std::is_arithmetic<decltype(2 * tensor.sum())>::value,
                "scalar arithmetic on a full reduction should stay lazy expressions");

  const Tensor<bool, 0, DataLayout> gt = tensor.sum() > 0;
  VERIFY(gt());
  const Tensor<bool, 0, DataLayout> eq = tensor.sum() == 24;
  VERIFY(eq());
  const Tensor<bool, 0, DataLayout> lt = tensor.sum() < 1L;
  VERIFY(!lt());
  const Tensor<bool, 0, DataLayout> ge = tensor.sum() >= 0.0f;
  VERIFY(ge());
  const Tensor<bool, 0, DataLayout> le = tensor.sum() <= 100u;
  VERIFY(le());
  const Tensor<bool, 0, DataLayout> ne = tensor.sum() != 0;
  VERIFY(ne());
  const Tensor<Scalar, 0, DataLayout> twice = 2 * tensor.sum();
  VERIFY_IS_EQUAL(twice(), Scalar(48));
  const Tensor<Scalar, 0, DataLayout> plus_one = 1 + tensor.sum();
  VERIFY_IS_EQUAL(plus_one(), Scalar(25));
  const Tensor<Scalar, 0, DataLayout> from_25 = 25 - tensor.sum();
  VERIFY_IS_EQUAL(from_25(), Scalar(1));
  const Tensor<Scalar, 0, DataLayout> ratio = 48 / tensor.sum();
  VERIFY_IS_EQUAL(ratio(), Scalar(2));
  const Tensor<Scalar, 0, DataLayout> scaled = 2.0 * tensor.sum();
  VERIFY_IS_EQUAL(scaled(), Scalar(48));

  // Unscoped enums must remain valid scalar operands.
  enum MixedOpsTestEnum { kEnumTwo = 2 };
  const Tensor<bool, 0, DataLayout> gt_enum = tensor.sum() > kEnumTwo;
  VERIFY(gt_enum());
  const Tensor<Scalar, 0, DataLayout> twice_enum = kEnumTwo * tensor.sum();
  VERIFY_IS_EQUAL(twice_enum(), Scalar(48));
}

static void test_scalar_reduction_mixed_operators_mod() {
  Tensor<int, 2> tensor(3, 4);
  tensor.setConstant(2);
  const Tensor<int, 0> rem = tensor.sum() % 5L;
  VERIFY_IS_EQUAL(rem(), 4);
}

template <typename Scalar, int DataLayout>
static void test_simple_reductions() {
  Tensor<Scalar, 4, DataLayout> tensor(2, 3, 5, 7);
  tensor.setRandom();
  // Add a little offset so that the product reductions won't be close to zero.
  tensor += tensor.constant(Scalar(0.5f));
  array<ptrdiff_t, 2> reduction_axis2;
  reduction_axis2[0] = 1;
  reduction_axis2[1] = 3;

  Tensor<Scalar, 2, DataLayout> result = tensor.sum(reduction_axis2);
  VERIFY_IS_EQUAL(result.dimension(0), 2);
  VERIFY_IS_EQUAL(result.dimension(1), 5);
  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 5; ++j) {
      Scalar sum = Scalar(0.0f);
      for (int k = 0; k < 3; ++k) {
        for (int l = 0; l < 7; ++l) {
          sum += tensor(i, k, j, l);
        }
      }
      VERIFY_IS_APPROX(result(i, j), sum);
    }
  }

  {
    Tensor<Scalar, 0, DataLayout> sum1 = tensor.sum();
    VERIFY_IS_EQUAL(sum1.rank(), 0);

    array<ptrdiff_t, 4> reduction_axis4;
    reduction_axis4[0] = 0;
    reduction_axis4[1] = 1;
    reduction_axis4[2] = 2;
    reduction_axis4[3] = 3;
    Tensor<Scalar, 0, DataLayout> sum2 = tensor.sum(reduction_axis4);
    VERIFY_IS_EQUAL(sum2.rank(), 0);

    VERIFY_IS_APPROX(sum1(), sum2());
  }

  reduction_axis2[0] = 0;
  reduction_axis2[1] = 2;
  result = tensor.prod(reduction_axis2);
  VERIFY_IS_EQUAL(result.dimension(0), 3);
  VERIFY_IS_EQUAL(result.dimension(1), 7);
  for (int i = 0; i < 3; ++i) {
    for (int j = 0; j < 7; ++j) {
      Scalar prod = Scalar(1.0f);
      for (int k = 0; k < 2; ++k) {
        for (int l = 0; l < 5; ++l) {
          prod *= tensor(k, i, l, j);
        }
      }
      VERIFY_IS_APPROX(result(i, j), prod);
    }
  }

  {
    Tensor<Scalar, 0, DataLayout> prod1 = tensor.prod();
    VERIFY_IS_EQUAL(prod1.rank(), 0);

    array<ptrdiff_t, 4> reduction_axis4;
    reduction_axis4[0] = 0;
    reduction_axis4[1] = 1;
    reduction_axis4[2] = 2;
    reduction_axis4[3] = 3;
    Tensor<Scalar, 0, DataLayout> prod2 = tensor.prod(reduction_axis4);
    VERIFY_IS_EQUAL(prod2.rank(), 0);

    VERIFY_IS_APPROX(prod1(), prod2());
  }

  reduction_axis2[0] = 0;
  reduction_axis2[1] = 2;
  result = tensor.maximum(reduction_axis2);
  VERIFY_IS_EQUAL(result.dimension(0), 3);
  VERIFY_IS_EQUAL(result.dimension(1), 7);
  for (int i = 0; i < 3; ++i) {
    for (int j = 0; j < 7; ++j) {
      Scalar max_val = std::numeric_limits<Scalar>::lowest();
      for (int k = 0; k < 2; ++k) {
        for (int l = 0; l < 5; ++l) {
          max_val = (std::max)(max_val, tensor(k, i, l, j));
        }
      }
      VERIFY_IS_APPROX(result(i, j), max_val);
    }
  }

  {
    Tensor<Scalar, 0, DataLayout> max1 = tensor.maximum();
    VERIFY_IS_EQUAL(max1.rank(), 0);

    array<ptrdiff_t, 4> reduction_axis4;
    reduction_axis4[0] = 0;
    reduction_axis4[1] = 1;
    reduction_axis4[2] = 2;
    reduction_axis4[3] = 3;
    Tensor<Scalar, 0, DataLayout> max2 = tensor.maximum(reduction_axis4);
    VERIFY_IS_EQUAL(max2.rank(), 0);

    VERIFY_IS_APPROX(max1(), max2());
  }

  reduction_axis2[0] = 0;
  reduction_axis2[1] = 1;
  result = tensor.minimum(reduction_axis2);
  VERIFY_IS_EQUAL(result.dimension(0), 5);
  VERIFY_IS_EQUAL(result.dimension(1), 7);
  for (int i = 0; i < 5; ++i) {
    for (int j = 0; j < 7; ++j) {
      Scalar min_val = (std::numeric_limits<Scalar>::max)();
      for (int k = 0; k < 2; ++k) {
        for (int l = 0; l < 3; ++l) {
          min_val = (std::min)(min_val, tensor(k, l, i, j));
        }
      }
      VERIFY_IS_APPROX(result(i, j), min_val);
    }
  }

  {
    Tensor<Scalar, 0, DataLayout> min1 = tensor.minimum();
    VERIFY_IS_EQUAL(min1.rank(), 0);

    array<ptrdiff_t, 4> reduction_axis4;
    reduction_axis4[0] = 0;
    reduction_axis4[1] = 1;
    reduction_axis4[2] = 2;
    reduction_axis4[3] = 3;
    Tensor<Scalar, 0, DataLayout> min2 = tensor.minimum(reduction_axis4);
    VERIFY_IS_EQUAL(min2.rank(), 0);

    VERIFY_IS_APPROX(min1(), min2());
  }

  reduction_axis2[0] = 0;
  reduction_axis2[1] = 1;
  result = tensor.mean(reduction_axis2);
  VERIFY_IS_EQUAL(result.dimension(0), 5);
  VERIFY_IS_EQUAL(result.dimension(1), 7);
  for (int i = 0; i < 5; ++i) {
    for (int j = 0; j < 7; ++j) {
      Scalar sum = Scalar(0.0f);
      int count = 0;
      for (int k = 0; k < 2; ++k) {
        for (int l = 0; l < 3; ++l) {
          sum += tensor(k, l, i, j);
          ++count;
        }
      }
      VERIFY_IS_APPROX(result(i, j), sum / Scalar(count));
    }
  }

  {
    Tensor<Scalar, 0, DataLayout> mean1 = tensor.mean();
    VERIFY_IS_EQUAL(mean1.rank(), 0);

    array<ptrdiff_t, 4> reduction_axis4;
    reduction_axis4[0] = 0;
    reduction_axis4[1] = 1;
    reduction_axis4[2] = 2;
    reduction_axis4[3] = 3;
    Tensor<Scalar, 0, DataLayout> mean2 = tensor.mean(reduction_axis4);
    VERIFY_IS_EQUAL(mean2.rank(), 0);

    VERIFY_IS_APPROX(mean1(), mean2());
  }

  {
    Tensor<int, 1> ints(10);
    std::iota(ints.data(), ints.data() + ints.dimension(0), 0);

    TensorFixedSize<bool, Sizes<>> all_;
    all_ = ints.all();
    VERIFY(!all_());
    all_ = (ints >= ints.constant(0)).all();
    VERIFY(all_());

    TensorFixedSize<bool, Sizes<>> any;
    any = (ints > ints.constant(10)).any();
    VERIFY(!any());
    any = (ints < ints.constant(1)).any();
    VERIFY(any());
  }
}

template <int DataLayout>
static void test_reductions_in_expr() {
  Tensor<float, 4, DataLayout> tensor(2, 3, 5, 7);
  tensor.setRandom();
  array<ptrdiff_t, 2> reduction_axis2;
  reduction_axis2[0] = 1;
  reduction_axis2[1] = 3;

  Tensor<float, 2, DataLayout> result(2, 5);
  result = result.constant(1.0f) - tensor.sum(reduction_axis2);
  VERIFY_IS_EQUAL(result.dimension(0), 2);
  VERIFY_IS_EQUAL(result.dimension(1), 5);
  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 5; ++j) {
      float sum = 0.0f;
      for (int k = 0; k < 3; ++k) {
        for (int l = 0; l < 7; ++l) {
          sum += tensor(i, k, j, l);
        }
      }
      VERIFY_IS_APPROX(result(i, j), 1.0f - sum);
    }
  }
}

template <int DataLayout>
static void test_full_reductions() {
  Tensor<float, 2, DataLayout> tensor(2, 3);
  tensor.setRandom();
  array<ptrdiff_t, 2> reduction_axis;
  reduction_axis[0] = 0;
  reduction_axis[1] = 1;

  Tensor<float, 0, DataLayout> result = tensor.sum(reduction_axis);
  VERIFY_IS_EQUAL(result.rank(), 0);

  float sum = 0.0f;
  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      sum += tensor(i, j);
    }
  }
  VERIFY_IS_APPROX(result(0), sum);

  result = tensor.square().sum(reduction_axis).sqrt();
  VERIFY_IS_EQUAL(result.rank(), 0);

  sum = 0.0f;
  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      sum += tensor(i, j) * tensor(i, j);
    }
  }
  VERIFY_IS_APPROX(result(), sqrtf(sum));
}

struct UserReducer {
  static const bool PacketAccess = false;
  UserReducer(float offset) : offset_(offset) {}
  void reduce(const float val, float* accum) { *accum += val * val; }
  float initialize() const { return 0; }
  float finalize(const float accum) const { return 1.0f / (accum + offset_); }

 private:
  const float offset_;
};

// A reducer that transforms each value it accepts (sum of squares), with reducer_traits
// declaring PacketAccess. Its reduce() cannot merge two partial accumulators — feeding a partial
// sum back through reduce() would square it — so the scalar unrolled path must not engage: the
// gate is reducer_can_reorder_accumulators, which such a reducer does not opt into, not
// PacketAccess. long double has no packets, which forces the scalar path while the declared
// PacketAccess stays true.
struct SquaredSumReducer {
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const long double t, long double* accum) const { *accum += t * t; }
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE long double initialize() const { return 0; }
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE long double finalize(const long double accum) const { return accum; }
};

namespace Eigen {
namespace internal {
template <typename Device>
struct reducer_traits<SquaredSumReducer, Device> {
  enum { Cost = 1, PacketAccess = true, IsStateful = false, IsExactlyAssociative = true };
};
}  // namespace internal
}  // namespace Eigen

template <int DataLayout>
static void test_value_transforming_reducer() {
  // The opt-ins that let the unrolled path interleave operands, pinned at both ends: a reducer
  // that transforms each value, and one that carries state, must stay out however they declare
  // PacketAccess.
  STATIC_CHECK((internal::reducer_can_reorder_accumulators<internal::SumReducer<float>>::value));
  STATIC_CHECK((internal::reducer_can_reorder_accumulators<internal::MinReducer<float>>::value));
  STATIC_CHECK((!internal::reducer_can_reorder_accumulators<SquaredSumReducer>::value));
  STATIC_CHECK((!internal::reducer_can_reorder_accumulators<internal::MeanReducer<float>>::value));

  Tensor<long double, 1, DataLayout> t(12);
  long double expected = 0;
  for (int i = 0; i < 12; ++i) {
    t(i) = static_cast<long double>(i + 1);
    expected += t(i) * t(i);
  }
  array<ptrdiff_t, 1> reduction_axis;
  reduction_axis[0] = 0;
  SquaredSumReducer reducer;
  Tensor<long double, 0, DataLayout> result = t.reduce(reduction_axis, reducer);
  VERIFY_IS_EQUAL(result(), expected);  // 1^2 + ... + 12^2 = 650, exact in long double
}

template <int DataLayout>
static void test_user_defined_reductions() {
  Tensor<float, 2, DataLayout> tensor(5, 7);
  tensor.setRandom();
  array<ptrdiff_t, 1> reduction_axis;
  reduction_axis[0] = 1;

  UserReducer reducer(10.0f);
  Tensor<float, 1, DataLayout> result = tensor.reduce(reduction_axis, reducer);
  VERIFY_IS_EQUAL(result.dimension(0), 5);
  for (int i = 0; i < 5; ++i) {
    float expected = 10.0f;
    for (int j = 0; j < 7; ++j) {
      expected += tensor(i, j) * tensor(i, j);
    }
    expected = 1.0f / expected;
    VERIFY_IS_APPROX(result(i), expected);
  }
}

template <int DataLayout>
static void test_tensor_maps() {
  int inputs[2 * 3 * 5 * 7];
  TensorMap<Tensor<int, 4, DataLayout>> tensor_map(inputs, 2, 3, 5, 7);
  TensorMap<Tensor<const int, 4, DataLayout>> tensor_map_const(inputs, 2, 3, 5, 7);
  const TensorMap<Tensor<const int, 4, DataLayout>> tensor_map_const_const(inputs, 2, 3, 5, 7);

  setRandomDataInRange(tensor_map, -1000, 1000);
  array<ptrdiff_t, 2> reduction_axis;
  reduction_axis[0] = 1;
  reduction_axis[1] = 3;

  Tensor<int, 2, DataLayout> result = tensor_map.sum(reduction_axis);
  Tensor<int, 2, DataLayout> result2 = tensor_map_const.sum(reduction_axis);
  Tensor<int, 2, DataLayout> result3 = tensor_map_const_const.sum(reduction_axis);

  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 5; ++j) {
      int sum = 0;
      for (int k = 0; k < 3; ++k) {
        for (int l = 0; l < 7; ++l) {
          sum += tensor_map(i, k, j, l);
        }
      }
      VERIFY_IS_EQUAL(result(i, j), sum);
      VERIFY_IS_EQUAL(result2(i, j), sum);
      VERIFY_IS_EQUAL(result3(i, j), sum);
    }
  }
}

template <int DataLayout>
static void test_static_dims() {
  Tensor<float, 4, DataLayout> in(72, 53, 97, 113);
  Tensor<float, 2, DataLayout> out(72, 97);
  in.setRandom();

  Eigen::IndexList<Eigen::type2index<1>, Eigen::type2index<3>> reduction_axis;

  out = in.maximum(reduction_axis);

  for (int i = 0; i < 72; ++i) {
    for (int j = 0; j < 97; ++j) {
      float expected = -1e10f;
      for (int k = 0; k < 53; ++k) {
        for (int l = 0; l < 113; ++l) {
          expected = (std::max)(expected, in(i, k, j, l));
        }
      }
      VERIFY_IS_EQUAL(out(i, j), expected);
    }
  }
}

template <int DataLayout>
static void test_innermost_last_dims() {
  Tensor<float, 4, DataLayout> in(72, 53, 97, 113);
  Tensor<float, 2, DataLayout> out(97, 113);
  in.setRandom();

  // Reduce on the innermost dimensions.
  // This triggers the use of packets for ColMajor.
  Eigen::IndexList<Eigen::type2index<0>, Eigen::type2index<1>> reduction_axis;

  out = in.maximum(reduction_axis);

  for (int i = 0; i < 97; ++i) {
    for (int j = 0; j < 113; ++j) {
      float expected = -1e10f;
      for (int k = 0; k < 53; ++k) {
        for (int l = 0; l < 72; ++l) {
          expected = (std::max)(expected, in(l, k, i, j));
        }
      }
      VERIFY_IS_EQUAL(out(i, j), expected);
    }
  }
}

template <int DataLayout>
static void test_innermost_first_dims() {
  Tensor<float, 4, DataLayout> in(72, 53, 97, 113);
  Tensor<float, 2, DataLayout> out(72, 53);
  in.setRandom();

  // Reduce on the innermost dimensions.
  // This triggers the use of packets for RowMajor.
  Eigen::IndexList<Eigen::type2index<2>, Eigen::type2index<3>> reduction_axis;

  out = in.maximum(reduction_axis);

  for (int i = 0; i < 72; ++i) {
    for (int j = 0; j < 53; ++j) {
      float expected = -1e10f;
      for (int k = 0; k < 97; ++k) {
        for (int l = 0; l < 113; ++l) {
          expected = (std::max)(expected, in(i, j, k, l));
        }
      }
      VERIFY_IS_EQUAL(out(i, j), expected);
    }
  }
}

template <int DataLayout>
static void test_reduce_middle_dims() {
  Tensor<float, 4, DataLayout> in(72, 53, 97, 113);
  Tensor<float, 2, DataLayout> out(72, 53);
  in.setRandom();

  // Reduce on the innermost dimensions.
  // This triggers the use of packets for RowMajor.
  Eigen::IndexList<Eigen::type2index<1>, Eigen::type2index<2>> reduction_axis;

  out = in.maximum(reduction_axis);

  for (int i = 0; i < 72; ++i) {
    for (int j = 0; j < 113; ++j) {
      float expected = -1e10f;
      for (int k = 0; k < 53; ++k) {
        for (int l = 0; l < 97; ++l) {
          expected = (std::max)(expected, in(i, k, l, j));
        }
      }
      VERIFY_IS_EQUAL(out(i, j), expected);
    }
  }
}

// Exercises the runtime innermost-dim detection added at TensorReduction.h:817
// (coeff path) and :861 (packet path). The static `ReducingInnerMostDims`
// predicate only fires when the reduce-dim container is an IndexList; with
// Eigen::array it stays false at compile time, but the m_reducingInnerMostDims
// runtime mirror picks up the same fast path. We verify correctness against a
// manual sum and check that the result is bit-identical to the IndexList form.
template <int DataLayout>
static void test_dynamic_innermost_dims() {
  // 2D: ColMajor reduces axis 0 (innermost); RowMajor reduces axis 1 (innermost).
  // Size chosen >> PacketSize so packet() drives the hot path.
  Tensor<float, 2, DataLayout> in_2d(127, 53);
  in_2d.setRandom();
  Eigen::array<int, 1> dyn_axis_2d{(DataLayout == ColMajor) ? 0 : 1};
  Tensor<float, 1, DataLayout> out_2d_dyn = in_2d.sum(dyn_axis_2d);

  Eigen::IndexList<Eigen::type2index<(DataLayout == ColMajor) ? 0 : 1>> static_axis_2d;
  Tensor<float, 1, DataLayout> out_2d_static = in_2d.sum(static_axis_2d);

  const int preserved_2d = (DataLayout == ColMajor) ? 53 : 127;
  const int reduced_2d = (DataLayout == ColMajor) ? 127 : 53;
  for (int i = 0; i < preserved_2d; ++i) {
    float expected = 0;
    for (int j = 0; j < reduced_2d; ++j) {
      expected += (DataLayout == ColMajor) ? in_2d(j, i) : in_2d(i, j);
    }
    VERIFY_IS_APPROX(out_2d_dyn(i), expected);
    VERIFY_IS_EQUAL(out_2d_dyn(i), out_2d_static(i));
  }

  // 4D: reduce the two innermost dims at once. This is the multi-reduce case
  // where the runtime mirror also needs to recognize the contiguous block.
  Tensor<float, 4, DataLayout> in_4d(31, 17, 19, 23);
  in_4d.setRandom();
  Eigen::array<int, 2> dyn_axes_4d = (DataLayout == ColMajor) ? Eigen::array<int, 2>{0, 1} : Eigen::array<int, 2>{2, 3};
  Tensor<float, 2, DataLayout> out_4d_dyn = in_4d.sum(dyn_axes_4d);
  Eigen::IndexList<Eigen::type2index<(DataLayout == ColMajor) ? 0 : 2>,
                   Eigen::type2index<(DataLayout == ColMajor) ? 1 : 3>>
      static_axes_4d;
  Tensor<float, 2, DataLayout> out_4d_static = in_4d.sum(static_axes_4d);

  if (DataLayout == ColMajor) {
    for (int i = 0; i < 19; ++i) {
      for (int j = 0; j < 23; ++j) {
        VERIFY_IS_EQUAL(out_4d_dyn(i, j), out_4d_static(i, j));
      }
    }
  } else {
    for (int i = 0; i < 31; ++i) {
      for (int j = 0; j < 17; ++j) {
        VERIFY_IS_EQUAL(out_4d_dyn(i, j), out_4d_static(i, j));
      }
    }
  }

  // Negative case: reducing a non-innermost dim must NOT take the runtime
  // fast path. We can't observe the dispatch directly, but bit-identity with
  // a manual sum + agreement with the IndexList form covers correctness.
  Eigen::array<int, 1> outer_axis{(DataLayout == ColMajor) ? 1 : 0};
  Tensor<float, 1, DataLayout> out_outer = in_2d.sum(outer_axis);
  const int preserved_outer = (DataLayout == ColMajor) ? 127 : 53;
  const int reduced_outer = (DataLayout == ColMajor) ? 53 : 127;
  for (int i = 0; i < preserved_outer; ++i) {
    float expected = 0;
    for (int j = 0; j < reduced_outer; ++j) {
      expected += (DataLayout == ColMajor) ? in_2d(i, j) : in_2d(j, i);
    }
    VERIFY_IS_APPROX(out_outer(i), expected);
  }
}

template <int DataLayout>
static void test_mixed_type_scalar_arithmetic() {
  Tensor<float, 2, DataLayout> in(2, 2);
  in.setConstant(1.0f);
  Tensor<float, 0, DataLayout> result;

  // Mixed type: double OP TensorReductionOp<float>
  result = 0.5 * in.sum();
  VERIFY_IS_EQUAL(result(), 2.0f);
  result = 0.5 + in.sum();
  VERIFY_IS_EQUAL(result(), 4.5f);
  result = 0.5 - in.sum();
  VERIFY_IS_EQUAL(result(), -3.5f);
  result = 2.0 / in.sum();
  VERIFY_IS_EQUAL(result(), 0.5f);

  // Mixed type: TensorReductionOp<float> OP double
  result = in.sum() * 0.5;
  VERIFY_IS_EQUAL(result(), 2.0f);
  result = in.sum() + 0.5;
  VERIFY_IS_EQUAL(result(), 4.5f);
  result = in.sum() - 0.5;
  VERIFY_IS_EQUAL(result(), 3.5f);
  result = in.sum() / 0.5;
  VERIFY_IS_EQUAL(result(), 8.0f);

  // Mixed type: TensorReductionOp<float> OP int
  result = in.sum() * 2;
  VERIFY_IS_EQUAL(result(), 8.0f);

  // Reductions on both sides must keep using the tensor-tensor operators rather than the
  // scalar overloads reachable through the rank-0 conversion to Scalar.
  result = in.sum() * in.sum();
  VERIFY_IS_EQUAL(result(), 16.0f);
  result = in.sum() - in.maximum();
  VERIFY_IS_EQUAL(result(), 3.0f);
}

template <typename ScalarType, int num_elements, int max_mean>
void test_sum_accuracy() {
  Tensor<double, 1> double_tensor(num_elements);
  Tensor<ScalarType, 1> tensor(num_elements);
  for (double prescribed_mean = 0; prescribed_mean <= max_mean;
       prescribed_mean = numext::maxi(1.0, prescribed_mean * 3.99)) {
    // FIXME: NormalRandomGenerator doesn't work in bfloat and half.
    double_tensor.setRandom<Eigen::internal::NormalRandomGenerator<double>>();
    double_tensor += double_tensor.constant(prescribed_mean);
    tensor = double_tensor.cast<ScalarType>();

    Tensor<ScalarType, 0> sum;
    sum = tensor.sum();

    // Compute the reference value in double precision.
    double expected_sum = 0.0;
    double abs_sum = 0.0;
    for (int i = 0; i < num_elements; ++i) {
      expected_sum += static_cast<double>(tensor(i));
      abs_sum += static_cast<double>(numext::abs(tensor(i)));
    }
    // Test against probabilistic forward error bound. In reality, the error is much smaller
    // when we use tree summation.
    double err = Eigen::numext::abs(static_cast<double>(sum()) - expected_sum);
    double tol = numext::sqrt(static_cast<double>(num_elements)) *
                 static_cast<double>(NumTraits<ScalarType>::epsilon()) * abs_sum;
    VERIFY_LE(err, tol);
  }
}

EIGEN_DECLARE_TEST(tensor_reduction) {
  CALL_SUBTEST(test_trivial_reductions<ColMajor>());
  CALL_SUBTEST(test_trivial_reductions<RowMajor>());
  CALL_SUBTEST((test_scalar_reduction_conversion<int, ColMajor>()));
  CALL_SUBTEST((test_scalar_reduction_conversion<int, RowMajor>()));
  CALL_SUBTEST((test_scalar_reduction_conversion<Eigen::half, ColMajor>()));
  CALL_SUBTEST((test_scalar_reduction_mixed_operators<float, ColMajor>()));
  CALL_SUBTEST((test_scalar_reduction_mixed_operators<double, ColMajor>()));
  CALL_SUBTEST((test_scalar_reduction_mixed_operators<int, RowMajor>()));
  CALL_SUBTEST(test_scalar_reduction_mixed_operators_mod());
  CALL_SUBTEST((test_simple_reductions<float, ColMajor>()));
  CALL_SUBTEST((test_simple_reductions<float, RowMajor>()));
  CALL_SUBTEST((test_simple_reductions<Eigen::half, ColMajor>()));
  CALL_SUBTEST((test_simple_reductions<Eigen::bfloat16, ColMajor>()));
  CALL_SUBTEST(test_reductions_in_expr<ColMajor>());
  CALL_SUBTEST(test_reductions_in_expr<RowMajor>());
  CALL_SUBTEST(test_full_reductions<ColMajor>());
  CALL_SUBTEST(test_full_reductions<RowMajor>());
  CALL_SUBTEST(test_user_defined_reductions<ColMajor>());
  CALL_SUBTEST(test_user_defined_reductions<RowMajor>());
  CALL_SUBTEST(test_value_transforming_reducer<ColMajor>());
  CALL_SUBTEST(test_value_transforming_reducer<RowMajor>());
  CALL_SUBTEST(test_tensor_maps<ColMajor>());
  CALL_SUBTEST(test_tensor_maps<RowMajor>());
  CALL_SUBTEST(test_static_dims<ColMajor>());
  CALL_SUBTEST(test_static_dims<RowMajor>());
  CALL_SUBTEST(test_innermost_last_dims<ColMajor>());
  CALL_SUBTEST(test_innermost_last_dims<RowMajor>());
  CALL_SUBTEST(test_innermost_first_dims<ColMajor>());
  CALL_SUBTEST(test_innermost_first_dims<RowMajor>());
  CALL_SUBTEST(test_reduce_middle_dims<ColMajor>());
  CALL_SUBTEST(test_reduce_middle_dims<RowMajor>());
  CALL_SUBTEST(test_dynamic_innermost_dims<ColMajor>());
  CALL_SUBTEST(test_dynamic_innermost_dims<RowMajor>());
  CALL_SUBTEST((test_sum_accuracy<float, 10 * 1024 * 1024, 8 * 1024>()));
  CALL_SUBTEST((test_sum_accuracy<Eigen::bfloat16, 10 * 1024 * 1024, 8 * 1024>()));
  // The range of half is limited to 65519 when using round-to-even,
  // so we are severely limited in the size and mean of the tensors
  // we can reduce without overflow.
  CALL_SUBTEST((test_sum_accuracy<Eigen::half, 4 * 1024, 16>()));
  CALL_SUBTEST((test_sum_accuracy<Eigen::half, 10 * 1024 * 1024, 0>()));
  CALL_SUBTEST(test_mixed_type_scalar_arithmetic<ColMajor>());
  CALL_SUBTEST(test_mixed_type_scalar_arithmetic<RowMajor>());
}