#ifndef JACOBI_HPP
#define JACOBI_HPP

#include <cassert>
#include <thread>
#include <vector>
#include "barrier.hpp"
#include "matrix.hpp"
#include "jacobi-thread.hpp"

template<typename Real>
class Jacobi {
   public:
      template<typename G>
      Jacobi(G g, unsigned int n, Real eps, unsigned int nofthreads) :
	    n(n),
	    a0(n+2, n+2, [=](unsigned int i, unsigned int j) -> double {
	       return g((Real)i/(n+1), (Real)j/(n+1));
	    }),
	    a1(n+2, n+2, [=](unsigned int i, unsigned int j) -> double {
	       return g((Real)i/(n+1), (Real)j/(n+1));
	    }),
	    eps(eps), nofthreads(nofthreads),
	    threads(nofthreads),
	    barrier([](Real a, Real b) -> Real {
	       return a > b? a: b;
	    }, 0, nofthreads),
	    waited(false) {
	 assert(n > 0); assert(eps > 0); assert(nofthreads < n);
	 /* 1d partitioning */
	 unsigned int remainder = n % nofthreads;
	 for (unsigned int index = 0; index < nofthreads; ++index) {
	    unsigned int nofrows = n / nofthreads;
	    unsigned int first_row = index * nofrows + 1;
	    if (index < remainder) {
	       ++nofrows;
	       if (index > 0) first_row += index;
	    } else {
	       first_row += remainder;
	    }
	    threads[index] = std::thread(JacobiThread<Real, Matrix<Real>>
	       (a0.get_block(first_row-1, 0, nofrows+2, n+2),
		  a1.get_block(first_row-1, 0, nofrows+2, n+2),
		  eps, barrier));
	 }
      }
      const Matrix<Real>& get_result() {
	 if (!waited) {
	    for (auto& t: threads) {
	       t.join();
	    }
	    waited = true;
	 }
	 if (barrier.get_iterations() % 2 == 0) {
	    return a0;
	 } else {
	    return a1;
	 }
      }
   private:
      unsigned int n;
      Matrix<Real> a0, a1;
      Real eps;
      unsigned int nofthreads;
      std::vector<std::thread> threads;
      AggregatingBarrier<Real> barrier;
      bool waited;
};

#endif