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

#define NXMAX 32
#define NYMAX 32

main (argc, argv)
int argc;
char *argv[];
{
   double sbuf[NYMAX], rbuf[NYMAX];
   double phi[NYMAX+2][NXMAX+2];
   double oldphi[NYMAX+2][NXMAX+2];
   int    mask[NYMAX+2][NXMAX+2];
   int    i, j, k;
   int    nptsx = 32, nptsy = 32;
   int    nsteps  = 20;
   int    rank, nprocs, dims[2], periods[2], coords[2];
   int    myposx, myposy, npx, npy;
   int    up, down, left, right;
   int    local_nptsx, local_nptsy;
   MPI_Comm new_comm;
   MPI_Status status;

/* Initialise and find rank and number of processes */
   MPI_Init (&argc, &argv);
   MPI_Comm_size (MPI_COMM_WORLD, &nprocs);
   MPI_Comm_rank (MPI_COMM_WORLD, &rank);

/* Work out number of processes in each direction of the process mesh */
   dims[0] = dims[1] = 0;
   MPI_Dims_create (nprocs, 2, dims);
   npy = dims[0];
   npx = dims[1];

   if (nptsx%npx==0 && nptsy%npy==0) {

/* Set up 2D topology */
      periods[0] = periods[1] = 0;
      MPI_Cart_create (MPI_COMM_WORLD, 2, dims, periods, 1, &new_comm);
      MPI_Cart_coords (new_comm, rank, 2, coords);
      myposy = coords[0];
      myposx = coords[1];

/* Determine neighbouring processes for communication shifts */
      MPI_Cart_shift  (new_comm, 0, -1,  &down,   &up);
      MPI_Cart_shift  (new_comm, 1, -1, &right, &left);

/* Initialise arrays */
      setup_grid (phi, npx, npy, myposx, myposy, nptsx, nptsy, 
                  &local_nptsx, &local_nptsy, mask);
 
/* Iterate to find solution */
      for(k=1;k<=nsteps;k++){
         for(j=1;j<=local_nptsy;j++)
            for(i=1;i<=local_nptsx;i++)
               oldphi[j][i] = phi[j][i];
         MPI_Sendrecv (&oldphi[1][1],  local_nptsx, MPI_DOUBLE,  up, 111,
            &oldphi[local_nptsy+1][1], local_nptsx, MPI_DOUBLE,down, 111,
              new_comm, &status);
         MPI_Sendrecv (&oldphi[local_nptsy][1],local_nptsx,MPI_DOUBLE,down,112,
            &oldphi[0][1], local_nptsx, MPI_DOUBLE, up, 112,
              new_comm, &status);
         for(i=1;i<=local_nptsy;i++) sbuf[i-1] = oldphi[i][1];
         MPI_Sendrecv (sbuf, local_nptsy, MPI_DOUBLE, left, 113,
                       rbuf, local_nptsy, MPI_DOUBLE, right, 113,
                       new_comm, &status);
         for(i=1;i<=local_nptsy;i++) oldphi[i][local_nptsx+1] = rbuf[i-1];   
         for(i=1;i<=local_nptsy;i++) sbuf[i-1] = oldphi[i][local_nptsx];
         MPI_Sendrecv (sbuf, local_nptsy, MPI_DOUBLE, right, 114,
                       rbuf, local_nptsy, MPI_DOUBLE, left, 114,
                       new_comm, &status);
         for(i=1;i<=local_nptsy;i++) oldphi[i][0] = rbuf[i-1];
         for(j=1;j<=local_nptsy;j++)
            for(i=1;i<=local_nptsx;i++)
               if (mask[j][i]) phi[j][i] = 0.25*(oldphi[j][i-1] +
                  oldphi[j][i+1] + oldphi[j-1][i] + oldphi[j+1][i]);
      }
      output_array (phi, rank, npx, npy, local_nptsx, local_nptsy, 
                    nprocs , myposx, myposy, new_comm);
   }
   else if (rank==0) {
      printf ("\nNumber of grid points in each direction must be exactly");
      printf ("\ndivisible by the number of processes in that direction");
   }
 
   MPI_Finalize();
}

setup_grid (phi, npx, npy, myposx, myposy, nptsx, nptsy, 
            local_ptrx, local_ptry, mask)
double  phi[NYMAX+2][NXMAX+2];
int     npx, npy, myposx, myposy, nptsx, nptsy;
int    *local_ptrx, *local_ptry;
int     mask[NYMAX+2][NXMAX+2];
{
    int i, j;
    int local_nptsx, local_nptsy, global_x, global_y;

    local_nptsx = nptsx/npx;
    local_nptsy = nptsy/npy;
  
    for(j=0;j<=local_nptsy+1;j++)
       for(i=0;i<=local_nptsx+1;i++){
          phi[j][i]  = 0.0;
          mask[j][i] = 1;
       }

    if (myposy == 0)
       for(i=0;i<=local_nptsx+1;i++) mask[1][i] = 0;

    if (myposy == npy-1)
       for(i=0;i<=local_nptsx+1;i++) mask[local_nptsy][i] = 0;

    if (myposx == 0)
       for(j=0;j<=local_nptsy+1;j++) mask[j][1] = 0;

    if (myposx == npx-1)
       for(j=0;j<=local_nptsy+1;j++) mask[j][local_nptsx] = 0;

    for(j=1;j<=local_nptsy;j++){
       global_y = local_nptsy*myposy + j - 1;
       if (global_y == nptsy/2 || global_y == nptsy/2-1){
          for(i=1;i<=local_nptsx;i++){
             global_x = local_nptsx*myposx + i - 1;
             if (global_x == nptsx/2 || global_x == nptsx/2-1){
                mask[j][i] = 0;
                phi[j][i]  = 1.0;
             }
          }
       }
    }

    *local_ptrx = local_nptsx;
    *local_ptry = local_nptsy;
}

output_array (phi, rank, npx, npy, local_nptsx, local_nptsy, 
              nprocs, myposx, myposy, new_comm)
double phi[NYMAX+2][NXMAX+2];
int    rank, npx, npy, local_nptsx, local_nptsy, nprocs, myposx, myposy;
MPI_Comm new_comm;
{
   int i, j, k=0, m, n;
   int source;
   int coords[2];
   MPI_Status status;

   for(m=0;m<npy;m++)
      for(j=1;j<=local_nptsy;j++)
         for(n=0;n<npx;n++){
            if (myposx == n && myposy == m)
               if (rank!=0) MPI_Send (&phi[j][1], local_nptsx,
                  MPI_DOUBLE, 0, 115,MPI_COMM_WORLD);
               else
                  for(i=1;i<=local_nptsx;i++,k++)
                     printf("%9.6f%c",phi[j][i],((k+1)%8==0) ? '\n' : ' ');
            else if (rank==0){
               coords[0] = m;
               coords[1] = n;
               MPI_Cart_rank (new_comm, coords, &source);
               MPI_Recv (&phi[1][1], local_nptsx, MPI_DOUBLE,
                         source, 115, MPI_COMM_WORLD, &status);
               for(i=1;i<=local_nptsx;i++,k++)
                  printf("%9.6f%c",phi[1][i],((k+1)%8==0) ? '\n' : ' ');
            }
         }
}