#ifndef JACOBI_HPP
#define JACOBI_HPP

#include <cmath>
#include <type_traits>
#include <hpc/matvec/isgematrix.h>

template<typename Matrix>
typename std::enable_if<hpc::matvec::IsGeMatrix<Matrix>::value, typename Matrix::ElementType>::type
single_jacobi_iteration(const Matrix& A, Matrix& B) {
   using T = typename Matrix::ElementType;
   using Index = typename Matrix::Index;
   T maxdiff{};
   #pragma omp parallel for reduction(max:maxdiff)
   for (Index i = 1; i < A.numRows - 1; ++i) {
      for (Index j = 1; j < A.numCols - 1; ++j) {
	 B(i, j) = 0.25 * (A(i-1, j) + A(i, j-1) + A(i, j+1) + A(i+1, j));
	 T diff = std::abs(A(i, j) - B(i, j));
	 if (diff > maxdiff) maxdiff = diff;
      }
   }
   return maxdiff;
}

template<typename G, typename Matrix>
typename std::enable_if<hpc::matvec::IsGeMatrix<Matrix>::value, void>::type
jacobi(G g, Matrix& A, typename Matrix::ElementType epsilon) {
   using T = typename Matrix::ElementType;
   using Index = typename Matrix::Index;
   // we do not copy-construct B here to avoid a double initialization
   using namespace hpc::matvec;
   GeMatrix<T> B(A.numRows, A.numCols, StorageOrder::RowMajor);
   #pragma omp parallel for
   for (Index i = 0; i < A.numRows; ++i) {
      for (Index j = 0; j < A.numCols; ++j) {
	 A(i, j) = B(i, j) =
	    g(static_cast<T>(i)/(A.numRows-1),
	       static_cast<T>(j)/(A.numCols-1));
      }
   }
   T maxdiff;
   do {
      single_jacobi_iteration(A, B);
      maxdiff = single_jacobi_iteration(B, A);
   } while (maxdiff > epsilon);
}

#endif