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

#define MAX_POINTS 1000
#define min(A,B) ((A)<(B) ? (A) : (B))
#define max(A,B) ((A)>(B) ? (A) : (B))
#define PI 3.1415926536
#define NSTEPS 100

main (argc, argv)
int argc;
char *argv[];
{
   double psi[MAX_POINTS+2];
   double new_psi[MAX_POINTS+2];
   double old_psi[MAX_POINTS+2];
   int rank, nprocs, mypos;
   int periods=0, reorder=1;
   int i, j, local_npts, lbound, istart, iend, a_tag=112;
   int npoints, left, right, proc;
   double x, tau=0.05;
   MPI_Comm new_comm;
   MPI_Status status;
   FILE *fp;

/* 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);

/* Read in input parameters */
   if (rank==0) {  
      fp = fopen ("input-wave.dat", "r");
      fscanf (fp, "%d", &npoints);
   }
   
   MPI_Bcast (&npoints, 1, MPI_INT, 0, MPI_COMM_WORLD);

   if (npoints%nprocs==0) {
/* Set up 1D topology */
      MPI_Cart_create (MPI_COMM_WORLD,1,&nprocs,&periods,reorder,&new_comm);
      MPI_Cart_coords (new_comm, rank, 1, &mypos);
      MPI_Cart_shift  (new_comm, 0, -1, &right, &left);

/* Initialise array */
      local_npts = npoints/nprocs;
      lbound     = mypos*local_npts;
      for(i=0;i<local_npts;i++){
          x = 2.0*PI*(double)(lbound+i)/(double)(npoints-1);
          x = sin(x);
          psi[i+1] = old_psi[i+1] = x;
      }

/* Do the updates */
      istart = (mypos==0) ? 2 : 1;
      iend   = (mypos==nprocs-1) ? local_npts-1 : local_npts;
      for(j=0;j<NSTEPS;j++){
         MPI_Sendrecv (&psi[1],            1, MPI_DOUBLE, left,  111, 
                       &psi[local_npts+1], 1, MPI_DOUBLE, right, 111,
                       new_comm, &status);
         MPI_Sendrecv (&psi[local_npts], 1, MPI_DOUBLE, right,  111, 
                       &psi[0],          1, MPI_DOUBLE, left,   111,
                       new_comm, &status);
         for(i=istart;i<=iend;i++){
            new_psi[i] = 2.0*psi[i] - old_psi[i] +
                         tau*tau*(psi[i-1]-2.0*psi[i]+psi[i+1]);
         }
         for(i=1;i<=local_npts;i++){
            old_psi[i] = psi[i];
            psi[i]     = new_psi[i];
         }
      }
/* Output the solution */
      for(j=0;j<nprocs;j++){
         if (mypos==0){
            MPI_Cart_rank (new_comm, &j, &proc);
            if (j>0) MPI_Recv (&psi[1], local_npts, MPI_DOUBLE,
                               proc, a_tag, new_comm, &status);
            for(i=1;i<=local_npts;i++)
               printf("%9.6f%c",psi[i],(i%8==0) ? '\n' : ' ');
          }
          else if (mypos==j) 
             MPI_Send (&psi[1], local_npts, MPI_DOUBLE, 0,
                       a_tag, new_comm);
      }
   }
   else if (rank==0) 
      printf("\n npiints must be exactly divisible by nprocs");

   MPI_Finalize ();
}