#include <thread>
#include <random>
#include <vector>
#include <hpc/matvec/gematrix.h>
#include <hpc/matvec/apply.h>
#include <hpc/matvec/print.h>
#include <hpc/aux/walltime.h>

template<typename MA>
typename std::enable_if<hpc::matvec::IsRealGeMatrix<MA>::value, void>::type
randomInit(MA& A) {
   using ElementType = typename MA::ElementType;
   using Index = typename MA::Index;

   std::random_device random;
   std::mt19937 mt(random());
   std::uniform_real_distribution<ElementType> uniform(-100,100);
   
   hpc::matvec::apply(A, [&](ElementType& val, Index i, Index j) -> void {
      val = uniform(mt);
   });
}

volatile double global;

int main() {
   using namespace hpc::matvec;
   using namespace hpc::aux;
   unsigned int nof_threads = std::thread::hardware_concurrency();

   WallTime<double> wall_time;

   constexpr unsigned int N = 5005;

   wall_time.tic();
   {
      GeMatrix<double> A(N, N);
      std::vector<std::thread> threads(nof_threads);
      unsigned int job_size = A.numRows / nof_threads;
      unsigned int remainder = A.numRows % nof_threads;
      GeMatrix<double>::Index start = 0;
      for (int index = 0; index < nof_threads; ++index) {
	 GeMatrix<double>::Index num_rows = job_size;
	 if (index < remainder) ++num_rows;
	 auto A_ = A(start, 0, num_rows, A.numCols);
	 threads[index] = std::thread([=]() mutable { randomInit(A_); });
	 start += num_rows;
      }
      for (int index = 0; index < nof_threads; ++index) {
	 threads[index].join();
      }
      global = A(7, 2);
   }
   double t1 = wall_time.toc();

   wall_time.tic();
   {
      GeMatrix<double> A(N, N);
      std::vector<std::thread> threads(nof_threads);
      unsigned int job_size = A.numRows / nof_threads;
      unsigned int remainder = A.numRows % nof_threads;
      for (int index = 0; index < nof_threads; ++index) {
	 GeMatrix<double>::Index num_rows = job_size;
	 if (index < remainder) ++num_rows;
	 auto A_ = GeMatrixView<double>(num_rows, A.numCols,
	    &A(index, 0), A.incRow * nof_threads, A.incCol);
	 threads[index] = std::thread([=]() mutable { randomInit(A_); });
      }
      for (int index = 0; index < nof_threads; ++index) {
	 threads[index].join();
      }
      global = A(7, 2);
   }
   double t2 = wall_time.toc();
   
   printf("t1 = %4.3lf, t2 = %4.3lf\n", t1, t2);
}