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


#define M_PI 3.14159265358979323846264338327
#define MAX(a,b) ((a)>(b)?(a):(b))
#define MIN(a,b) ((a)<(b)?(a):(b))

#define KEISI 0


double u0(double x, double y) {
  double u=0.0;
  switch ( KEISI ) {
  case 1:
    if ( x >= 0.25 && x <=0.5 && y>=0.5 && y<=0.75 ) 
      u=1.0;
    break;
  case 2:
    u=x*(1-x)*y*(1-y);
    break;
  default:
    u=sin(M_PI*x)*sin(M_PI*y);
  }
  return u;
}

//
// exact solution of the heat equation (if you happen to know it
// for the given u0):

double uexact(double x, double y, double t) {
  return exp(-2*M_PI*M_PI*t)*sin(M_PI*x)*sin(M_PI*y);
}

  
void output_solution(int n, double u[n][n]) {
  int i, j;
  double h;
  FILE *fp;
  h=1.0/(n-1);
  fp=fopen("u.dat","w");   // for plotting with gnuplot 
  for ( i=0; i<n; i++ ) {
    for ( j=0; j<n; j++ )
      fprintf(fp,"%lf %lf %lf\n",i*h,j*h,u[i][j]);
    fprintf(fp,"\n");
  }
  fclose(fp);
}

void init_zero2d(int n, double arr[n][n]) {
  for (int i=0; i<n; i++)
    for (int j=0; j<n; j++)
      arr[i][j]=0.0;
}




//
// Solve heat equation 
//      du/dt-uxx-uyy=0 in ]0,1[ x ]0,1[
//       u=0 on boundary
//       u=u0(x,y) at t=0
//
int main(int argc, char *argv[]) {
  double h, dt, dth2, c1, c2, err, umax, T_final=0.01;
  int is, ie, js, je, step, i1, i2, j1, j2;
  int N, NSTEPS;
  int np, me, ier, p, ibuf[4], dims[2];
  int periods[2]={0,0}, reorder=0, debug=0;
  int tag1=555, tag2=556, tag3=557;
  int coords[2], source, dest, nloc;
  double t;
  MPI_Status stat;
  MPI_Comm comm_2d;
  MPI_Request req;
  if ( argc != 3 ) {
    printf("Usage: mpi_heat N NSTEPS\n E.g. N=600, NSTEPS=15000\n");
    exit(1);
  }
  MPI_Init(&argc,&argv);
  MPI_Comm_size(MPI_COMM_WORLD,&np);
  MPI_Comm_rank(MPI_COMM_WORLD,&me);
  N = atoi(argv[1]); NSTEPS = atoi(argv[2]);
  if (me==0) printf("N=%d, NSTEPS=%d\n",N,NSTEPS);
  p = round(sqrt(np));
  if ( (p*p !=np) || ( N%p != 0) ) {
    if ( me==0 ) 
      printf("np should be a perfect square and divides n = %d\n",N);
    MPI_Finalize();
    exit(1);
  }
  h = 1.0/(N-1);
  dt = T_final/NSTEPS;
  if ( dt > 0.25*h*h ) {
    if ( me==0 ) 
      printf("Scheme is not stable. dt=%.3le, should be <= %.3le\n",dt,0.25*h*h);
    MPI_Finalize();
    exit(1);
  }
  dth2=dt/(h*h);
  c1 = 1-4*dth2;
  c2 = dth2;
  //
  if (me==0) 
    t = MPI_Wtime();
  //
  // Computational domain is divided between processors using 2D domain
  // decomposition.
  //
  // Create virtual p x p Cartesian processor grid:
  //
  dims[0]=p; dims[1]=p;
  MPI_Cart_create(MPI_COMM_WORLD,2,dims,periods,reorder,&comm_2d);
  MPI_Cart_coords(comm_2d,me,2,coords);
  nloc = N/p;
  is = coords[0]*nloc;
  ie = (coords[0]+1)*nloc-1;
  js = coords[1]*nloc;
  je = (coords[1]+1)*nloc-1;
  if ( debug ) 
    printf(" me=%d would take care of nodes (%d,%d)...(%d,%d)\n",me,is,js,ie,je);
  double sbuf[nloc], rbuf[nloc];
  double u[nloc+2][nloc+2], unew[nloc][nloc], buf2[nloc][nloc];
  //
  init_zero2d(nloc,unew);
  for ( int i=0; i<nloc+2; i++ )
    for ( int j=0; j<nloc+2; j++)
      u[i][j]=0.0;
  //
  // Initial condition: 
  //
  for ( int i=0; i<nloc; i++ ) 
    for ( int j=0; j<nloc; j++ ) 
      u[1+i][1+j]=u0((is+i)*h,(js+j)*h);
  //
  // Time stepping:
  //
  for ( step=1; step <= NSTEPS; step++ ) {
    MPI_Cart_shift(comm_2d,0,1,&source,&dest);
    if ( source >= 0 ) {
        // exchange data with neighbour above
      for ( int j=0; j<nloc; j++ ) 
	sbuf[j]=u[1][j+1];
      MPI_Isend(sbuf,nloc,MPI_DOUBLE,source,tag1,MPI_COMM_WORLD,&req);
      MPI_Recv(rbuf,nloc,MPI_DOUBLE,source,tag1,MPI_COMM_WORLD,&stat);
      for ( int j=0; j<nloc; j++ )
	u[0][j+1]=rbuf[j];
    }
    if ( dest >= 0 ) {
      // exchange data with neighbour below
      for ( int j=0; j<nloc; j++ )
	sbuf[j]=u[nloc][j+1];
      MPI_Isend(sbuf,nloc,MPI_DOUBLE,dest,tag1,MPI_COMM_WORLD,&req);
      MPI_Recv(rbuf,nloc,MPI_DOUBLE,dest,tag1,MPI_COMM_WORLD,&stat);
      for ( int j=0; j<nloc; j++ ) 
	u[nloc+1][j+1]=rbuf[j];
    }
    MPI_Cart_shift(comm_2d,1,1,&source,&dest) ;
    if ( source >= 0 ) {
      // exchange data with neighbour left
      for ( int i=0; i<nloc; i++ )
	sbuf[i]=u[i+1][1];
      MPI_Isend(sbuf,nloc,MPI_DOUBLE,source,tag1,MPI_COMM_WORLD,&req);
      MPI_Recv(rbuf,nloc,MPI_DOUBLE,source,tag1,MPI_COMM_WORLD,&stat);
      for ( int i=0; i<nloc; i++ )
	u[i+1][0]=rbuf[i];
    }
    if ( dest >= 0 ) {
      // exchange data with neighbour right
      for ( int i=0; i<nloc; i++ )
	sbuf[i]=u[i+1][nloc];
      MPI_Isend(sbuf,nloc,MPI_DOUBLE,dest,tag1,MPI_COMM_WORLD,&req);
      MPI_Recv(rbuf,nloc,MPI_DOUBLE,dest,tag1,MPI_COMM_WORLD,&stat);
      for ( int i=0; i<nloc; i++ )
	u[i+1][nloc+1]=rbuf[i];
    }
    //
    // Update points inside block (i.e. one time step):
    //
    i1 = 1; i2 = nloc; j1=1; j2=nloc;
    if ( is==0   ) i1 = 2;
    if ( ie==N-1 ) i2 = nloc-1;
    if ( js==0   ) j1 = 2;
    if ( je==N-1 ) j2 = nloc-1;
    for ( int i=i1; i<=i2; i++ )
      for ( int j=j1; j<=j2; j++ ) 
	unew[i-1][j-1]=c1*u[i][j]+c2*(u[i+1][j]+u[i-1][j]+u[i][j+1]+u[i][j-1]);
    //
    // Copy unew to u
    //
    for ( int i=1; i<=nloc; i++ )
      for ( int j=1; j<=nloc; j++ )
	u[i][j]=unew[i-1][j-1];
  }
  //
  // Send solution blocks to master for writing to file.
  //
  if ( me==0 ) {
    double uu[N][N];
    // master's own block
    for ( int i=0; i<nloc; i++ )
      for ( int j=0; j<nloc; j++ )
	uu[is+i][js+j]=unew[i][j];   
    // then receive other blocks from other tasks:
    for ( int j=1; j<np; j++ ) {
      MPI_Recv(ibuf,4,MPI_INTEGER,j,tag2,MPI_COMM_WORLD,&stat);
      is = ibuf[0];
      ie = ibuf[1];
      js = ibuf[2];
      je = ibuf[3];
      MPI_Recv(buf2[0],nloc*nloc,MPI_DOUBLE,j,tag3,MPI_COMM_WORLD,&stat);
      for ( int i=0; i<nloc; i++ )
	for ( int j=0; j<nloc; j++ )
	  uu[is+i][js+j]=buf2[i][j];
    }
    t = MPI_Wtime()-t;
    printf("np=%d. Elapsed time = %.1lf seconds.\n",np,t);
    output_solution(N,uu); 
    //
    if ( KEISI==0 ) {
      err = 0.0;
      umax = -1e6;
      for ( int i=0; i<N; i++)
	for ( int j=0; j<N; j++ ) {
	  err = MAX(err,fabs(uu[i][j]-uexact(i*h,j*h,NSTEPS*dt)));
	  umax = MAX(umax,uu[i][j]);
	}
      printf("Umax = %lf, Max. error= %le\n",umax,err);
    }
  }
  else {
    // send block to master
    ibuf[0]=is; ibuf[1]=ie; ibuf[2]=js; ibuf[3]=je;
    MPI_Send(ibuf,4,MPI_INTEGER,0,tag2,MPI_COMM_WORLD);
    for ( int i=0; i<nloc; i++ )
      for ( int j=0; j<nloc; j++ )
	buf2[i][j] = unew[i][j];
    MPI_Send(buf2[0],nloc*nloc,MPI_DOUBLE,0,tag3,MPI_COMM_WORLD);
  }
  MPI_Finalize(); 
}