//==================================================================================================
/*
  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 i0, rotgen::Index j0,
                             rotgen::Index ni, rotgen::Index nj)
{
  using T = typename MatrixType::value_type;
  TTS_EQUAL(block.rows(), ni);
  TTS_EQUAL(block.cols(), nj);
  TTS_EQUAL(block.size(), ni * nj);

  // test block values
  for_each_element(block, [&](auto i, auto j, auto val) {
    TTS_EQUAL(val, matrix(i0 + i, j0 + 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(i0 + i, j0 + j), value++);
    });

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

    for_each_element(block, [&](auto&, auto&, auto val) {
      TTS_EQUAL(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 refs = []()
  {
    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;

      return std::make_tuple( m,
              std::make_tuple
              (
                std::make_tuple(head(m, 1), 0, 0, 1, 1, 1, rotgen::Dynamic)
              , std::make_tuple(head(m, 5), 0, 0, 1, 5, 1, rotgen::Dynamic)
              , std::make_tuple(head(m,11), 0, 0, 1,11, 1, rotgen::Dynamic)
              , std::make_tuple(head<1>(m), 0, 0, 1, 1, 1, 1)
              , std::make_tuple(head<4>(m), 0, 0, 1, 4, 1, 4)
              , std::make_tuple(head<9>(m), 0, 0, 1, 9, 1, 9)
              , std::make_tuple(tail(m, 1), 0,10, 1, 1, 1, rotgen::Dynamic)
              , std::make_tuple(tail(m, 5), 0, 6, 1, 5, 1, rotgen::Dynamic)
              , std::make_tuple(tail(m,11), 0, 0, 1,11, 1, rotgen::Dynamic)
              , std::make_tuple(tail<1>(m), 0,10, 1, 1, 1, 1)
              , std::make_tuple(tail<4>(m), 0, 7, 1, 4, 1, 4)
              , std::make_tuple(tail<9>(m), 0, 2, 1, 9, 1, 9)
              , std::make_tuple(segment(m,0,11), 0, 0, 1, 11, 1, rotgen::Dynamic)
              , std::make_tuple(segment(m,0, 7), 0, 0, 1,  7, 1, rotgen::Dynamic)
              , std::make_tuple(segment(m,5, 6), 0, 5, 1,  6, 1, rotgen::Dynamic)
              , std::make_tuple(segment<11>(m,0), 0, 0, 1, 11, 1, 11)
              , std::make_tuple(segment< 7>(m,0), 0, 0, 1,  7, 1,  7)
              , std::make_tuple(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;

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

  auto matrix = std::get<0>(refs);
  auto blocks = std::get<1>(refs);

  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);
};