#include <cmath>
#include <cassert>
#include <cstring>
#include <iostream>
#include <sstream>
#include <mpi.h>

using namespace std;

// single Jacobi iteration step
double single_jacobi_iteration(double** A, double** B, int n, int m) {
   for (int i = 1; i <= n; ++i) {
      for (int j = 1; j <= m; ++j) {
	 B[i][j] = 0.25 * (A[i-1][j] + A[i][j-1] + A[i][j+1] + A[i+1][j]);
      }
   }
   double maxdiff = 0;
   for (int i = 1; i <= n; ++i) {
      for (int j = 1; j <= m; ++j) {
	 double diff = fabs(A[i][j] - B[i][j]);
	 if (diff > maxdiff) maxdiff = diff;
	 A[i][j] = B[i][j];
      }
   }
   return maxdiff;
}

void initialize_A(double& Aij, int i, int j, int N) {
   const static double E_POWER_MINUS_PI = pow(M_E, -M_PI);
   if (j == 0) {
      Aij = sin(M_PI * ((double)i/(N-1)));
   } else if (j == N-1) {
      Aij = sin(M_PI * ((double)i/(N-1))) * E_POWER_MINUS_PI;
   } else {
      Aij = 0;
   }
}

// 1D-partitioned task
double** run_jacobi_iteration(int rank, int nofprocesses, int N, double eps) {
   int n = N-2;
   assert(nofprocesses <= n);
   int nofrows = n / nofprocesses;
   int remainder = n % nofprocesses;
   int first_row = rank * nofrows + 1;
   if (rank < remainder) {
      ++nofrows;
      if (rank > 0) first_row += rank;
   } else {
      first_row += remainder;
   }
   int last_row = first_row + nofrows - 1;
   double** A = new double*[nofrows+2];
   for (int i = 0; i <= nofrows+1; ++i) {
      A[i] = new double[N];
      for (int j = 0; j < N; ++j) {
	 initialize_A(A[i][j], i + first_row, j, N);
      }
   }
   double** B = new double*[nofrows+1];
   for (int i = 1; i <= nofrows; ++i) {
      B[i] = new double[N-1];
   }

   int previous = rank == 0? MPI_PROC_NULL: rank-1;
   int next = rank == nofprocesses-1? MPI_PROC_NULL: rank+1;
   for(;;) {
      double maxdiff = single_jacobi_iteration(A, B, nofrows, n);
      double global_max;
      MPI_Reduce(&maxdiff, &global_max, 1, MPI_DOUBLE,
	 MPI_MAX, 0, MPI_COMM_WORLD);
      MPI_Bcast(&global_max, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
      if (global_max < eps) break;
      MPI_Status status;
      // send highest row to the process which is next in rank
      MPI_Sendrecv(A[nofrows] + 1, n, MPI_DOUBLE, next, 0,
	 A[0] + 1, n, MPI_DOUBLE, previous, 0, MPI_COMM_WORLD, &status);
      // send lowest row to the process which is previous in rank
      MPI_Sendrecv(A[1] + 1, n, MPI_DOUBLE, previous, 0,
	 A[nofrows+1] + 1, n, MPI_DOUBLE, next, 0, MPI_COMM_WORLD, &status);
   }

   // collect results in process 0
   double** Result = 0;
   if (rank == 0) {
      Result = new double*[N]; assert(Result);
      for (int i = 0; i < N; ++i) {
	 Result[i] = new double[N]; assert(Result[i]);
	 for (int j = 0; j < N; ++j) {
	    initialize_A(Result[i][j], i, j, N);
	 }
      }
      for (int i = 1; i <= last_row; ++i) {
	 memcpy(Result[i] + 1, A[i] + 1, n * sizeof(double));
      }
      for (int i = last_row+1; i <= n; ++i) {
	 MPI_Status status;
	 MPI_Recv(Result[i] + 1, n, MPI_DOUBLE,
	    MPI_ANY_SOURCE, i, MPI_COMM_WORLD, &status);
      }
   } else {
      for (int i = 1; i <= nofrows; ++i) {
	 MPI_Send(A[i] + 1, n, MPI_DOUBLE, 0, first_row + i - 1,
	    MPI_COMM_WORLD);
      }
   }
   return Result;
}

char* cmdname;
void usage() {
   cerr << "Usage: " << cmdname << " [N [eps]] " << endl;
   MPI_Abort(MPI_COMM_WORLD, 1);
}

int main(int argc, char** argv) {
   int N = 10;
   double eps = 1e-6;

   MPI_Init(&argc, &argv);
   int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank);
   int nofprocesses; MPI_Comm_size(MPI_COMM_WORLD, &nofprocesses);

   if (rank == 0) {
      cmdname = *argv++; --argc;
      if (argc > 2) usage();
      if (argc > 0) {
	 istringstream arg(*argv++); --argc;
	 if (!(arg >> N) || N < 3 || nofprocesses > N-2) usage();
      }
      if (argc > 0) {
	 istringstream arg(*argv++); --argc;
	 if (!(arg >> eps) || eps <= 0) usage();
      }
   }
   MPI_Bcast(&N, 1, MPI_INT, 0, MPI_COMM_WORLD);
   MPI_Bcast(&eps, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);

   double** A = run_jacobi_iteration(rank, nofprocesses, N, eps);

   if (rank == 0) {
      cout << N << endl;
      for (int i = 0; i < N; ++i) {
	 for (int j = 0; j < N; ++j) {
	    cout << " " << A[i][j];
	 }
	 cout << endl;
	 delete[] A[i];
      }
      delete[] A;
   }

   MPI_Finalize();
}