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

int main(int argc, char *argv[])
{
  int nlocal,offset;
  int x,n,max=0,local_max=0;
  double sum=0,local_sum=0;
  int numprocs,myrank;
  int rank_left,rank_right;
  int tag1=1,tag2=2;
  int *tmp,*buffer,*nextbuffer;

  MPI_Status status;
  MPI_Request request1,request2;

  MPI_Init(&argc,&argv);
  MPI_Comm_size(MPI_COMM_WORLD, &numprocs);
  MPI_Comm_rank(MPI_COMM_WORLD, &myrank);

  /* Exit if # of arguments < 2 */
  if (argc < 3) {
    if (myrank==0)
      printf("Usage: %s <Ndim> <Niter>\n",argv[0]);
    MPI_Finalize();
    exit(1);
  }

  int N=atoi(argv[1]);
  int Niter=atoi(argv[2]);

  if (myrank==0) {
    printf("\nSize of the array:    %d\n",N);
    printf("Number of iterations: %d\n\n",Niter);
  }

  /* Calculate the size of local array. */
  nlocal=N/numprocs;
  offset=myrank*nlocal;

  buffer = (int*)malloc((nlocal+2) * sizeof(int));
  nextbuffer = (int*)malloc((nlocal+2) * sizeof(int));

  /* Initialize the array. */
  for (x=0; x<nlocal+2; x++) {
    buffer[x]=(x+offset-1)%10;
  }
  if (myrank==0) buffer[0]=(N-1)%10;
  if (myrank==numprocs-1) buffer[nlocal+1]=0;

  rank_left=myrank-1;
  rank_right=myrank+1;

  /* Main loop */
  for (n=0; n < Niter; n++) {
    for (x=1; x<nlocal+1; x++) {
      nextbuffer[x]=(buffer[x-1]+buffer[x+1])%10;
    }
    /* Swap the pointers. */
    tmp=buffer; buffer=nextbuffer; nextbuffer=tmp;
    /* Exchange boundary elements with neighbors. */
    MPI_Isend(&buffer[1],1,MPI_INT,rank_left,tag1,MPI_COMM_WORLD,&request1);
    MPI_Isend(&buffer[nlocal],1,MPI_INT,rank_right,tag2,MPI_COMM_WORLD,&request2);
    MPI_Recv(&buffer[0], 1, MPI_INT, rank_left,tag1,MPI_COMM_WORLD,&status);
    MPI_Recv(&buffer[nlocal+1], 1, MPI_INT, rank_right,tag2,MPI_COMM_WORLD,&status);
    MPI_Wait(&request1,&status);
    MPI_Wait(&request2,&status);
  }

  /* Find the maximum and average. */
  for (x=1; x<nlocal+1; x++) {
    local_sum += buffer[x];
    if (buffer[x]>local_max)
      local_max=buffer[x];
  }

  /* Reduce the result to the root process. */
  MPI_Reduce(&local_sum,&sum,1,MPI_INT,MPI_SUM,0,MPI_COMM_WORLD);
  MPI_Reduce(&local_max,&max,1,MPI_INT,MPI_MAX,0,MPI_COMM_WORLD);

  if (myrank==0) {
    printf("Maximum: %d\n",max);
    printf("Average: %f\n\n",sum/(double)N);
  }

  free(buffer);
  free(nextbuffer);

  MPI_Finalize();

}