/* poisson2d.c 
 *
 * Solve 2D Poisson equation u_{xx} + u_{yy} = f(x,y) using
 * point-Jacobi iteration method with Dirichlet boundary conditions.
 * Currently we assume f=0 for simplicity.
 *
 * Note: This code is written to show basic MPI usage. It is *not*
 * suitable for production runs.
 *
 */

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

void init_poisson(int, int, double *, double);
void swap_ptr(double **, double **);
void setbc_poisson(int, int, double *, int, int);
void poisson_ghost_exch(int, int, double *, int, int);
void print_array(int, int, double *, int, int);
void print_one(int, int, double, int);
void write_array(int, int, double *, int, int, double, double);
double poisson_jacobi(int, int, double *, double *, int, int);

       /* problem size etc. */

#define M1          40         /* global size in x */
#define M2          40         /* global size in y */
#define ITMAX       5000       /* max iteration no. */

int main(int argc, char *argv[])
{
     double *phi[2];
     int n1, n2, rank, nproc, it, i;
     double residual,t1,t2;

     /* MPI initialization */
     MPI_Init(&argc, &argv);
     MPI_Comm_rank(MPI_COMM_WORLD, &rank);
     MPI_Comm_size(MPI_COMM_WORLD, &nproc);

     /* local problem size */
     n1 = M1;
     n2 = M2/nproc;

     /* allocate some arrays */
     phi[0] = (double *) malloc(sizeof(double)*n1*(n2+2));  
     phi[1] = (double *) malloc(sizeof(double)*n1*(n2+2));

       /* set start-up condition for iterations */

     init_poisson(n1,n2+2,phi[0],0.0);
     setbc_poisson(n1,n2,phi[0],rank,nproc);

     /* main loop starts here */

     for (it=1; it<ITMAX; it++) {
          residual = poisson_jacobi(n1,n2,phi[0],phi[1], rank, nproc);
          setbc_poisson(n1,n2,phi[1],rank,nproc);
          print_one(it,20,residual,rank); 
          poisson_ghost_exch(n1,n2,phi[1],rank,nproc);
          swap_ptr(&phi[0],&phi[1]);
     }

     /* write result to disk */
     write_array(n1,n2,phi[1],rank,nproc,1.0/(M1-1), 1.0/(M2-1));

     /* clean up */

     for(i=0; i<2; i++) free(phi[i]);
     MPI_Finalize();
     return 0;
}

       /* exchange ghost point values among processors */
void poisson_ghost_exch(int n1, int n2, double *phi, int rank, int nproc)
{
     int i, j, ij, south, north;
     double *ps, *pn, *gs, *gn;
     MPI_Status status;

     pn = phi + n1*n2;
     gs = phi;
     ps = phi + n1;
     gn = phi + n1*(n2+1);

     south = rank-1;
     north = rank+1;
     if (south < 0) south = MPI_PROC_NULL;
     if (north > nproc-1) north = MPI_PROC_NULL;

     /* send to north */
     MPI_Sendrecv(pn,n1,MPI_DOUBLE,north,0,
                  gs,n1,MPI_DOUBLE,south,0,MPI_COMM_WORLD,&status);
     /* send to south */
     MPI_Sendrecv(ps,n1,MPI_DOUBLE,south,0,
                  gn,n1,MPI_DOUBLE,north,0,MPI_COMM_WORLD,&status);
}

       /* set boundary values along boundaries */
void setbc_poisson(int n1, int n2, double *p, int rank, int nproc)
{
     int i,j,offset;
     /* west */
     i=0;
     for (j=0; j<n2; j++) 
          p[i+n1*j] = 0.0;

     /* east */
     i = n1-1;
     for (j=0; j<n2; j++) 
          p[i+n1*j] = 0.0;

     /* south */
     if (rank == 0) {
          offset = n1;
          for (i=0; i<n1; i++) p[i+offset] = 0.0;
     }

     /* north */
     if (rank == nproc-1) {
          offset = n1*n2;
          for (i=0; i<n1; i++) 
               p[i+offset] = 1.0;
     }
}

       /* Point-Jacobi iteration */
double poisson_jacobi(int n1, int n2, double *f1, double *f2, 
                      int rank, int nproc)
{
     int i,j,ij;
     int istart = 1, iend = n1-2;
     int jstart = 1, jend = n2;
     int east = 1, west = -1, north = n1, south = -n1;
     double res, tmp, res0=0.0;
     double hinv = 0.25;

     if (rank == 0) jstart = 2;
     if (rank == nproc-1) jend = n2-1;

     for (j=jstart; j<jend+1; j++) {
          for (i=istart; i<iend+1; i++) {
               ij = i+n1*j;
               f2[ij] = hinv*(f1[ij+east] + f1[ij+west] 
                              + f1[ij+south] + f1[ij+north]);
          }
     }

     res=0.0;
     for (j=jstart; j<jend+1; j++) {
          for (i=istart; i<iend+1; i++) {
               ij = i+n1*j;
               tmp = f2[ij] - hinv*(f2[ij+east] + f2[ij+west] 
                                    + f2[ij+south] + f2[ij+north]);
               res += tmp*tmp;
          }
     }

     res = sqrt(res/(n1*n2));
     MPI_Reduce(&res,&res0,1,MPI_DOUBLE,MPI_MAX,0,MPI_COMM_WORLD);
     return (res0);
}

       /* initial condition for iterations */
void init_poisson(int n1, int n2, double *phi, double value)
{
     int i,j;
     for (j=0; j<n2; j++)
          for (i=0; i<n1; i++) 
               phi[i+n1*j] = value;
}

void swap_ptr(double **p1, double **p2)
{
     double *tmp;
     tmp = *p1;
     *p1  = *p2;
     *p2  = tmp;
}

/* write results to disk file */
void write_array(int n1, int n2, double *p, int rank, int nproc,
                 double sx, double sy)
{
     int i,j,ij,offset,jend;
     char fn[80];
     FILE *fp;
     sprintf(fn,"p%d.dat",rank);
     fp=fopen(fn,"w");
     offset = n1;
     jend=n2+1; 
     if (rank==nproc-1) jend=n2;
     for (j=0; j<jend; j++) {
          for (i=0; i<n1; i++) {
               ij = i+n1*j;
               fprintf(fp,"%g %g %g\n",i*sx,(j+rank*n2)*sy,p[offset+ij]);
          }
          fprintf(fp,"\n");
     }    
     fclose(fp);
}

void print_one(int it, int period, double r, int rank)
{
     if ((rank==0) && (it%period==0))
          printf("it = %d: r = %g\n",it,r);
}

void print_array(int n1, int n2, double *p, int rank, int it)
{
     int i, j, jstart=1, jend=n2;
     printf("rank = %d, iter = %d\n",rank, it);
     for (j=jstart; j<jend+1; j++) {

          for (i=0; i<n1; i++) 
               printf("%g ",p[i+n1*j]);

          printf("\n");
     }
}