#include <mpi.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

int main(int argc, char *argv[])
{
  char idstr[32];
  char buff[128];
  int numprocs;
  int myid;
  double *a, *b, *c;
  int n = 6;                      // horyzon nsteps
  double time = 1.0, tau = 0.001;
  int i;
  double k = 0.06;                // heat conduction
  double h = 1.0 / (n - 1);       // horyzon step
  int remain;

  int ncur;                    // num of horyzon points in cur proc

  MPI_Status stat; 
  MPI_Init (&argc, &argv); 
  MPI_Comm_size (MPI_COMM_WORLD, &numprocs); 
  MPI_Comm_rank (MPI_COMM_WORLD, &myid); 
  
  ncur = n / numprocs;
  remain = n - numprocs * ncur;
  ncur += 2;
  while (remain--)
    if (remain == myid)
      ncur++;
  h = 1.0 / (ncur - 2 - 1);
 
  printf ("%d %d\n", myid, ncur - 2);

  a = calloc (ncur, sizeof (double));
  b = calloc (ncur, sizeof (double));
  bzero (a, ncur * sizeof (double));
  bzero (b, ncur * sizeof (double));
  if (myid == numprocs - 1) {
    a[ncur - 2] = 1.0;
    b[ncur - 2] = 1.0;
  }
    
  
  while ((time -= tau) > 0) {
    printf ("time=%f ", time);

    if (!(myid % 2)) {
      if (myid != numprocs - 1)
        MPI_Send (&a[ncur - 2], 1, MPI_DOUBLE, myid + 1, 0, MPI_COMM_WORLD);
      if (myid != 0)
        MPI_Send (&a[1], 1, MPI_DOUBLE, myid - 1, 0, MPI_COMM_WORLD);
      if (myid != numprocs - 1)
        MPI_Recv (&a[ncur - 1], 1, MPI_DOUBLE, myid + 1, 0, MPI_COMM_WORLD, &stat);
      if (myid != 0)
        MPI_Recv (&a[0], 1, MPI_DOUBLE, myid - 1, 0, MPI_COMM_WORLD, &stat);
    }
    else {
      if (myid != 0)
        MPI_Recv (&a[0], 1, MPI_DOUBLE, myid - 1, 0, MPI_COMM_WORLD, &stat);
      if (myid != numprocs - 1)
        MPI_Recv (&a[ncur - 1], 1, MPI_DOUBLE, myid + 1, 0, MPI_COMM_WORLD, &stat);
      if (myid != 0)
        MPI_Send (&a[1], 1, MPI_DOUBLE, myid - 1, 0, MPI_COMM_WORLD);
      if (myid != numprocs - 1)
        MPI_Send (&a[ncur - 2], 1, MPI_DOUBLE, myid + 1, 0, MPI_COMM_WORLD);
    }
    
    printf ("runk=%d ", myid);
    
    if (myid != 0)
      b[1] = a[1] + k * tau / h / h * (a[2] + a[0] - 2 * a[1]);
    printf ("b[1]=%f ", b[1]);
    
    for (i = 2; i < ncur - 2; ++i) {
      b[i] = a[i] + k * tau / h / h * (a[i + 1] + a[i - 1] - 2 * a[i]);
      printf ("b[%d]=%f ", i, b[i]);
    }
    
    if (myid != numprocs - 1)
      b[ncur - 2] = a[ncur - 2] + k * tau / h / h * (a[ncur - 1] + a[ncur - 3] - 2 * a[ncur - 2]);
    printf ("b[%d]=%f\n", ncur - 2, b[ncur - 2]);
    
    c = a; a = b; b = c;
  }
    

  MPI_Finalize();
  return EXIT_SUCCESS;
}