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

using namespace std;

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;
   double global_max;
   do {
      // compute border zones
      for (int j = 1; j <= n; ++j) {
	 B[1][j] = 0.25 * (A[0][j] + A[1][j-1] + A[1][j+1] + A[2][j]);

	 B[nofrows][j] = 0.25 * (A[nofrows-1][j] + A[nofrows][j-1] +
	                         A[nofrows][j+1] + A[nofrows+1][j]);
      }

      // initiate non-blocking communication
      MPI_Request req[4];
      MPI_Irecv(A[0] + 1, n, MPI_DOUBLE, previous, 0,
	 MPI_COMM_WORLD, &req[0]);
      MPI_Irecv(A[nofrows+1] + 1, n, MPI_DOUBLE, next, 0,
	 MPI_COMM_WORLD, &req[1]);
      MPI_Isend(B[1] + 1, n, MPI_DOUBLE, previous, 0,
	 MPI_COMM_WORLD, &req[2]);
      MPI_Isend(B[nofrows] + 1, n, MPI_DOUBLE, next, 0,
	 MPI_COMM_WORLD, &req[3]);

      // computer inner zone
      for (int i = 2; i < nofrows; ++i) {
	 for (int j = 1; j <= n; ++j) {
	    B[i][j] = 0.25 * (A[i-1][j] + A[i][j-1] + A[i][j+1] + A[i+1][j]);
	 }
      }
      // prepare next iteration and compute maxdiff
      double maxdiff = 0;
      for (int i = 1; i <= nofrows; ++i) {
	 for (int j = 1; j <= n; ++j) {
	    double diff = fabs(A[i][j] - B[i][j]);
	    if (diff > maxdiff) maxdiff = diff;
	    A[i][j] = B[i][j];
	 }
      }

      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);

      // block until initiated communication is finished
      for (int i = 0; i < 4; ++i) {
	 MPI_Status status;
	 MPI_Wait(&req[i], &status);
      }
   } while (global_max > eps);

   // 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();
}