/* diffusion_par2.c: parallel 2d-diffusion simulation, with
 * "optimized" write_frame.
 * Author: Andrew R. Siegel
 * Modified by Stephen F. Siegel <siegel@cis.udel.edu>, August 2007
 * for EuroPVM/MPI 2007 tutorial
 */

#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include "mpi.h"
#define SQUARE(x) ((x)*(x))

/* Constants: the following should be defined at compilation:
 *
 *       D = diffusion constant
 *      dt = time step
 *      dx = distance between two lattice points
 *       r = radius of initial circle in center
 *  nsteps = number of time steps
 *      nx = number of points in x direction
 *      ny = number of points in y direction
 * nprocsx = number of procs in the x direction
 * nprocsy = number of procs in the y direction
 *     nxl = horizontal extent of one process; divides nprocsx
 *     nyl = vertical extent of one processl; divides nprocsy
 */

/* Global variables */

int nprocs;    /* number of processes */
int rank;      /* the rank of this process */
int xcoord;    /* the horizontal coordinate of this proc */
int ycoord;    /* the vertical coordinate of this proc */
int up;        /* rank of upper neighbor on torus */
int down;      /* rank of lower neighbor on torus */
int left;      /* rank of left neighbor on torus */
int right;     /* rank of right neighbor on torus */
double u[nxl+2][nyl+2];  /* values of the discretized function */

/* init: initializes global variables. */
void init() { 
  int i, j; 
  
  MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  assert(nprocs == nprocsx*nprocsy);
  assert(nx == nxl*nprocsx);
  assert(ny == nyl*nprocsy);
  xcoord = rank % nprocsx;
  ycoord = rank / nprocsx;
  up = xcoord + nprocsx*((ycoord + 1)%nprocsy);
  down = xcoord + nprocsx*((ycoord + nprocsy - 1)%nprocsy);
  left = (xcoord + nprocsx - 1)%nprocsx + nprocsx*ycoord;
  right = (xcoord + 1)%nprocsx + nprocsx*ycoord;
  for (i = 1; i <= nxl; i++)
    for (j = 1; j <= nyl; j++)
      u[i][j] = 
	(SQUARE(dx*(nxl*xcoord+i-1) - 0.5*dx*(nx - 1)) +
	 SQUARE(dx*(nyl*ycoord+j-1) - 0.5*dx*(ny - 1))
	 < r*r ? 100 : 0);
}

/* write_frame: writes current values of u to file */
void write_frame(int time) {
  int i, j;
  double buf[nxl];

  if (rank != 0) {
    for (j = 1; j <= nyl; j++) {
      for (i = 1; i <= nxl; i++) buf[i-1] = u[i][j];
      MPI_Send(buf, nxl, MPI_DOUBLE, 0, j-1, MPI_COMM_WORLD);
    }
  } else {
    MPI_Status status;
    char filename[50];
    char command[50];
    FILE *file;
    int linelen = nprocsx*nxl*8+1; /* length of one line, incl. \n */
    int k, m, n, from;

    sprintf(filename, "./parout/out_%d", time);
    file = fopen(filename, "w");
    assert(file);
    for (n = 0; n < nprocsy; n++) {
      for (j = 1; j <= nyl; j++) {
        for (m = 0; m < nprocsx; m++) {
          from = n*nprocsx + m;
          if (from != 0)
            for (i = 1; i <= nxl; i++) fprintf(file, "        ");
	  else
            for (i = 1; i <= nxl; i++) fprintf(file, "%8.2f", u[i][j]);
	}
	fprintf(file, "\n");
      }
    }
    for (k = 0; k < nyl*(nprocs - 1); k++) {
      int procx, procy, row;
      
      MPI_Recv(buf, nxl, MPI_DOUBLE, MPI_ANY_SOURCE, MPI_ANY_TAG,
               MPI_COMM_WORLD, &status);
      procx = status.MPI_SOURCE % nprocsx;
      procy = status.MPI_SOURCE / nprocsx;
      row = procy*nyl + status.MPI_TAG;
      fseek(file, row*linelen + procx*nxl*8, SEEK_SET);
      for (i = 1; i <= nxl; i++) fprintf(file, "%8.2f", buf[i-1]);
    }
    fclose(file);
    sprintf(command, "cat %s", filename);
    system(command);
  }
  sleep(3);
}

/* exchange_ghost_cells: updates ghost cells using MPI communication */
void exchange_ghost_cells() {
  int i, j;
  double sbufx[nxl], rbufx[nxl], sbufy[nyl], rbufy[nyl];

  for (i = 1; i <= nxl; ++i) sbufx[i-1] = u[i][1];
  MPI_Sendrecv(sbufx, nxl, MPI_DOUBLE, down, 0, rbufx, nxl,
               MPI_DOUBLE, up, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
  for (i = 1; i <= nxl; ++i) u[i][nyl+1] = rbufx[i-1];
  for (i = 1; i <= nxl; ++i) sbufx[i-1] = u[i][nyl];
  MPI_Sendrecv(sbufx, nxl, MPI_DOUBLE, up, 0, rbufx, nxl,
               MPI_DOUBLE, down, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
  for (i = 1; i <= nxl; ++i) u[i][0] = rbufx[i-1];
  for (j = 1; j <= nyl; ++j) sbufy[j-1] = u[1][j];  
  MPI_Sendrecv(sbufy, nyl, MPI_DOUBLE, left, 0, rbufy, nyl,
               MPI_DOUBLE, right, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
  for (j = 1; j <= nyl; ++j) u[nxl+1][j] = rbufy[j-1];
  for (j = 1; j <= nyl; ++j) sbufy[j-1] = u[nxl][j];
  MPI_Sendrecv(sbufy, nyl, MPI_DOUBLE, right, 0, rbufy, nyl,
               MPI_DOUBLE, left, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
  for (j = 1; j <= nyl; ++j) u[0][j] = rbufy[j-1];
}

/* update: updates u.  Uses ghost cells.  Purely local operation. */
void update() {
  int i, j;
  double k = D*dt/(dx*dx);
  double u_new[nxl+2][nyl+2];

  for (i = 1; i <= nxl; i++)
    for (j = 1; j <= nyl; j++)
      u_new[i][j] = u[i][j]
        + k*(u[i+1][j]+u[i-1][j]+u[i][j+1]+u[i][j-1]-4*u[i][j]);
  for (i = 1; i <= nxl; i++)
    for (j = 1; j <= nyl; j++) 
      u[i][j] = u_new[i][j];
}

/* main: executes simulation, creates one output file for each time
 * step */
int main(int argc,char *argv[]) {
  int iter;

  MPI_Init(&argc, &argv);
  init();
  write_frame(0);
  for (iter = 1; iter <= nsteps; iter++) {
    exchange_ghost_cells();
    update();
    write_frame(iter);
  }
  MPI_Finalize();
  return 0;
}
