//==================================================================================================
/*
  ROTGEN - Runtime Overlay for Eigen
  Copyright : CODE RECKONS
  SPDX-License-Identifier: BSL-1.0
*/
//==================================================================================================
#include "unit/tests.hpp"
#include <rotgen/matrix.hpp>
#include <rotgen/extract.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 self_ops)
{
  using EigenMatrix = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>;

  MatrixType a(matrix_construct.rows, matrix_construct.cols);
  MatrixType b(matrix_construct.rows, matrix_construct.cols);
  EigenMatrix ref_a(matrix_construct.rows, matrix_construct.cols);
  EigenMatrix ref_b(matrix_construct.rows, matrix_construct.cols);

  for (std::size_t r = 0; r < matrix_construct.rows; ++r)
  {
    for (std::size_t c = 0; c < matrix_construct.cols; ++c)
    {
      ref_a(r,c) = a(r,c) = static_cast<T>(matrix_construct.init_fn(r, c));
      ref_b(r,c) = b(r,c) = static_cast<T>(b_init_fn(r, c));
    }
  }

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

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

  for (std::size_t r = 0; r < matrix_construct.ni; ++r)
    for (std::size_t c = 0; c < matrix_construct.nj; ++c)
      TTS_EQUAL(result_block(r, c), ref_result_block(r, c));*/

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

  for (std::size_t r = 0; r < matrix_construct.ni; ++r)
    for (std::size_t c = 0; c < matrix_construct.nj; ++c)
      TTS_EQUAL(a_block(r, c), ref_a_block(r, c));

  for (std::size_t r = 0; r < matrix_construct.rows; ++r)
  {
    for (std::size_t c = 0; c < matrix_construct.cols; ++c)
    {
      TTS_EQUAL(a(r, c), ref_a(r, c));
      TTS_EQUAL(b(r, c), ref_b(r, c));
    }
  }
}

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

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

  for (std::size_t r = 0; r < matrix_construct.rows; ++r)
    for (std::size_t c = 0; c < matrix_construct.cols; ++c)
      ref_a(r,c) = a(r,c) = static_cast<T>(matrix_construct.init_fn(r, c));

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

  //TTS_EQUAL(ops(a, s), ref);

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

  for (std::size_t r = 0; r < matrix_construct.ni; ++r)
    for (std::size_t c = 0; c < matrix_construct.nj; ++c)
      TTS_EQUAL(a_block(r, c), ref_a_block(r, c));

  for (std::size_t r = 0; r < matrix_construct.rows; ++r)
    for (std::size_t c = 0; c < matrix_construct.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>;

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

  for (std::size_t r = 0; r < matrix_construct.rows; ++r)
    for (std::size_t c = 0; c < matrix_construct.cols; ++c)
      ref_a(r,c) = a(r,c) = static_cast<T>(matrix_construct.init_fn(r, c));

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

  //TTS_EQUAL(a * s, ref);
  //TTS_EQUAL(s * a, ref);

  a_block *= scalar;
  ref_a_block *= scalar;

  for (std::size_t r = 0; r < matrix_construct.ni; ++r)
    for (std::size_t c = 0; c < matrix_construct.nj; ++c)
      TTS_EQUAL(a_block(r, c), ref_a_block(r, c));

  for (std::size_t r = 0; r < matrix_construct.rows; ++r)
    for (std::size_t c = 0; c < matrix_construct.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>;

  MatrixType a(a_matrix_construct.rows, a_matrix_construct.cols);
  MatrixType b(b_matrix_construct.rows, b_matrix_construct.cols);
  EigenMatrix ref_a(a_matrix_construct.rows, a_matrix_construct.cols);
  EigenMatrix ref_b(b_matrix_construct.rows, b_matrix_construct.cols);

  for (std::size_t r = 0; r < a_matrix_construct.rows; ++r)
    for (std::size_t c = 0; c < a_matrix_construct.cols; ++c)
      ref_a(r,c) = a(r,c) = static_cast<T>(a_matrix_construct.init_fn(r, c));

  for (std::size_t r = 0; r < b_matrix_construct.rows; ++r)
    for (std::size_t c = 0; c < b_matrix_construct.cols; ++c)
      ref_b(r,c) = b(r,c) = static_cast<T>(b_matrix_construct.init_fn(r, c));


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

  // a * b
  a_block *= b_block;
  ref_a_block *= ref_b_block;

  for (std::size_t r = 0; r < a_matrix_construct.ni; ++r)
    for (std::size_t c = 0; c < a_matrix_construct.nj; ++c)
      TTS_EQUAL(a_block(r, c), ref_a_block(r, c));

  for (std::size_t r = 0; r < a_matrix_construct.rows; ++r)
    for (std::size_t c = 0; c < a_matrix_construct.cols; ++c)
      TTS_EQUAL(a(r, c), ref_a(r, c));

  for (std::size_t r = 0; r < b_matrix_construct.rows; ++r)
    for (std::size_t c = 0; c < b_matrix_construct.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>({1 , 1, init_a, 0, 0, 1, 1}, init_b, s_op);
  test_block_matrix_operations<mat_t, T>({13 , 15, init_a, 1, 2, 3, 4}, init_b, s_op);
  test_block_matrix_operations<mat_t, T>({5 , 9, init_a, 2, 2, 2, 2}, init_b, s_op);
  test_block_matrix_operations<mat_t, T>({15 , 15, init_a, 3, 4, 5, 5}, init_b, s_op);
  test_block_matrix_operations<mat_t, T>({5 , 5, init_b, 1, 0, 3, 2}, init_a, s_op);
  test_block_matrix_operations<mat_t, T>({10, 1, init_a, 0, 0, 5, 1}, init_b, s_op);
  test_block_matrix_operations<mat_t, T>({1 , 10, init_a, 0, 0, 1, 5}, init_b, s_op);
  test_block_matrix_operations<mat_t, T>({21 , 5, init_0, 4, 4, 10, 1}, init_b, s_op);
  test_block_matrix_operations<mat_t, T>({11 , 7, init_a, 2, 0, 7, 5}, init_0, 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, s_op);
  test_block_matrix_operations<mat_t, T>({13 , 15, init_a, 1, 2, 3, 4}, init_b, s_op);
  test_block_matrix_operations<mat_t, T>({5 , 9, init_a, 2, 2, 2, 2}, init_b, s_op);
  test_block_matrix_operations<mat_t, T>({15 , 15, init_a, 3, 4, 5, 5}, init_b, s_op);
  test_block_matrix_operations<mat_t, T>({5 , 5, init_b, 1, 0, 3, 2}, init_a, s_op);
  test_block_matrix_operations<mat_t, T>({10, 1, init_a, 0, 0, 5, 1}, init_b, s_op);
  test_block_matrix_operations<mat_t, T>({1 , 10, init_a, 0, 0, 1, 5}, init_b,s_op);
  test_block_matrix_operations<mat_t, T>({21 , 5, init_0, 4, 4, 10, 1}, init_b, s_op);
  test_block_matrix_operations<mat_t, T>({11 , 7, init_a, 2, 0, 7, 5}, init_0, 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 = [](auto  r, auto 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>({1 , 1, init_a, 0, 0, 1, 1}, T{ 3.5}, s_op);
  test_block_scalar_operations<mat_t, T>({13 , 15, init_a, 1, 2, 3, 4}, T{-2.5}, s_op);
  test_block_scalar_operations<mat_t, T>({5 , 9, init_a, 2, 2, 2, 2}, T{ 42. }, s_op);
  test_block_scalar_operations<mat_t, T>({15 , 15, init_a, 3, 4, 5, 5}, T{-5. }, s_op);
  test_block_scalar_operations<mat_t, T>({5 , 5, init_b, 1, 0, 3, 2}, T{ 1. }, s_op);
  test_block_scalar_operations<mat_t, T>({10, 1, init_a, 0, 0, 5, 1}, T{ 0. }, s_op);
  test_block_scalar_operations<mat_t, T>({1 , 10, init_a, 0, 0, 1, 5}, T{ 6. },s_op);
  test_block_scalar_operations<mat_t, T>({21 , 5, init_0, 4, 4, 10, 1}, T{ 10.}, s_op);
  test_block_scalar_operations<mat_t, T>({11 , 7, init_a, 2, 0, 7, 5}, T{-0.5}, s_op);
};