/****************************************************************************
 * 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/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\n ", argv[0]);
            fprintf(stderr,"\t\twhere N is GRID size, EPSILON is  Tolerance\n"); 
            fprintf(stderr,"\t\texample N = 100, EPSILON = 0.1\n"); 
            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\n");
       return -1;
   }

   if(rank==0) {
   
       printf ( "\n" );
       printf ( "HEATED_PLATE\n" );
       printf ( "  Parallel MPI version using %d processors \n",size );
       printf ( "  A program to solve for the steady state temperature distribution\n" );
       printf ( "  over a rectangular plate.\n" );
       printf ( "\n" );
       printf ( "  Spatial grid of %d by %d points.\n", N, N );
       printf ( "  Each processor will use grid of %d +2 by %d points.\n", M-2, N );
       printf ( "  The iteration will end until tolerance <= %f\n\n",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");


   /* 
     *   iterate until the  new solution unew differs from the old solution u
     *     by no more than EPSILON.
     *     */

   float diff=1.0;
   int iter=0;

   while(diff> EPSILON)  {

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

        if(rank==0)
            if(iter%ITER_PRINT==0) 
                printf("Processor #%d, iteration %d, epsilon = %f\n ", rank,iter,diff);
        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;
    float globaldiff; 
    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 && rank< size-1)
    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> 0 && 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] );
            }
        }

    }


  /**
   *  COMPUTE GLOBAL CONVERGENCE
   *
   * */

    MPI_Allreduce(&diff, &globaldiff , 1, MPI_FLOAT, MPI_MAX, MPI_COMM_WORLD); 


   /**
    *  COPY OLD DATA 
    * */

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


 return globaldiff;   
}

/*****************************************************************************
 *  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 = 1; ix < nx-1; ix++){ 
        u[ix*ny]=100.0; 

    }

    //boundary right
    for (ix = 1; ix < nx-1; ix++){ 
        u[ix*ny+ (ny-1)]=100.0; 

    }

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

    //boundary top
     for (iy = 0; iy < ny; iy++){ 
        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 ( '\n', 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);

}