#include <stdio.h>
#include <mpi.h>
#include <math.h>
#define EPSILON 0.00001
#define N 100
#define time_steps 100

int main (int argc, char *argv[])
 {
  int i,j;
  int myid, numprocs;
  int from,to;
  int step;
  double time;
  double eps, enew, global_eps;
  double time_max = 3.0;
  double alpha = 0.06;
  double dx = 1.0/N;
  double dy = 1.0/time_steps;
  double dt = time_max/time_steps;
  double dxinv = 1.0/dx;
  double dyinv = 1.0/dy;
  double dtinv = 1.0/dt;
  double divinv = 1.0/(dtinv + 2 * alpha * (dxinv * dxinv + dyinv * dyinv));
  double **t,*tS;
  double **told,*toldS;
  double **tstep,*tstepS;
  double *sendR1,*sendR2;
  double *recvR1,*recvR2;
  double minval, maxval;
  long clock(),cputime;
  char fname[40];
  FILE *out;
  MPI_Status status;

  MPI_Init(&argc, &argv);
  MPI_Comm_size(MPI_COMM_WORLD,&numprocs);
  MPI_Comm_rank(MPI_COMM_WORLD,&myid); 
  
  if (((N % numprocs) == 0) && (numprocs > 1))
   {
    if (myid == 0) 
     {
      clock();
      tstepS = (double *)  malloc(N * N * sizeof(double));
      tstep =  (double **) malloc(N * sizeof(double *));

      for (i=0; i<N; i++)
       tstep[i] = &tstepS[i*N];
     }
    //allocate enough memory for each fraction of the mesh
    tS =    (double *)  malloc(((N/numprocs)+2) * N * sizeof(double));
    toldS = (double *)  malloc(((N/numprocs)+2) * N * sizeof(double));
    t =     (double **) malloc(((N/numprocs)+2) * sizeof(double *));
    told =  (double **) malloc(((N/numprocs)+2) * sizeof(double *));

    printf("rows = %d\n", (N/numprocs)+2);
    for (i=0; i<(N/numprocs)+2; i++)
     {
      t[i] = &tS[i*N];
      told[i] = &toldS[i*N];
     }
    // set initial boundary conditions
    for (i=0; i<(N/numprocs)+2; i++)
     for (j=0; j<N; j++)
      told[i][j] = 0.0; 
    // for all time steps
    for (step = 1; step <= time_steps; step++)
     {
      time = step * (time_max/time_steps);
      // reset top boundary condition each timestep
      if (myid == 0) 
       for (j=0; j<N; j++)
         told[0][j] = 2.0 * sin(time);
      do
       {
        // exchange the rows you need to share with other processes
        if (myid == 0)
         {
          //allocate memory for message arrays
          sendR1 = (double *) malloc(N * sizeof(double));
          recvR1 = (double *) malloc(N * sizeof(double));
          //send to my last computed row to myid+1
          for(j=0; j<N; j++) sendR1[j]=told[(N/numprocs)][j];
          MPI_Send(sendR1,N,MPI_DOUBLE,myid+1,0,MPI_COMM_WORLD);
          //receive my last row from myid+1
          MPI_Recv(recvR1,N,MPI_DOUBLE,myid+1,0,MPI_COMM_WORLD,&status);
          for(j=0; j<N; j++) told[(N/numprocs)+1][j]=recvR1[j];      
          free(sendR1);
          free(recvR1);   
         }
        else if ((myid > 0) && (myid < (numprocs-1)))
         {
          //allocate memory for message arrays
          sendR1 = (double *) malloc(N * sizeof(double));
          recvR1 = (double *) malloc(N * sizeof(double));
          sendR2 = (double *) malloc(N * sizeof(double));
          recvR2 = (double *) malloc(N * sizeof(double));
          //send my first computed row to myid-1
          for(j=0; j<N; j++) sendR1[j]=told[1][j];
          MPI_Send(sendR1,N,MPI_DOUBLE,myid-1,0,MPI_COMM_WORLD);
          //send my last computed row to myid+1
          for(j=0; j<N; j++) sendR2[j]=told[(N/numprocs)][j];
          MPI_Send(sendR2,N,MPI_DOUBLE,myid+1,0,MPI_COMM_WORLD);
          //receive my first row from myid-1
          MPI_Recv(recvR1,N,MPI_DOUBLE,myid-1,0,MPI_COMM_WORLD,&status);
          for(j=0; j<N; j++) told[0][j]=recvR1[j]; 
          //receive my last row from myid+1
          MPI_Recv(recvR2,N,MPI_DOUBLE,myid+1,0,MPI_COMM_WORLD,&status);
          for(j=0; j<N; j++) told[(N/numprocs)+1][j]=recvR2[j]; 
          free(sendR1);
          free(sendR2);
          free(recvR1);
          free(recvR2); 
         }
        if (myid == (numprocs-1))
         {
          //allocate memory for message arrays
          sendR1 = (double *) malloc(N * sizeof(double));
          recvR1 = (double *) malloc(N * sizeof(double));
          //send my first computed row myid-1
          for(j=0; j<N; j++) sendR1[j]=told[1][j];
          MPI_Send(sendR1,N,MPI_DOUBLE,myid-1,0,MPI_COMM_WORLD);
          //receive my first row myid-1
          MPI_Recv(recvR1,N,MPI_DOUBLE,myid-1,0,MPI_COMM_WORLD,&status);
          for(j=0; j<N; j++) told[0][j]=recvR1[j];
          free(sendR1);
          free(recvR1); 
         }
        eps = 0.0;
        for (i=1; i<=(N/numprocs); i++)
          for (j=1; j<(N-1); j++)
          t[i][j]=((told[i][j+1]+told[i][j-1])*alpha*dyinv*dyinv+(told[i+1][j]+told[i-1][j])*alpha*dxinv*dxinv+(told[i][j]*dtinv))*divinv;
        for (i=1; i<=(N/numprocs); i++)
         {
          for (j=1; j<(N-1); j++)
           {
            enew = fabs(t[i][j] - told[i][j]);
            if (enew > eps) { eps = enew; }
           }
         }
        for (i=0; i<((N/numprocs)+2); i++)
         for (j=0; j<N; j++)
          told[i][j] = t[i][j];
        MPI_Allreduce(&eps,&global_eps,1,MPI_DOUBLE,MPI_MAX,MPI_COMM_WORLD);
       }
      while(global_eps > EPSILON);
      // gather all the subsets of the grid
      MPI_Gather(told[1],N*(N/numprocs),MPI_DOUBLE,tstepS,N*(N/numprocs),MPI_DOUBLE,0,MPI_COMM_WORLD);
      // Produce a raster image for this time step
      if (myid == 0)
       {
        minval = 0.0;
        maxval = 0.0;
        for (i=0; i<N; i++)
         {
          for (j=0; j<N; j++)
           {   
            if (tstep[i][j] < minval) { minval = tstep[i][j]; }
            if (tstep[i][j] > maxval) { maxval = tstep[i][j]; }
           }
         }
        sprintf(fname,"Output\\heat%03d.raw",step); 
        out = fopen(fname,"w+b");
        for (i=0; i<N; i++)
          for (j=0; j<N; j++)
            fprintf(out,"%c",(int)(((tstep[i][j]-minval)*255.0)/(maxval - minval)));
        fclose(out);
        printf("Time step: %d\r",step);
       }
      MPI_Barrier(MPI_COMM_WORLD);
     } // for all time steps 
    free(t);
    free(tS);
    free(told);
    free(toldS);
    if (myid == 0)
     {
      free(tstep);
      free(tstepS);
      cputime = clock();
      printf("%d time steps in %.2f seconds\n",step-1,cputime/1.0e+3);
     }
    } // if ((N % numprocs) == 0)) 
   else
    if (myid == 0)
     {
      if (numprocs < 2)
       printf("ERROR: Must have at least 2 processes\n");
      if ((N % numprocs) != 0)
       printf("ERROR: Grid must be evenly divisible by the number of processes\n");
     }
   MPI_Finalize();
 }