/****************************************************************************
   * DESCRIPTION:  
   *   Serial HEAT2D Example - C Version
   *   This example is based on a simplified 
   *   two-dimensional heat equation domain decomposition.  The initial 
   *   temperature is computed to be high in the middle of the domain and 
   *   zero at the boundaries.  The boundaries are held at zero throughout 
   *   the simulation.  During the time-stepping, an array containing two 
   *   domains is used; these domains alternate between old data and new data.
   *
   *  The physical region, and the boundary conditions, are suggested
      by this diagram;
  
                     u = 0
               +------------------+
               |                  |
      u = 100  |                  | u = 100
               |                  |
               |                  |
               |                  |
               |                  |
               +------------------+
                     u = 100
  
  Interrior point :
    u[Central] = (1/4) * ( u[North] + u[South] + u[East] + u[West] )
  
  
  PARALLEL MPI VERSION :
  
             +-------------------+
             |                   | P0   m=(n-2)/P +2
             +-------------------+
             |                   | P1
             +-------------------+
         n   |                   | ..
             +-------------------+
             |                   | Pq
             +-------------------+
                   
             <-------- n -------->  
              <-------n-2 ------>
  
  
   ****************************************************************************/
  #include <stdio.h>
  #include <stdlib.h>
  #include <math.h>

  #include <mpi.h>

  #define NN 50
  #define MM 50  
  
  #define RING 100
  #define ITER_PRINT 100
  #define PRINT_DATA 1
  #define MAX_ITER 1000
  #define _EPSILON 0.01
  
  
  float update(int rank,int size, int nx,int ny, float *u, float *unew);
  void inidat(int rank, int size, int nx, int ny, float *u, float *unew); 
  void prtdat(int rank, int size, int nx, int ny, float *u,const char *fnam);
  
  
  
  
  int main(int argc, char *argv[])
  {
      int N=NN,M=MM;
  
      int rank,size;
  
      float EPSILON=_EPSILON;

      /* INITIALIZE MPI */

      MPI_Init(&argc, &argv);

  
      /* GET THE PROCESSOR ID AND NUMBER OF PROCESSORS */

      MPI_Comm_rank(MPI_COMM_WORLD, &rank);
      MPI_Comm_size(MPI_COMM_WORLD, &size);
  
      //Only Rank 0 read application parameters
      if(rank==0) {
  
          if(argc !=3)
          {
              fprintf(stderr,"usage %s N EPSILON
   ", argv[0]);
              fprintf(stderr,"\t\twhere N is GRID size, EPSILON is  Tolerance
  "); 
              fprintf(stderr,"\t\texample N = 100, EPSILON = 0.1
  "); 
              return -1;
          } 

          N = M = atoi(argv[1]);
          EPSILON = atof(argv[2]);
  
          if(N % size!=0)
          {
              fprintf(stderr,"Grid Size MUST be divisible by the number of processors !");
              return -1;
          }
  
      }

      //Wait for rank 0 , all process start here 
      MPI_Barrier(MPI_COMM_WORLD);
  
      //Exchange N
      MPI_Bcast(&N , 1, MPI_FLOAT, 0 , MPI_COMM_WORLD);
      //Exchange EPSILON  
      MPI_Bcast(&EPSILON , 1, MPI_FLOAT, 0 , MPI_COMM_WORLD);
  
      //local size
      M = (N-2) / size + 2 ;
  
      float *u     = (float *)malloc(N * M * sizeof(float));
      float *unew  = (float *)malloc(N * M * sizeof(float));
  
      if(u==0 || unew ==0)
      {
          perror("Can't allocated data
  ");
          return -1;
      }
  
      if(rank==0) {
  
          printf ( "
  " );
          printf ( "HEATED_PLATE
  " );
          printf ( "  Parallel MPI version using %d processors 
  ",size );
          printf ( "  A program to solve for the steady state temperature distribution
  " );
          printf ( "  over a rectangular plate.
  " );
          printf ( "
  " );
          printf ( "  Spatial grid of %d by %d points.
  ", N, N );
          printf ( "  Each processor will use grid of %d  by %d points.
  ", M, N );
          printf ( "  The iteration will end until tolerance <= %f
  
  ",EPSILON);
  
      }
  
      /* Initialize grid and create input file 
       * each process initialize its part
       * */
  
      inidat(rank,size,M,N,u,unew);
  
      prtdat(rank,size,M,N,u, "initial.dat");
  
  
      float diff, globaldiff=1.0;
      int iter=0;
  
      /*
       *   iterate (JACOBI ITERATION) until the new solution unew differs from the old solution u

       *     by no more than EPSILON.

       **/
  
      while(globaldiff> EPSILON)  {

          diff= update(rank,size,M,N, u, unew);

          /**
           *   COMPUTE GLOBAL CONVERGENCE
           * */
  
          MPI_Allreduce(&diff, &globaldiff , 1, MPI_FLOAT, MPI_MAX, MPI_COMM_WORLD); 

  
          if(rank==0)
              if(iter%ITER_PRINT==0) 

                  printf("Processor #%d, iteration %d, epsilon = %f
   ", rank,iter,globaldiff);

          iter++;
      }
  
      prtdat(rank,size,M,N, u, "final.dat");
      free(u);
      free(unew);
      MPI_Finalize();
  }
  
  
  
  /****************************************************************************
   *  subroutine update
   ****************************************************************************/
  float update(int rank, int size, int nx,int ny, float *u, float *unew){
      int ix, iy;
      float  diff=0.0;

      MPI_Status status;

      /*
       * EXCHANGE GHOST CELL
       */
      if (rank > 0 && rank< size-1)
      {
          MPI_Sendrecv(&u[ny*(nx-2)], ny, MPI_FLOAT, rank+1, 0,
                  &u[ny*0],     ny, MPI_FLOAT, rank-1, 0, MPI_COMM_WORLD, &status);
          MPI_Sendrecv(&u[ny*1],     ny, MPI_FLOAT, rank-1, 1,
                  &u[ny*(nx-1)], ny, MPI_FLOAT, rank+1, 1, MPI_COMM_WORLD, &status);
      }

      else if (rank == 0)
          MPI_Sendrecv(&u[ny*(nx-2)], ny, MPI_FLOAT, rank+1, 0,
                  &u[ny*(nx-1)], ny, MPI_FLOAT, rank+1, 1, MPI_COMM_WORLD, &status);
      else if (rank == size-1)
          MPI_Sendrecv(&u[ny*1],     ny, MPI_FLOAT, rank-1, 1,
                  &u[ny*0],     ny, MPI_FLOAT, rank-1, 0, MPI_COMM_WORLD, &status);

  
  
  
      /**
       * PERFORM LOCAL COMPUTATION
       * */
  
      for (ix = 1; ix < nx-1; ix++) {
          for (iy = 1; iy < ny-1; iy++) {
              unew[ix*ny+iy] = (u[(ix+1)*ny+iy] +  u[(ix-1)*ny+iy] + u[ix*ny+iy+1] +  u[ix*ny+iy-1] )/4.0
                  ;
              if (diff < fabs (unew[ix*ny+iy] - u[ix*ny+iy] ))
              {
                  diff = fabs ( unew[ix*ny+iy] - u[ix*ny+iy] );
              }
          }
  
      }

      for (ix = 1; ix < nx-1; ix++) {
          for (iy = 1; iy < ny-1; iy++) {
              u[ix*ny+iy] = unew[ix*ny+iy]; 
          }
      }  

      return diff;   

  }
  
  /*****************************************************************************
   *  Initialize Data
   *****************************************************************************/
  
  void inidat(int rank, int size,int nx, int ny, float *u, float *unew) 
  {
      int ix, iy;
  
      /*
       *Set boundary data and interrior values
       * */
  
  
      // interior points
      for (ix = 1; ix < nx-1; ix++) 
          for (iy = 1; iy < ny-1; iy++) { 
              u[ix*ny+iy]=0.0; 
          }

      //boundary left
      for (ix = 0; ix < nx; ix++){ 

          u[ix*ny]=100.0; 
  
      }
  
      //boundary right

      for (ix = 0; ix < nx; ix++){ 

          u[ix*ny+ (ny-1)]=100.0; 
  
      }
  
      //boundary down

      for (iy = 1; iy < ny-1; iy++){ 
  
          if(rank==size-1) {
              u[(nx-1)*(ny)+iy]=100.0; 
              u[0]=100.0;
  
          }else
          {
              u[(nx-1)*(ny)+iy]=0.0;
          }

      }
  
      //boundary top

      for (iy = 1; iy < ny-1; iy++){ 
         if(rank==0)
          u[iy]=100.0;
        else
           u[iy]=0.0; 

  
      }
  
  }
  
  
  /***************************************************************************
   * Print Data to files
   **************************************************************************/
  
  void print2file(int rank, int nx, int ny, float *u,const char *fname)
  {
      int ix, iy;
      FILE *fp;

      char str[255];

      sprintf(str, "%d_%s", rank,fname); 
  
      fp = fopen(str, "w");

      for (ix = 0 ; ix < nx; ix++) {
          for (iy =0; iy < ny; iy++) {
  
              fprintf(fp, "%6.2f ", u[ix*ny+iy]);
          }
          fputc ( '
  ', fp);
      }
  
      fclose(fp);
  
  }
  
  
  
  void prtdat(int rank, int size,int nx, int ny, float *u,const char *fname)
  {
  
  
      if(ITER_PRINT==0)return;
  
  
      /**
       * USE TOKEN RING to write to unique file
       * 0 will rite first and then 1, 2, ... size-1
       *
       * * */

      print2file(rank,nx,ny,u,fname);

  
  }