//==============================================================================
/*
  ROTGEN - Runtime Overlay for Eigen
  Copyright : CODE RECKONS
  SPDX-License-Identifier: BSL-1.0
*/
//==============================================================================
#pragma once

#include <rotgen/rotgen.hpp>

#include <Eigen/Dense>
#include <tuple>
#include <vector>

namespace rotgen::tests
{
  struct matrix_descriptor
  {
    rotgen::Index rows, cols;
    std::function<double(std::size_t, std::size_t)> init_fn;
  };

  template<rotgen::Index Rows, rotgen::Index Cols>
  struct static_matrix_descriptor
  {
    static constexpr rotgen::Index rows = Rows;
    static constexpr rotgen::Index cols = Cols;
    std::function<double(std::size_t, std::size_t)> init_fn;
  };

  template<typename MatrixType> struct matrix_block_test_case
  {
    matrix_descriptor matrix_desc;
    rotgen::Index i0, j0, ni, nj;
  };

  template<typename MatrixType, rotgen::Index NI, rotgen::Index NJ>
  struct static_matrix_block_test_case
  {
    matrix_descriptor matrix_desc;
    static constexpr rotgen::Index ni = NI;
    static constexpr rotgen::Index nj = NJ;
    rotgen::Index i0, j0;
  };

  void prepare(rotgen::Index rows,
               rotgen::Index cols,
               auto const& init_fn,
               auto& output)
  {
    for (rotgen::Index r = 0; r < rows; ++r)
      for (rotgen::Index c = 0; c < cols; ++c) output(r, c) = init_fn(r, c);
  }

  void prepare(auto fn, auto& output)
  {
    for (rotgen::Index r = 0; r < output.rows(); ++r)
      for (rotgen::Index c = 0; c < output.cols(); ++c) output(r, c) = fn(r, c);
  }

  auto default_init_function = [](auto row, auto col) {
    return row + 3 * col - 2.5;
  };

  auto generate_matrix_references()
  {
    std::vector<rotgen::tests::matrix_descriptor> cases = {
      // Singular matrix
      {1, 1, default_init_function},
      // Square matrix below MAX_SIZE
      {3, 3, default_init_function},
      // Square matrix at MAX_SIZE
      {4, 4, default_init_function},
      // Square matrix above MAX_SIZE
      {7, 7, default_init_function},
      // Tall matrix below MAX_SIZE
      {5, 2, default_init_function},
      // Tall matrix at MAX_SIZE
      {8, 2, default_init_function},
      // Tall matrix above MAX_SIZE
      {10, 3, default_init_function},
      // Thick matrix below MAX_SIZE
      {2, 5, default_init_function},
      // Thick matrix at MAX_SIZE
      {2, 8, default_init_function},
      // Thick matrix above MAX_SIZE
      {3, 10, default_init_function}};

    return cases;
  }

  auto generate_static_matrix_references()
  {
    auto fn = [](auto r, auto c) { return r + 3 * c - 2.5; };
    std::tuple cases = {// Singular matrix
                        static_matrix_descriptor<1, 1>{fn},
                        // Square matrix below MAX_SIZE
                        static_matrix_descriptor<3, 3>{fn},
                        // Square matrix at MAX_SIZE
                        static_matrix_descriptor<4, 4>{fn},
                        // Square matrix above MAX_SIZE
                        static_matrix_descriptor<7, 7>{fn},
                        // Tall matrix below MAX_SIZE
                        static_matrix_descriptor<5, 2>{fn},
                        // Tall matrix at MAX_SIZE
                        static_matrix_descriptor<8, 2>{fn},
                        // Tall matrix above MAX_SIZE
                        static_matrix_descriptor<10, 3>{fn},
                        // Thick matrix below MAX_SIZE
                        static_matrix_descriptor<2, 5>{fn},
                        // Thick matrix at MAX_SIZE
                        static_matrix_descriptor<2, 8>{fn},
                        // Thick matrix above MAX_SIZE
                        static_matrix_descriptor<3, 10>{fn}};

    return cases;
  }

  template<typename T, typename O> auto generate_block_references()
  {
    rotgen::tests::matrix_descriptor base{
      16, 16, [](auto r, auto c) { return r + 3 * c - 2.5; }};

    using mat_t = rotgen::matrix<T, rotgen::Dynamic, rotgen::Dynamic, O::value>;

    std::vector<rotgen::tests::matrix_block_test_case<mat_t>> cases = {
      // Singular blocks
      {base, 0, 0, 1, 1},
      {base, 3, 2, 1, 1},
      {base, 4, 11, 1, 1},
      {base, 15, 15, 1, 1},
      // Square block below MAX_SIZE
      {base, 7, 9, 3, 3},
      // Square block at MAX_SIZE
      {base, 6, 3, 4, 4},
      // Square block above MAX_SIZE
      {base, 6, 3, 6, 6},
      // Tall block below MAX_SIZE
      {base, 6, 3, 5, 2},
      // Tall block at MAX_SIZE
      {base, 6, 3, 8, 2},
      // Tall block above MAX_SIZE
      {base, 6, 3, 7, 3},
      // Thick block below MAX_SIZE
      {base, 6, 3, 2, 5},
      // Thick block at MAX_SIZE
      {base, 6, 3, 2, 8},
      // Thick block above MAX_SIZE
      {base, 6, 3, 3, 7}};

    return cases;
  }

  template<typename T, typename O> auto generate_static_block_references()
  {
    rotgen::tests::matrix_descriptor base{
      16, 16, [](auto r, auto c) { return r + 3 * c - 2.5; }};

    using mat_t = rotgen::matrix<T, rotgen::Dynamic, rotgen::Dynamic, O::value>;

    std::tuple cases = {
      // Singular blocks
      static_matrix_block_test_case<mat_t, 1, 1>{base, 0, 0},
      static_matrix_block_test_case<mat_t, 1, 1>{base, 3, 2},
      static_matrix_block_test_case<mat_t, 1, 1>{base, 4, 11},
      static_matrix_block_test_case<mat_t, 1, 1>{base, 15, 15},
      // Square block below MAX_SIZE
      static_matrix_block_test_case<mat_t, 3, 3>{base, 7, 9},
      // Square block at MAX_SIZE
      static_matrix_block_test_case<mat_t, 4, 4>{base, 6, 3},
      // Square block above MAX_SIZE
      static_matrix_block_test_case<mat_t, 6, 6>{base, 6, 3},
      // Tall block below MAX_SIZE
      static_matrix_block_test_case<mat_t, 5, 2>{base, 6, 3},
      // Tall block at MAX_SIZE
      static_matrix_block_test_case<mat_t, 8, 2>{base, 6, 3},
      // Tall block above MAX_SIZE
      static_matrix_block_test_case<mat_t, 7, 3>{base, 6, 3},
      // Thick block below MAX_SIZE
      static_matrix_block_test_case<mat_t, 2, 5>{base, 6, 3},
      // Thick block at MAX_SIZE
      static_matrix_block_test_case<mat_t, 2, 8>{base, 6, 3},
      // Thick block above MAX_SIZE
      static_matrix_block_test_case<mat_t, 3, 7>{base, 6, 3}};

    return cases;
  }
}