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

template<typename EigenType, typename F>
void for_each_element(EigenType const& m, F&& f)
{
  for (rotgen::Index i = 0; i < m.rows(); ++i)
    for (rotgen::Index j = 0; j < m.cols(); ++j) f(i, j, m(i, j));
}

template<typename EigenType, typename F>
void for_each_element(EigenType& m, F&& f)
{
  for (rotgen::Index i = 0; i < m.rows(); ++i)
    for (rotgen::Index j = 0; j < m.cols(); ++j) f(i, j, m(i, j));
}

template<typename MatrixType, typename T>
MatrixType make_initialized_matrix(
  rotgen::tests::matrix_block_test_case<MatrixType> const& matrix_construct)
{
  auto [d, i0, j0, ni, nj] = matrix_construct;
  auto [r, c, fn] = d;

  MatrixType matrix;

  if constexpr (MatrixType::RowsAtCompileTime == -1 &&
                MatrixType::ColsAtCompileTime == -1)
    rotgen::resize(matrix, r, c);

  for (rotgen::Index i = 0; i < r; ++i)
    for (rotgen::Index j = 0; j < c; ++j)
      matrix(i, j) = static_cast<T>(fn(i, j));

  return matrix;
}

template<typename MatrixType, typename BlockType>
void validate_block_behavior(MatrixType& matrix,
                             BlockType& block,
                             rotgen::Index block_i,
                             rotgen::Index block_j,
                             rotgen::Index block_m,
                             rotgen::Index block_n)
{
  using T = typename MatrixType::value_type;
  TTS_EQUAL(block.rows(), block_m);
  TTS_EQUAL(block.cols(), block_n);
  TTS_EQUAL(block.size(), block_m * block_n);

  // test block values
  for_each_element(block, [&](auto i, auto j, auto block_val) {
    TTS_EQUAL(block_val, matrix(block_i + i, block_j + j));
  });

  // test aliasing if non immutable
  if constexpr (!BlockType::is_immutable)
  {
    T value = 1;
    for_each_element(block,
                     [&](auto i, auto j, auto&) { block(i, j) = value++; });

    value = 1;
    for_each_element(block, [&](auto i, auto j, auto&) {
      TTS_EQUAL(matrix(block_i + i, block_j + j), value++);
    });

    for_each_element(block, [&](auto i, auto j, auto&) {
      matrix(block_i + i, block_j + j) = T(42);
    });

    for_each_element(block, [&](auto&, auto&, auto block_val) {
      TTS_EQUAL(block_val, T(42));
    });
  }
}

template<typename MatrixType, typename T>
void test_dynamic_block_extraction(
  rotgen::tests::matrix_block_test_case<MatrixType> const& matrix_construct)
{
  MatrixType matrix = make_initialized_matrix<MatrixType, T>(matrix_construct);

  MatrixType const c_matrix = matrix;

  auto c_block_main =
    rotgen::extract(c_matrix, matrix_construct.i0, matrix_construct.j0,
                    matrix_construct.ni, matrix_construct.nj);
  auto c_block_top_left_corner =
    rotgen::topLeftCorner(c_matrix, matrix_construct.ni, matrix_construct.nj);
  auto c_block_top_right_corner =
    rotgen::topRightCorner(c_matrix, matrix_construct.ni, matrix_construct.nj);
  auto c_block_bottom_left_corner = rotgen::bottomLeftCorner(
    c_matrix, matrix_construct.ni, matrix_construct.nj);
  auto c_block_bottom_right_corner = rotgen::bottomRightCorner(
    c_matrix, matrix_construct.ni, matrix_construct.nj);

  auto c_block_top_rows = rotgen::topRows(c_matrix, matrix_construct.ni);
  auto c_block_middle_rows =
    rotgen::middleRows(c_matrix, matrix_construct.i0, matrix_construct.ni);
  auto c_block_bottom_rows = rotgen::bottomRows(c_matrix, matrix_construct.ni);

  auto c_block_left_cols = rotgen::leftCols(c_matrix, matrix_construct.nj);
  auto c_block_middle_cols =
    rotgen::middleCols(c_matrix, matrix_construct.j0, matrix_construct.nj);
  auto c_block_right_cols = rotgen::rightCols(c_matrix, matrix_construct.nj);

  auto block_main =
    rotgen::extract(matrix, matrix_construct.i0, matrix_construct.j0,
                    matrix_construct.ni, matrix_construct.nj);
  auto block_top_left_corner =
    rotgen::topLeftCorner(matrix, matrix_construct.ni, matrix_construct.nj);
  auto block_top_right_corner =
    rotgen::topRightCorner(matrix, matrix_construct.ni, matrix_construct.nj);
  auto block_bottom_left_corner =
    rotgen::bottomLeftCorner(matrix, matrix_construct.ni, matrix_construct.nj);
  auto block_bottom_right_corner =
    rotgen::bottomRightCorner(matrix, matrix_construct.ni, matrix_construct.nj);

  auto block_top_rows = rotgen::topRows(matrix, matrix_construct.ni);
  auto block_middle_rows =
    rotgen::middleRows(matrix, matrix_construct.i0, matrix_construct.ni);
  auto block_bottom_rows = rotgen::bottomRows(matrix, matrix_construct.ni);

  auto block_left_cols = rotgen::leftCols(matrix, matrix_construct.nj);
  auto block_middle_cols =
    rotgen::middleCols(matrix, matrix_construct.j0, matrix_construct.nj);
  auto block_right_cols = rotgen::rightCols(matrix, matrix_construct.nj);

  auto blocks = std::make_tuple(
    // --- CONST TESTS
    std::make_tuple(c_block_main, matrix_construct.i0, matrix_construct.j0,
                    matrix_construct.ni, matrix_construct.nj),
    std::make_tuple(c_block_top_left_corner, 0, 0, matrix_construct.ni,
                    matrix_construct.nj),
    std::make_tuple(c_block_top_right_corner, 0,
                    matrix.cols() - matrix_construct.nj, matrix_construct.ni,
                    matrix_construct.nj),
    std::make_tuple(c_block_bottom_left_corner,
                    matrix.rows() - matrix_construct.ni, 0, matrix_construct.ni,
                    matrix_construct.nj),
    std::make_tuple(c_block_bottom_right_corner,
                    matrix.rows() - matrix_construct.ni,
                    matrix.cols() - matrix_construct.nj, matrix_construct.ni,
                    matrix_construct.nj),

    std::make_tuple(c_block_top_rows, 0, 0, matrix_construct.ni, matrix.cols()),
    std::make_tuple(c_block_middle_rows, matrix_construct.i0, 0,
                    matrix_construct.ni, matrix.cols()),
    std::make_tuple(c_block_bottom_rows, matrix.rows() - matrix_construct.ni, 0,
                    matrix_construct.ni, matrix.cols()),

    std::make_tuple(c_block_left_cols, 0, 0, matrix.rows(),
                    matrix_construct.nj),
    std::make_tuple(c_block_middle_cols, 0, matrix_construct.j0, matrix.rows(),
                    matrix_construct.nj),
    std::make_tuple(c_block_right_cols, 0, matrix.cols() - matrix_construct.nj,
                    matrix.rows(), matrix_construct.nj),
    // --- REGULAR TESTS
    std::make_tuple(block_main, matrix_construct.i0, matrix_construct.j0,
                    matrix_construct.ni, matrix_construct.nj),
    std::make_tuple(block_top_left_corner, 0, 0, matrix_construct.ni,
                    matrix_construct.nj),
    std::make_tuple(block_top_right_corner, 0,
                    matrix.cols() - matrix_construct.nj, matrix_construct.ni,
                    matrix_construct.nj),
    std::make_tuple(block_bottom_left_corner,
                    matrix.rows() - matrix_construct.ni, 0, matrix_construct.ni,
                    matrix_construct.nj),
    std::make_tuple(block_bottom_right_corner,
                    matrix.rows() - matrix_construct.ni,
                    matrix.cols() - matrix_construct.nj, matrix_construct.ni,
                    matrix_construct.nj),

    std::make_tuple(block_top_rows, 0, 0, matrix_construct.ni, matrix.cols()),
    std::make_tuple(block_middle_rows, matrix_construct.i0, 0,
                    matrix_construct.ni, matrix.cols()),
    std::make_tuple(block_bottom_rows, matrix.rows() - matrix_construct.ni, 0,
                    matrix_construct.ni, matrix.cols()),

    std::make_tuple(block_left_cols, 0, 0, matrix.rows(), matrix_construct.nj),
    std::make_tuple(block_middle_cols, 0, matrix_construct.j0, matrix.rows(),
                    matrix_construct.nj),
    std::make_tuple(block_right_cols, 0, matrix.cols() - matrix_construct.nj,
                    matrix.rows(), matrix_construct.nj));

  std::apply(
    [&](auto&&... block_entries) {
      (([&] {
         auto&& [block, i_offset, j_offset, ni, nj] = block_entries;
         using block_t = std::remove_reference_t<decltype(block)>;

         TTS_EQUAL(block_t::RowsAtCompileTime, rotgen::Dynamic);
         TTS_EQUAL(block_t::ColsAtCompileTime, rotgen::Dynamic);
         TTS_EQUAL(block_t::storage_order, MatrixType::storage_order);

         validate_block_behavior(matrix, block, i_offset, j_offset, ni, nj);
       })(),
       ...);
    },
    blocks);
}

template<typename MatrixType, typename T, rotgen::Index NI, rotgen::Index NJ>
void test_static_block_extraction(
  rotgen::tests::static_matrix_block_test_case<MatrixType, NI, NJ> const&
    matrix_construct)
{
  auto [d, i0, j0] = matrix_construct;
  auto [r, c, fn] = d;

  MatrixType matrix(r, c);

  for (rotgen::Index i = 0; i < r; ++i)
    for (rotgen::Index j = 0; j < c; ++j)
      matrix(i, j) = static_cast<T>(fn(i, j));

  MatrixType const c_matrix = matrix;

  auto c_block_main = rotgen::extract<NI, NJ>(c_matrix, i0, j0);
  auto c_block_top_left_corner = rotgen::topLeftCorner<NI, NJ>(c_matrix);
  auto c_block_top_right_corner = rotgen::topRightCorner<NI, NJ>(c_matrix);
  auto c_block_bottom_left_corner = rotgen::bottomLeftCorner<NI, NJ>(c_matrix);
  auto c_block_bottom_right_corner =
    rotgen::bottomRightCorner<NI, NJ>(c_matrix);
  auto c_block_top_rows = rotgen::topRows<NI>(c_matrix);
  auto c_block_middle_rows = rotgen::middleRows<NI>(c_matrix, i0);
  auto c_block_bottom_rows = rotgen::bottomRows<NI>(c_matrix);
  auto c_block_left_cols = rotgen::leftCols<NJ>(c_matrix);
  auto c_block_middle_cols = rotgen::middleCols<NJ>(c_matrix, j0);
  auto c_block_right_cols = rotgen::rightCols<NJ>(c_matrix);
  auto c_block_row = rotgen::row(c_matrix, i0);
  auto c_block_col = rotgen::col(c_matrix, j0);

  auto block_main = rotgen::extract<NI, NJ>(matrix, i0, j0);
  auto block_top_left_corner = rotgen::topLeftCorner<NI, NJ>(matrix);
  auto block_top_right_corner = rotgen::topRightCorner<NI, NJ>(matrix);
  auto block_bottom_left_corner = rotgen::bottomLeftCorner<NI, NJ>(matrix);
  auto block_bottom_right_corner = rotgen::bottomRightCorner<NI, NJ>(matrix);
  auto block_top_rows = rotgen::topRows<NI>(matrix);
  auto block_middle_rows = rotgen::middleRows<NI>(matrix, i0);
  auto block_bottom_rows = rotgen::bottomRows<NI>(matrix);
  auto block_left_cols = rotgen::leftCols<NJ>(matrix);
  auto block_middle_cols = rotgen::middleCols<NJ>(matrix, j0);
  auto block_right_cols = rotgen::rightCols<NJ>(matrix);
  auto block_row = rotgen::row(matrix, i0);
  auto block_col = rotgen::col(matrix, j0);

  auto blocks = std::make_tuple(

    // ----- TEST ON BLOCK FROM CONST MATRIX
    std::make_tuple(c_block_main, i0, j0, matrix_construct.ni,
                    matrix_construct.nj, int(NI), int(NJ)),
    std::make_tuple(c_block_top_left_corner, 0, 0, matrix_construct.ni,
                    matrix_construct.nj, int(NI), int(NJ)),
    std::make_tuple(c_block_top_right_corner, 0,
                    matrix.cols() - matrix_construct.nj, matrix_construct.ni,
                    matrix_construct.nj, int(NI), int(NJ)),
    std::make_tuple(c_block_bottom_left_corner,
                    matrix.rows() - matrix_construct.ni, 0, matrix_construct.ni,
                    matrix_construct.nj, int(NI), int(NJ)),
    std::make_tuple(c_block_bottom_right_corner,
                    matrix.rows() - matrix_construct.ni,
                    matrix.cols() - matrix_construct.nj, matrix_construct.ni,
                    matrix_construct.nj, int(NI), int(NJ)),

    std::make_tuple(c_block_top_rows, 0, 0, matrix_construct.ni, matrix.cols(),
                    int(NI), rotgen::Dynamic),
    std::make_tuple(c_block_middle_rows, i0, 0, matrix_construct.ni,
                    matrix.cols(), int(NI), rotgen::Dynamic),
    std::make_tuple(c_block_bottom_rows, matrix.rows() - matrix_construct.ni, 0,
                    matrix_construct.ni, matrix.cols(), int(NI),
                    rotgen::Dynamic),

    std::make_tuple(c_block_left_cols, 0, 0, matrix.rows(), matrix_construct.nj,
                    rotgen::Dynamic, int(NJ)),
    std::make_tuple(c_block_middle_cols, 0, j0, matrix.rows(),
                    matrix_construct.nj, rotgen::Dynamic, int(NJ)),
    std::make_tuple(c_block_right_cols, 0, matrix.cols() - matrix_construct.nj,
                    matrix.rows(), matrix_construct.nj, rotgen::Dynamic,
                    int(NJ)),

    std::make_tuple(c_block_row, i0, 0, 1, matrix.cols(), 1, rotgen::Dynamic),
    std::make_tuple(c_block_col, 0, j0, matrix.rows(), 1, rotgen::Dynamic, 1),
    // -- Block to NON CONST
    std::make_tuple(block_main, i0, j0, matrix_construct.ni,
                    matrix_construct.nj, int(NI), int(NJ)),
    std::make_tuple(block_top_left_corner, 0, 0, matrix_construct.ni,
                    matrix_construct.nj, int(NI), int(NJ)),
    std::make_tuple(block_top_right_corner, 0,
                    matrix.cols() - matrix_construct.nj, matrix_construct.ni,
                    matrix_construct.nj, int(NI), int(NJ)),
    std::make_tuple(block_bottom_left_corner,
                    matrix.rows() - matrix_construct.ni, 0, matrix_construct.ni,
                    matrix_construct.nj, int(NI), int(NJ)),
    std::make_tuple(block_bottom_right_corner,
                    matrix.rows() - matrix_construct.ni,
                    matrix.cols() - matrix_construct.nj, matrix_construct.ni,
                    matrix_construct.nj, int(NI), int(NJ)),

    std::make_tuple(block_top_rows, 0, 0, matrix_construct.ni, matrix.cols(),
                    int(NI), rotgen::Dynamic),
    std::make_tuple(block_middle_rows, i0, 0, matrix_construct.ni,
                    matrix.cols(), int(NI), rotgen::Dynamic),
    std::make_tuple(block_bottom_rows, matrix.rows() - matrix_construct.ni, 0,
                    matrix_construct.ni, matrix.cols(), int(NI),
                    rotgen::Dynamic),

    std::make_tuple(block_left_cols, 0, 0, matrix.rows(), matrix_construct.nj,
                    rotgen::Dynamic, int(NJ)),
    std::make_tuple(block_middle_cols, 0, j0, matrix.rows(),
                    matrix_construct.nj, rotgen::Dynamic, int(NJ)),
    std::make_tuple(block_right_cols, 0, matrix.cols() - matrix_construct.nj,
                    matrix.rows(), matrix_construct.nj, rotgen::Dynamic,
                    int(NJ)),

    std::make_tuple(block_row, i0, 0, 1, matrix.cols(), 1, rotgen::Dynamic),
    std::make_tuple(block_col, 0, j0, matrix.rows(), 1, rotgen::Dynamic, 1));

  std::apply(
    [&](auto&&... block_entries) {
      (([&] {
         auto&& [block, i_offset, j_offset, ni, nj, rows_ct, cols_ct] =
           block_entries;
         using block_t = std::remove_reference_t<decltype(block)>;

         TTS_EQUAL(block_t::RowsAtCompileTime, rows_ct);
         TTS_EQUAL(block_t::ColsAtCompileTime, cols_ct);

         validate_block_behavior(matrix, block, i_offset, j_offset, ni, nj);
       })(),
       ...);
    },
    blocks);
}

TTS_CASE_TPL(
  "Check all dynamic block extractions on a dynamic row-major matrix",
  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, 1>;

  std::vector<rotgen::tests::matrix_block_test_case<mat_t>> cases = {
    {{6, 11, [](rotgen::Index i, rotgen::Index j) { return T(i * 10 + j); }},
     1,
     2,
     3,
     2},
    {{7, 10,
      [](rotgen::Index i, rotgen::Index j) { return T(std::sin(i + j)); }},
     4,
     4,
     3,
     3},
    {{5, 5, [](rotgen::Index i, rotgen::Index j) { return T((i + j) % 7); }},
     0,
     0,
     5,
     5},
    {{9, 14, [](rotgen::Index i, rotgen::Index j) { return T(i + j + 3 * j); }},
     3,
     7,
     1,
     1}};

  for (auto const& matrix_case : cases)
  {
    test_dynamic_block_extraction<mat_t, T>(matrix_case);
  }
};

TTS_CASE_TPL(
  "Check all dynamic block extractions on a static column-major matrix",
  rotgen::tests::types)<typename T, typename O>(tts::type<tts::types<T, O>>)
{
  using mat_t = rotgen::matrix<T, 4, 5, O::value, 0>;

  std::vector<rotgen::tests::matrix_block_test_case<mat_t>> cases = {
    {{4, 5, [](auto i, auto j) { return T(2 * i + j * j * j - 42); }},
     1,
     2,
     3,
     2},
    {{4, 5, [](auto i, auto j) { return T(std::tan(i * i * j)); }}, 0, 1, 2, 1},
    {{4, 5, [](auto i, auto j) { return T((i * i + j * j) / 6); }},
     2,
     0,
     0,
     0}};

  for (auto const& matrix_case : cases)
  {
    test_dynamic_block_extraction<mat_t, T>(matrix_case);
  }
};

TTS_CASE_TPL("Check all static block extractions",
             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_static_block_extraction<mat_t, T, 1, 2>(
    rotgen::tests::static_matrix_block_test_case<mat_t, 1, 2>{
      {11, 11,
       [](rotgen::Index i, rotgen::Index j) {
         return T(i * i * i + 3 * j - 127);
       }},
      3,
      2});

  test_static_block_extraction<mat_t, double, 4, 3>(
    rotgen::tests::static_matrix_block_test_case<mat_t, 4, 3>{
      {14, 15,
       [](rotgen::Index i, rotgen::Index j) { return T(std::cos(i * j * 2)); }},
      5,
      1});

  test_static_block_extraction<mat_t, double, 0, 0>(
    rotgen::tests::static_matrix_block_test_case<mat_t, 0, 0>{
      {5, 5, [](rotgen::Index i, rotgen::Index j) { return T((i + j) % 9); }},
      0,
      0});
};

TTS_CASE_TPL("Check vector-only extractions",
             rotgen::tests::types)<typename T, typename O>(
  tts::type<tts::types<T, O>>)
{
  auto run_case = [](auto&& matrix, auto&& block, int i_offset, int j_offset,
                     int ni, int nj, auto const& rows_ct, auto const& cols_ct) {
    using block_t = std::remove_reference_t<decltype(block)>;

    TTS_EQUAL(block_t::RowsAtCompileTime, rows_ct);
    TTS_EQUAL(block_t::ColsAtCompileTime, cols_ct);

    validate_block_behavior(matrix, block, i_offset, j_offset, ni, nj);
  };

  if constexpr (O::value == rotgen::RowMajor)
  {
    using mat_t = rotgen::matrix<T, 1, rotgen::Dynamic, O::value>;
    mat_t m(1, 11);

    for (rotgen::Index i = 0; i < 11; ++i) m(i) = 1 + i;

    run_case(m, head(m, 1), 0, 0, 1, 1, 1, rotgen::Dynamic);
    run_case(m, head(m, 5), 0, 0, 1, 5, 1, rotgen::Dynamic);
    run_case(m, head(m, 11), 0, 0, 1, 11, 1, rotgen::Dynamic);
    run_case(m, head<1>(m), 0, 0, 1, 1, 1, 1);
    run_case(m, head<4>(m), 0, 0, 1, 4, 1, 4);
    run_case(m, head<9>(m), 0, 0, 1, 9, 1, 9);
    run_case(m, tail(m, 1), 0, 10, 1, 1, 1, rotgen::Dynamic);
    run_case(m, tail(m, 5), 0, 6, 1, 5, 1, rotgen::Dynamic);
    run_case(m, tail(m, 11), 0, 0, 1, 11, 1, rotgen::Dynamic);
    run_case(m, tail<1>(m), 0, 10, 1, 1, 1, 1);
    run_case(m, tail<4>(m), 0, 7, 1, 4, 1, 4);
    run_case(m, tail<9>(m), 0, 2, 1, 9, 1, 9);
    run_case(m, segment(m, 0, 11), 0, 0, 1, 11, 1, rotgen::Dynamic);
    run_case(m, segment(m, 0, 7), 0, 0, 1, 7, 1, rotgen::Dynamic);
    run_case(m, segment(m, 5, 6), 0, 5, 1, 6, 1, rotgen::Dynamic);
    run_case(m, segment<11>(m, 0), 0, 0, 1, 11, 1, 11);
    run_case(m, segment<7>(m, 0), 0, 0, 1, 7, 1, 7);
    run_case(m, segment<6>(m, 5), 0, 5, 1, 6, 1, 6);
  }
  else
  {
    using mat_t = rotgen::matrix<T, rotgen::Dynamic, 1, O::value>;
    mat_t m(11, 1);

    for (rotgen::Index i = 0; i < 11; ++i) m(i) = 1 + i;

    run_case(m, head(m, 1), 0, 0, 1, 1, rotgen::Dynamic, 1);
    run_case(m, head(m, 5), 0, 0, 5, 1, rotgen::Dynamic, 1);
    run_case(m, head(m, 11), 0, 0, 11, 1, rotgen::Dynamic, 1);
    run_case(m, head<1>(m), 0, 0, 1, 1, 1, 1);
    run_case(m, head<4>(m), 0, 0, 4, 1, 4, 1);
    run_case(m, head<9>(m), 0, 0, 9, 1, 9, 1);
    run_case(m, tail(m, 1), 10, 0, 1, 1, rotgen::Dynamic, 1);
    run_case(m, tail(m, 5), 6, 0, 5, 1, rotgen::Dynamic, 1);
    run_case(m, tail(m, 11), 0, 0, 11, 1, rotgen::Dynamic, 1);
    run_case(m, tail<1>(m), 10, 0, 1, 1, 1, 1);
    run_case(m, tail<4>(m), 7, 0, 4, 1, 4, 1);
    run_case(m, tail<9>(m), 2, 0, 9, 1, 9, 1);
  }
};