#include <cstdlib>
#include <random>
#include <thread>
#include <vector>
#include <hpc/matvec/gematrix.hpp>
#include <hpc/matvec/iterators.hpp>
#include <hpc/matvec/print.hpp>
#include <hpc/matvec/traits.hpp>

using namespace hpc;
using namespace hpc::matvec;

template<typename T>
struct Slices {
   Slices(T nof_threads, T problem_size) :
	 nof_threads((assert(nof_threads > 0), nof_threads)),
	 problem_size(problem_size),
	 remainder(problem_size % nof_threads),
	 slice_size(problem_size / nof_threads) {
   }
   T offset(T index) {
      assert(index < nof_threads);
      if (index < remainder) {
	 return index * (slice_size + 1);
      } else {
	 return remainder * (slice_size + 1) +
	        (index - remainder) * slice_size;
      }
   }
   T size(T index) {
      assert(index < nof_threads);
      if (index < remainder) {
	 return slice_size + 1;
      } else {
	 return slice_size;
      }
   }
   T nof_threads; T problem_size;
   T remainder; T slice_size;
};

template <
   template<typename> class MatrixA, typename T,
   Require< Ge<MatrixA<T>> > = true
>
void randomInit(MatrixA<T>& A) {
   std::random_device random;
   std::mt19937 mt{random()};
   std::uniform_real_distribution<T> uniform(-100, 100);
   
   for (auto [i, j, Aij]: A) {
      Aij = uniform(mt);
      (void) i; (void) j; // suppress gcc warning
   }
}

int main() {
   GeMatrix<double> A(51, 7);
   auto nof_threads = std::thread::hardware_concurrency();

   std::vector<std::thread> threads(nof_threads);
   Slices<std::size_t> slices(nof_threads, A.numRows());
   for (unsigned int index = 0; index < nof_threads; ++index) {
      auto firstRow = slices.offset(index);
      auto numRows = slices.size(index);
      threads[index] = std::thread(
	 [A_ = A.block(firstRow, 0).dim(numRows, A.numCols())]() mutable {
	    randomInit(A_);
	 }
      );
   }
   for (unsigned int index = 0; index < nof_threads; ++index) {
      threads[index].join();
   }
   fmt::printf("A:\n"); print(A);
}