//==================================================================================================
/*
  ROTGEN - Runtime Overlay for Eigen
  Copyright : CODE RECKONS
  SPDX-License-Identifier: BSL-1.0
*/
//==================================================================================================
#include "unit/tests.hpp"
#include <rotgen/rotgen.hpp>
#include <Eigen/Dense>

template<typename MatrixType, typename T>
void test_block_matrix_operations(
  rotgen::tests::matrix_block_test_case<MatrixType> const& matrix_construct,
  auto b_init_fn,
  auto ops,
  auto self_ops)
{
  using EigenMatrix = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>;

  auto [d, i0, j0, ni, nj] = matrix_construct;
  auto [r, c, fn] = d;

  MatrixType a(r, c);
  MatrixType b(r, c);
  EigenMatrix ref_a(r, c);
  EigenMatrix ref_b(r, c);

  TTS_EXPECT(verify_rotgen_reentrance(ops(a, b)));

  for (rotgen::Index rr = 0; rr < r; ++rr)
  {
    for (rotgen::Index cc = 0; cc < c; ++cc)
    {
      ref_a(rr, cc) = a(rr, cc) = static_cast<T>(fn(rr, cc));
      ref_b(rr, cc) = b(rr, cc) = static_cast<T>(b_init_fn(rr, cc));
    }
  }

  auto a_block = rotgen::extract(a, i0, j0, ni, nj);
  auto b_block = rotgen::extract(b, i0, j0, ni, nj);
  auto ref_a_block = ref_a.block(i0, j0, ni, nj);
  auto ref_b_block = ref_b.block(i0, j0, ni, nj);

  auto result_block = ops(a_block, b_block);
  auto ref_result_block = ops(ref_a_block, ref_b_block);

  for (rotgen::Index rr = 0; rr < ni; ++rr)
    for (rotgen::Index cc = 0; cc < nj; ++cc)
      TTS_EQUAL(result_block(rr, cc), ref_result_block(rr, cc));

  self_ops(a_block, b_block);
  self_ops(ref_a_block, ref_b_block);

  for (rotgen::Index rr = 0; rr < ni; ++rr)
    for (rotgen::Index cc = 0; cc < nj; ++cc)
      TTS_EQUAL(a_block(rr, cc), ref_a_block(rr, cc));

  for (rotgen::Index rr = 0; rr < ni; ++rr)
  {
    for (rotgen::Index cc = 0; cc < nj; ++cc)
    {
      TTS_EQUAL(a(rr, cc), ref_a(rr, cc));
      TTS_EQUAL(b(rr, cc), ref_b(rr, cc));
    }
  }
}

template<typename MatrixType, typename T>
void test_block_scalar_operations(
  rotgen::tests::matrix_block_test_case<MatrixType> const& matrix_construct,
  auto scalar,
  auto ops,
  auto self_ops)
{
  using EigenMatrix = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>;

  auto [d, i0, j0, ni, nj] = matrix_construct;
  auto [rows, cols, fn] = d;

  MatrixType a(rows, cols);
  EigenMatrix ref_a(rows, cols);

  TTS_EXPECT(verify_rotgen_reentrance(ops(a, scalar)));

  for (rotgen::Index r = 0; r < rows; ++r)
    for (rotgen::Index c = 0; c < cols; ++c)
      ref_a(r, c) = a(r, c) = static_cast<T>(fn(r, c));

  auto a_block = rotgen::extract(a, i0, j0, ni, nj);
  auto ref_a_block = ref_a.block(i0, j0, ni, nj);

  auto result = ops(a_block, scalar);
  auto ref_result = ops(ref_a_block, scalar);

  for (rotgen::Index r = 0; r < ni; ++r)
    for (rotgen::Index c = 0; c < nj; ++c)
      TTS_EQUAL(result(r, c), ref_result(r, c));

  self_ops(a_block, scalar);
  self_ops(ref_a_block, scalar);

  for (rotgen::Index r = 0; r < ni; ++r)
    for (rotgen::Index c = 0; c < nj; ++c)
      TTS_EQUAL(a_block(r, c), ref_a_block(r, c));

  for (rotgen::Index r = 0; r < rows; ++r)
    for (rotgen::Index c = 0; c < cols; ++c) TTS_EQUAL(a(r, c), ref_a(r, c));
}

template<typename MatrixType, typename T>
void test_scalar_block_multiplications(
  rotgen::tests::matrix_block_test_case<MatrixType> const& matrix_construct,
  T scalar)
{
  using EigenMatrix = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>;

  auto [d, i0, j0, ni, nj] = matrix_construct;
  auto [rows, cols, fn] = d;

  MatrixType a(rows, cols);
  EigenMatrix ref_a(rows, cols);

  TTS_EXPECT(verify_rotgen_reentrance(a * scalar));
  TTS_EXPECT(verify_rotgen_reentrance(scalar * a));

  for (rotgen::Index r = 0; r < rows; ++r)
    for (rotgen::Index c = 0; c < cols; ++c)
      ref_a(r, c) = a(r, c) = static_cast<T>(fn(r, c));

  auto a_block = rotgen::extract(a, i0, j0, ni, nj);
  auto ref_a_block = ref_a.block(i0, j0, ni, nj);

  auto a_scalar_multiplication = a_block * scalar;
  auto scalar_a_multiplication = scalar * a_block;
  auto a_scalar_multiplication_ref = ref_a_block * scalar;
  auto scalar_a_multiplication_ref = scalar * ref_a_block;

  for (rotgen::Index r = 0; r < ni; ++r)
  {
    for (rotgen::Index c = 0; c < nj; ++c)
    {
      TTS_EQUAL(a_scalar_multiplication(r, c),
                a_scalar_multiplication_ref(r, c));
      TTS_EQUAL(scalar_a_multiplication(r, c),
                scalar_a_multiplication_ref(r, c));
    }
  }

  a_block *= scalar;
  ref_a_block *= scalar;

  for (rotgen::Index r = 0; r < ni; ++r)
    for (rotgen::Index c = 0; c < nj; ++c)
      TTS_EQUAL(a_block(r, c), ref_a_block(r, c));

  for (rotgen::Index r = 0; r < rows; ++r)
    for (rotgen::Index c = 0; c < cols; ++c) TTS_EQUAL(a(r, c), ref_a(r, c));
}

template<typename MatrixType, typename T>
void test_block_multiplication(
  rotgen::tests::matrix_block_test_case<MatrixType> const& a_matrix_construct,
  rotgen::tests::matrix_block_test_case<MatrixType> const& b_matrix_construct)
{
  using EigenMatrix = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>;

  auto [a_d, a_i0, a_j0, a_ni, a_nj] = a_matrix_construct;
  auto [a_rows, a_cols, a_fn] = a_d;

  auto [b_d, b_i0, b_j0, b_ni, b_nj] = b_matrix_construct;
  auto [b_rows, b_cols, b_fn] = b_d;

  MatrixType a(a_rows, a_cols);
  MatrixType b(b_rows, b_cols);
  EigenMatrix ref_a(a_rows, a_cols);
  EigenMatrix ref_b(b_rows, b_cols);

  for (rotgen::Index r = 0; r < a_rows; ++r)
    for (rotgen::Index c = 0; c < a_cols; ++c)
      ref_a(r, c) = a(r, c) = static_cast<T>(a_fn(r, c));

  for (rotgen::Index r = 0; r < b_rows; ++r)
    for (rotgen::Index c = 0; c < b_cols; ++c)
      ref_b(r, c) = b(r, c) = static_cast<T>(b_fn(r, c));

  auto a_block = rotgen::extract(a, a_i0, a_j0, a_ni, a_nj);
  auto b_block = rotgen::extract(b, b_i0, b_j0, b_ni, b_nj);
  auto ref_a_block = ref_a.block(a_i0, a_j0, a_ni, a_nj);
  auto ref_b_block = ref_b.block(b_i0, b_j0, b_ni, b_nj);

  TTS_EXPECT(verify_rotgen_reentrance(a_block * b_block));

  auto a_b_product_original = a_block * b_block;
  auto a_b_product_ref = ref_a_block * ref_b_block;

  for (rotgen::Index r = 0; r < a_ni; ++r)
    for (rotgen::Index c = 0; c < a_nj; ++c)
      TTS_EQUAL(a_b_product_original(r, c), a_b_product_ref(r, c));

  a_block *= b_block;
  ref_a_block *= ref_b_block;

  for (rotgen::Index r = 0; r < a_ni; ++r)
    for (rotgen::Index c = 0; c < a_nj; ++c)
      TTS_EQUAL(a_block(r, c), ref_a_block(r, c));

  for (rotgen::Index r = 0; r < a_rows; ++r)
    for (rotgen::Index c = 0; c < a_cols; ++c) TTS_EQUAL(a(r, c), ref_a(r, c));

  for (rotgen::Index r = 0; r < b_rows; ++r)
    for (rotgen::Index c = 0; c < b_cols; ++c) TTS_EQUAL(b(r, c), ref_b(r, c));
}

// Basic initializers
inline constexpr auto init_a = [](auto r, auto c) {
  return 9.9 * r * r * r - 6 * c - 12;
};
inline constexpr auto init_b = [](auto r, auto c) {
  return 3.1 * r + 4.2 * c - 12.3;
};
inline constexpr auto init_0 = [](auto, auto) { return 0; };

TTS_CASE_TPL("Check block addition",
             rotgen::tests::types)<typename T, typename O>(
  tts::type<tts::types<T, O>>)
{
  using mat_t = rotgen::matrix<T, rotgen::Dynamic, rotgen::Dynamic, O::value>;
  auto op = [](auto a, auto b) { return a + b; };
  auto s_op = [](auto& a, auto b) { return a += b; };

  test_block_matrix_operations<mat_t, T>({{1, 1, init_a}, 0, 0, 1, 1}, init_b,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{13, 15, init_a}, 1, 2, 3, 4}, init_b,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{5, 9, init_a}, 2, 2, 2, 2}, init_b,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{15, 15, init_a}, 3, 4, 5, 5}, init_b,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{5, 5, init_b}, 1, 0, 3, 2}, init_a,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{10, 1, init_a}, 0, 0, 5, 1}, init_b,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{1, 10, init_a}, 0, 0, 1, 5}, init_b,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{21, 5, init_0}, 4, 4, 10, 1}, init_b,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{11, 7, init_a}, 2, 0, 7, 5}, init_0,
                                         op, s_op);
};

TTS_CASE_TPL("Check block subtraction",
             rotgen::tests::types)<typename T, typename O>(
  tts::type<tts::types<T, O>>)
{
  using mat_t = rotgen::matrix<T, rotgen::Dynamic, rotgen::Dynamic, O::value>;
  auto op = [](auto a, auto b) { return a - b; };
  auto s_op = [](auto& a, auto b) { return a -= b; };

  test_block_matrix_operations<mat_t, T>({{1, 1, init_a}, 0, 0, 1, 1}, init_b,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{1, 1, init_a}, 0, 0, 1, 1}, init_b,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{13, 15, init_a}, 1, 2, 3, 4}, init_b,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{5, 9, init_a}, 2, 2, 2, 2}, init_b,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{15, 15, init_a}, 3, 4, 5, 5}, init_b,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{5, 5, init_b}, 1, 0, 3, 2}, init_a,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{10, 1, init_a}, 0, 0, 5, 1}, init_b,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{1, 10, init_a}, 0, 0, 1, 5}, init_b,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{21, 5, init_0}, 4, 4, 10, 1}, init_b,
                                         op, s_op);
  test_block_matrix_operations<mat_t, T>({{11, 7, init_a}, 2, 0, 7, 5}, init_0,
                                         op, s_op);
};

TTS_CASE_TPL("Check block multiplications",
             rotgen::tests::types)<typename T, typename O>(
  tts::type<tts::types<T, O>>)
{
  using mat_t = rotgen::matrix<T, rotgen::Dynamic, rotgen::Dynamic, O::value>;
  auto init_id = [](rotgen::Index r, rotgen::Index c) {
    return r == c ? 1 : 0;
  };

  test_block_multiplication<mat_t, T>({{1, 1, init_a}, 0, 0, 1, 1},
                                      {{1, 1, init_b}, 0, 0, 1, 1});
  test_block_multiplication<mat_t, T>({{13, 15, init_b}, 1, 2, 4, 4},
                                      {{13, 15, init_b}, 0, 1, 4, 4});
  test_block_multiplication<mat_t, T>({{6, 9, init_a}, 2, 2, 2, 2},
                                      {{5, 9, init_b}, 2, 2, 2, 2});
  test_block_multiplication<mat_t, T>({{15, 15, init_a}, 3, 4, 5, 5},
                                      {{15, 15, init_b}, 0, 0, 5, 5});
  test_block_multiplication<mat_t, T>({{5, 5, init_b}, 1, 0, 3, 3},
                                      {{5, 5, init_a}, 0, 0, 3, 3});
  test_block_multiplication<mat_t, T>({{11, 3, init_id}, 2, 0, 2, 2},
                                      {{18, 7, init_a}, 2, 0, 2, 2});
  test_block_multiplication<mat_t, T>({{10, 1, init_a}, 0, 0, 1, 1},
                                      {{10, 1, init_a}, 0, 0, 1, 1});
  test_block_multiplication<mat_t, T>({{1, 10, init_a}, 0, 0, 1, 1},
                                      {{1, 10, init_id}, 0, 0, 1, 1});
  test_block_multiplication<mat_t, T>({{21, 5, init_0}, 1, 1, 3, 3},
                                      {{12, 7, init_0}, 4, 4, 3, 3});
  test_block_multiplication<mat_t, T>({{11, 7, init_a}, 2, 0, 7, 7},
                                      {{11, 11, init_a}, 2, 1, 7, 7});
  test_block_multiplication<mat_t, T>({{11, 7, init_a}, 2, 0, 5, 5},
                                      {{11, 12, init_id}, 0, 0, 5, 5});
};

TTS_CASE_TPL("Check block multiplication with scalar",
             rotgen::tests::types)<typename T, typename O>(
  tts::type<tts::types<T, O>>)
{
  using mat_t = rotgen::matrix<T, rotgen::Dynamic, rotgen::Dynamic, O::value>;

  test_scalar_block_multiplications<mat_t, T>({{1, 1, init_a}, 0, 0, 1, 1},
                                              T{3.5});
  test_scalar_block_multiplications<mat_t, T>({{13, 15, init_a}, 1, 2, 3, 4},
                                              T{0.});
  test_scalar_block_multiplications<mat_t, T>({{5, 9, init_a}, 2, 2, 2, 2},
                                              T{-2.5});
  test_scalar_block_multiplications<mat_t, T>({{15, 15, init_a}, 3, 4, 5, 5},
                                              T{42.});
  test_scalar_block_multiplications<mat_t, T>({{5, 5, init_b}, 1, 0, 3, 2},
                                              T{-5.});
  test_scalar_block_multiplications<mat_t, T>({{10, 1, init_a}, 0, 0, 5, 1},
                                              T{1.});
  test_scalar_block_multiplications<mat_t, T>({{1, 10, init_a}, 0, 0, 1, 5},
                                              T{6.});
  test_scalar_block_multiplications<mat_t, T>({{21, 5, init_0}, 4, 4, 10, 1},
                                              T{10.1});
  test_scalar_block_multiplications<mat_t, T>({{11, 7, init_a}, 2, 0, 7, 5},
                                              T{-0.5});
};

TTS_CASE_TPL("Check block division with scalar",
             rotgen::tests::types)<typename T, typename O>(
  tts::type<tts::types<T, O>>)
{
  using mat_t = rotgen::matrix<T, rotgen::Dynamic, rotgen::Dynamic, O::value>;
  auto op = [](auto a, auto b) { return a / b; };
  auto s_op = [](auto& a, auto b) { return a /= b; };

  test_block_scalar_operations<mat_t, T>({{1, 1, init_a}, 0, 0, 1, 1}, T{3.5},
                                         op, s_op);
  test_block_scalar_operations<mat_t, T>({{13, 15, init_a}, 1, 2, 3, 4},
                                         T{-2.5}, op, s_op);
  test_block_scalar_operations<mat_t, T>({{5, 9, init_a}, 2, 2, 2, 2}, T{42.},
                                         op, s_op);
  test_block_scalar_operations<mat_t, T>({{15, 15, init_a}, 3, 4, 5, 5}, T{-5.},
                                         op, s_op);
  test_block_scalar_operations<mat_t, T>({{5, 5, init_b}, 1, 0, 3, 2}, T{1.},
                                         op, s_op);
  test_block_scalar_operations<mat_t, T>({{10, 1, init_a}, 0, 0, 5, 1}, T{0.},
                                         op, s_op);
  test_block_scalar_operations<mat_t, T>({{1, 10, init_a}, 0, 0, 1, 5}, T{6.},
                                         op, s_op);
  test_block_scalar_operations<mat_t, T>({{21, 5, init_0}, 4, 4, 10, 1}, T{10.},
                                         op, s_op);
  test_block_scalar_operations<mat_t, T>({{11, 7, init_a}, 2, 0, 7, 5}, T{-0.5},
                                         op, s_op);
};