#pragma TASS input 
int NX_BOUND;
#pragma TASS input 
int NY_BOUND;
#pragma TASS input 
int TIME_BOUND;
#pragma TASS input {nx > 2 && nx < NX_BOUND} 
int nx;  // number of x coordinates (including boundary)
#pragma TASS input {ny > 2 && ny < NY_BOUND} 
int ny;  // number of rows including boundary
#pragma TASS input {epsilon > 0 && epsilon < 1} 
double epsilon; // total error tolerance
#pragma TASS input 
double initialValues[ny][nx]; // initial values
#pragma TASS output 
int t;
#pragma TASS output 
double out[ny][nx];

int ALLREDUCE_TAG = -1;
int upper;
int lower;
int firstRow;
int nyl;
int start;
int stop;
int new = 0;   // flag to select which dimension of grid is the new grid
/* layout of local grid:
 *   row\col | 012=nx-1
 *  ---------+----
 *   nyl+1=2 | 678
 *         1 | 345
 *         0 | 012
 * grid_data = (double*)malloc(nx*(nyl+2)*sizeof(double));
 * grid = (double**)malloc((nyl+2)*sizeof(double*));
 * i-th row, j-th column:
 * grid[i][j] = grid_data[i*nx+j]
 * grid[i] = grid_data + i*nx
 *
 * Now there are actually two grids (one old, one new), so
 * need arrays of length 2 for both of above:
 *
 * (double*)[2] grid_data;
 * (double**)[2] grid;
 */
double** grid[2];
double* grid_data[2];
double* buf;

double square(double x){
  return x * x;
}

void MPI_Send(double* buffer, int count, int destination, int tag) {
  double tempBuffer[count];
  int i;
  
  if (destination >= 0) {
    for (i=0; i<count; i++) tempBuffer[i] = buffer[i];
    send(tempBuffer, destination, tag);
  }
}

void MPI_Recv(double* buffer, int count, int source, int tag) {
  double tempBuffer[count];
  int i;
  
  if (source >= 0) {
    recv(tempBuffer, source, tag);
    for (i=0; i<count; i++) buffer[i] = tempBuffer[i];
  }
}

/* This implementation is not quite right because it always sends first then receives.
 * However it is not possible to implement a true MPI_Sendrecv until a sendrecv location
 * is added to the model package.  This solution will still work without deadlocking in this
 * application because the process on the extreme end of the grid will only have to receive.
 */
void MPI_Sendrecv(double* sendbuf, int sendcount, int dest, int sendtag,
                  double* recvbuf, int recvcount, int source, int recvtag) {
  double tempSendBuffer[sendcount];
  double tempRecvBuffer[recvcount];
  int i;
  
  if (dest >= 0) {
    for (i=0; i<sendcount; i++) tempSendBuffer[i]=sendbuf[i];
    send(tempSendBuffer, dest, sendtag); 
  }
  if (source >= 0) {
    recv(tempRecvBuffer, source, recvtag);
    for (i=0; i<recvcount; i++) recvbuf[i]=tempRecvBuffer[i];
  }
}

/* This is an implementation of allreduce for the addition operator.*/
void MPI_Allreduce(double* sendbuf, double* recvbuf, int count) {
  int i;
  int j;
  double tempBuffer[count];
  
  if (PID == 0) {
    for (j=0; j<count; j++) tempBuffer[j]=sendbuf[j];
  	for (i=1; i<NPROCS; i++) {
  	  MPI_Recv(recvbuf, count, i, ALLREDUCE_TAG);
  	  for(j=0; j<count; j++) tempBuffer[j] += recvbuf[j];
  	}
  	for (j=0; j<count; j++) recvbuf[j]=tempBuffer[j];
  	for (i=1; i<NPROCS; i++) {
  	  MPI_Send(recvbuf, count, i, ALLREDUCE_TAG);
  	}
  } else {
    MPI_Send(sendbuf, count, 0, ALLREDUCE_TAG);
    MPI_Recv(recvbuf, count, 0, ALLREDUCE_TAG);
  }
}

int owner(int row) { return (NPROCS*(row+1)-1)/ny; }

int firstForProc(int rank) { return (rank*ny)/NPROCS; }

int countForProc(int rank) {
  int a;
  int b;
  
  a = firstForProc(rank+1);
  b = firstForProc(rank);
  return a-b;
}

void init() {
  int tmp;
  int row;
  int col;
  int s;
  int i;
  
  for (i=0; i<ny; i++);
  firstRow = firstForProc(PID);
  nyl = countForProc(PID);
  for(s=0; s<2; s++) {
    grid_data[s] = malloc((nx*(nyl+2)) * sizeof(double));
    grid[s] = malloc((nyl+2) * sizeof(double *));
    for (i=0; i<nyl+2; i++) grid[s][i]=grid_data[s]+i*nx;
  }
  if (firstRow == 0 || nyl == 0)
    lower = -1;
  else
    lower = owner(firstRow-1);
  if (firstRow+nyl >= ny || nyl == 0)
    upper = -1;
  else
    upper = owner(firstRow+nyl);
  tmp = owner(0);
  if (PID == tmp) start = 2; else start = 1;
  tmp = owner(ny-1);
  if (PID == tmp) stop = nyl-1; else stop=nyl;
  for (row=1; row<=nyl; row++) {
    for (col=0; col<nx; col++) {
      grid[0][row][col] = initialValues[firstRow+row-1][col];
      grid[1][row][col] = grid[0][row][col];
    }
  }
  buf = malloc(nx * sizeof(double));
}

void cleanup() {
  int s;
  
  free(buf);
  for (s=0; s<2; s++) {
    free(grid[s]);
    free(grid_data[s]);
  }
}

void write_result(int time, double** data) {
  int row;
  int col;
  int proc;
  
  t = time;
  if (PID != 0) {
    for (row=nyl; row>=1; row--)
      MPI_Send(grid[new][row], nx, 0, 0);
  }
  else {
    for (row=ny-1; row>=0; row--) {
      proc = owner(row);
      if (proc != 0) {
        MPI_Recv(buf, nx, proc, 0);
        for (col=0; col<nx; col++) out[row][col]=buf[col];
      } else {
        for (col=0; col<nx; col++) out[row][col]=grid[new][row+1][col];
      }
    }
  }
}

void exchange() {
  MPI_Sendrecv(grid[new][1], nx, lower, 0, grid[new][nyl+1], nx, upper, 0);
  MPI_Sendrecv(grid[new][nyl], nx, upper, 0, grid[new][0], nx, lower, 0);
}

/* returns local error */
double local_update() {
  int row;
  int col;
  int old=1-new;
  double error = 0.0;
  double result;
  
  for (row=start; row<=stop; row++) {
    for (col=1; col<nx-1; col++) {
      grid[new][row][col] = (grid[old][row-1][col]+grid[old][row+1][col]+
          grid[old][row][col-1]+grid[old][row][col+1])/4.0;
      result = square(grid[old][row][col] - grid[new][row][col]);
      error += result;
    }
  }
  return error;
}

void main() {
  double global_error = epsilon;
  double error;
  int time = 0;
  
  init();
  while (global_error >= epsilon && time < TIME_BOUND) {
    exchange();
    new = 1-new;
    error = local_update();
    MPI_Allreduce(&error, &global_error, 1);
    time++;
  }
  write_result(time, grid[new]);
}
