#include <mpi.h>
#include <math.h>
#include <time.h>
#include <stdio.h>
#include <stdlib.h>

#define TOL    1E-10
#define WIDTH  10

#define FMIN   5000.0
#define FMAX   9999.0

int main(int argc, char** argv) {
  double A, diff, prev, temp;
  int world_rank, world_size, row, col, neighbors[4], lflag, gflag, i;

  srand(time(NULL));

  diff = 0;
  lflag = 0;
  gflag = 0;
  A = 999999.99;
  // A = FMIN + ((double) rand() / RAND_MAX) * (FMAX - FMIN);

  MPI_Init(NULL, NULL);
  MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
  MPI_Comm_size(MPI_COMM_WORLD, &world_size);

  row = world_rank / WIDTH;
  col = world_rank % WIDTH;

  neighbors[0] = (row == 0) ? -1 : (row - 1) * WIDTH + col; // North
  neighbors[1] = (col == 0) ? -1 : row * WIDTH + (col - 1); // West
  neighbors[2] = (col == (WIDTH - 1)) ? -1 : row * WIDTH + (col + 1); // East
  neighbors[3] = (row == (WIDTH - 1)) ? -1 : (row + 1) * WIDTH + col; // South

  MPI_Barrier(MPI_COMM_WORLD);

  while (1) {
    prev = A;

    for (i = 0; i < 4; ++i) {
      if (neighbors[i] >= 0) {
        MPI_Send(&A, 1, MPI_DOUBLE, neighbors[i], 0, MPI_COMM_WORLD);
      }
    }

    for (i = 0; i < 4; ++i) {
      if (neighbors[i] >= 0) {
        MPI_Recv(&temp, 1, MPI_DOUBLE, neighbors[i], 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        A += temp;
      }
    }

    A *= 0.2;
    diff = fabs(A - prev);
    lflag = (diff < TOL);

    MPI_Barrier(MPI_COMM_WORLD);
    MPI_Allreduce(&lflag, &gflag, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);
    MPI_Barrier(MPI_COMM_WORLD);

    if (gflag > 0) {
      break;
    }
  }

  printf("diff at A[%d, %d] = %.32f\n", row, col, diff);

  MPI_Barrier(MPI_COMM_WORLD);
  MPI_Finalize();

  return EXIT_SUCCESS;
}