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

/* Data size in REAL's, per process (override with compiler flag if you'd like) */
/* Note: in most real applications this would be quite small (just local norms, say) */
#define DATASIZE 32

/* How many burnin and trial runs to perform */
#define NBURNIN 5
#define NTRIALS 100

/* to scale local work time */
#define LOCALREPS 50000

/* which kind of reduction to test */
#define REDOP MPI_SUM

/* datatype to test with (double is probably most relevant) */
#define REAL double
#define MPI_MY_REAL MPI_DOUBLE

#define CHKERR(x) do{if(x>0){printf("MPI Error on line __LINE__"); exit(1);}}while(0)

/* Do some useless work on a scalar */
#define WORKFUNCTION(x) d*(d + 1.2232*(232/(d+232.3 - 0.1*d+(d-3.2))))/(21232.34 + d*d/(2323.44/(d*(d+2.232))))

/***************************************************************************************\
  Ctx to operate on
\***************************************************************************************/
/* A structure to define problem data to be operated upon */
typedef struct {
  REAL * array_send;
  REAL * array_recv;
  REAL * array_work;
} Data;

/* Forward Declarations */
void trial_blockingredonly(Data*);
void trial_nonblockingredonly(Data*);
void trial_loconly(Data*);
void trial_blocking(Data*);
void trial_nonblocking(Data*);

/* Create and destroy data to operate on */
void create(Data * pData)
{
  int i;

  pData->array_recv = (REAL*) malloc(DATASIZE * sizeof(REAL));
  pData->array_send = (REAL*) malloc(DATASIZE * sizeof(REAL));
  for(i=0;i<DATASIZE;++i){
    pData->array_send[i] = 4389034389*i + 34933434244;
  }
  pData->array_work = (REAL*) malloc(DATASIZE * sizeof(REAL));
}

void destroy(Data * pData)
{
  free(pData->array_recv);
  free(pData->array_send);
  free(pData->array_work);
}

/***************************************************************************************\
  Main
\***************************************************************************************/
int main(int argc,char ** argv)
{
  int i,ierr;
  REAL avg_time_blocking, avg_time_nonblocking, avg_time_blockingredonly, avg_time_nonblockingredonly, avg_time_loconly;
  REAL totals[5],time_min[5],time_max[5];
  int rank,size;

  /* Setup */
  ierr = MPI_Init(&argc,&argv); CHKERR(ierr);
  ierr = MPI_Comm_rank(MPI_COMM_WORLD,&rank);CHKERR(ierr);
  ierr = MPI_Comm_size(MPI_COMM_WORLD,&size);CHKERR(ierr);

  if (!rank) {
    printf("Testing with %d MPI ranks\nreducing an array of size %d (%ld bytes)\n",size,DATASIZE,DATASIZE * sizeof(REAL));
    printf("Running %d burnin runs and %d tests ... ",NBURNIN,NTRIALS);
  }

  /* Burn in 1  */
  {
    Data data;
    create(&data);
    for(i=0;i<NBURNIN;++i){
      trial_blockingredonly(&data);
    }
    destroy(&data);
  }
  {
    Data data;
    create(&data);
    for(i=0;i<NBURNIN;++i){
      trial_nonblockingredonly(&data);
    }
    destroy(&data);
  }
  {
    Data data;
    create(&data);
    for(i=0;i<NBURNIN;++i){
      trial_loconly(&data);
    }
    destroy(&data);
  }
  {
    Data data;
    create(&data);
    for(i=0;i<NBURNIN;++i){
      trial_blocking(&data);
    }
    destroy(&data);
  }
  {
    Data data;
    create(&data);
    for(i=0;i<NBURNIN;++i){
      trial_nonblocking(&data);
    }
    destroy(&data);
  }

  /* Time runs  */

  /* Blocking Reduction Only */
  avg_time_blockingredonly = 0.0;
  {
    Data data;
    double t;
    create(&data);
    for(i=0;i<NBURNIN;++i){
      trial_blockingredonly(&data);
    }
    t= MPI_Wtime();
    for(i=0;i<NTRIALS;++i){
      trial_blockingredonly(&data);
    }
    t = MPI_Wtime() - t;
    destroy(&data);
    avg_time_blockingredonly = t/NTRIALS;
  }
  /* Non-Blocking Reduction Only */
  avg_time_nonblockingredonly = 0.0;
  {
    Data data;
    double t;
    create(&data);
    for(i=0;i<NBURNIN;++i){
      trial_nonblockingredonly(&data);
    }
    t= MPI_Wtime();
    for(i=0;i<NTRIALS;++i){
      trial_nonblockingredonly(&data);
    }
    t = MPI_Wtime() - t;
    destroy(&data);
    avg_time_nonblockingredonly = t/NTRIALS;
  }

  /* Local Work Only */
  avg_time_loconly = 0.0;
  {
    Data data;
    double t;
    create(&data);
    for(i=0;i<NBURNIN;++i){
      trial_loconly(&data);
    }
    t= MPI_Wtime();
    for(i=0;i<NTRIALS;++i){
      trial_loconly(&data);
    }
    t= MPI_Wtime() - t;
    destroy(&data);
    avg_time_loconly = t/NTRIALS;
  }
  
  /* Blocking */
  avg_time_blocking = 0.0;
  {
    Data data;
    double t;
    create(&data);
    for(i=0;i<NBURNIN;++i){
      trial_blocking(&data);
    }
    t= MPI_Wtime();
    for(i=0;i<NTRIALS;++i){
      trial_blocking(&data);
    }
    t= MPI_Wtime() - t;
    destroy(&data);
    avg_time_blocking = t/NTRIALS;
  }

  /* Nonblocking */
  avg_time_nonblocking       = 0.0;
  {
    Data data;
    double t;
    create(&data);
    for(i=0;i<NBURNIN;++i){
      trial_nonblocking(&data);
    }
    t= MPI_Wtime();
    for(i=0;i<NTRIALS;++i){
      trial_nonblocking(&data);
    }
    t= MPI_Wtime() - t;
    destroy(&data);
    avg_time_nonblocking = t/NTRIALS;
  }
 
  /* Collect Data */
  totals[0] = avg_time_blockingredonly;
  totals[1] = avg_time_nonblockingredonly;
  totals[2] = avg_time_loconly;
  totals[3] = avg_time_blocking;
  totals[4] = avg_time_nonblocking;
  ierr = MPI_Allreduce(totals,time_min,5,MPI_MY_REAL,MPI_MIN,MPI_COMM_WORLD);CHKERR(ierr);
  ierr = MPI_Allreduce(totals,time_max,5,MPI_MY_REAL,MPI_MAX,MPI_COMM_WORLD);CHKERR(ierr);

  /* On rank 0, print averages */
  if(!rank){
    printf(" Done.\n");
    printf("For %d runs with %d burnin runs, on %d MPI processes, min/max times over all ranks:\n",NTRIALS,NBURNIN,size);
    printf("MPI timer resolution:         %1.4e seconds\n",MPI_Wtick());
    printf("MPI timer resolution/#trials: %1.4e seconds\n",MPI_Wtick()/((double)NTRIALS));
    printf("B.   Red. Only (min/max):    %1.6e / %1.6e seconds\n",time_min[0],time_max[0]);
    printf("N.B. Red. Only (min/max):    %1.6e / %1.6e seconds\n",time_min[1],time_max[1]);
    printf("Loc. Only (min/max):         %1.6e / %1.6e seconds\n",time_min[2],time_max[2]);

    printf("Blocking (min/max):          %1.6e / %1.6e seconds\n",time_min[3],time_max[3]);
    printf("Non-Blocking (min/max):      %1.6e / %1.6e seconds\n",time_min[4],time_max[4]);

    printf("Observe to see if the local time is enough to hide the reduction, and see if the reduction is indeed hidden\n");
  }

  /* Finalize */
  MPI_Finalize();
  return 0;
}

/***************************************************************************************\
  Overlapping Work
\***************************************************************************************/
/* Define the work to be overlapped with the reduction */
void other_work(Data * pData)
{
  /*
  int rank,size,left,right,ierr;
  ierr = MPI_Comm_rank(MPI_COMM_WORLD,&rank); CHKERR(ierr);
  ierr = MPI_Comm_size(MPI_COMM_WORLD,&size); CHKERR(ierr);
  right = rank < size ? rank+1 : 0;
  left  = rank > 0    ? rank-1 : size-1;
  */

  /* Local work */
  {
    int j;
    int dummy;
    for(j=0;j<LOCALREPS;++j){
      int i = j % DATASIZE; 
      REAL d = pData->array_work[i];
      //pData->array_work[i] = d*(d + 1.2232*(232/(d+232.3 - 0.1*d+(d-3.2))))/(21232.34 + d*d/(2323.44/(d*(d+2.232))));
      pData->array_work[i] = WORKFUNCTION(d);
      //pData->array_work[i] = simple_rand_n(d,10);
    }
  }
}

/***************************************************************************************\
  Blocking Trial
\***************************************************************************************/
void trial_blocking(Data * pData)
{
  int ierr;

  other_work(pData);
  ierr = MPI_Allreduce(pData->array_send,pData->array_recv,DATASIZE,MPI_MY_REAL,MPI_SUM,MPI_COMM_WORLD);CHKERR(ierr);
}

/***************************************************************************************\
  Nonblocking Trial
\***************************************************************************************/
void trial_nonblocking(Data * pData)
{
  int ierr;
  MPI_Request request;

  ierr = MPI_Iallreduce(pData->array_send,pData->array_recv,DATASIZE,MPI_MY_REAL,MPI_SUM,MPI_COMM_WORLD,&request);CHKERR(ierr);
  other_work(pData);
  ierr = MPI_Wait(&request,MPI_STATUS_IGNORE);
}

/***************************************************************************************\
  Reduction Only Controls
\***************************************************************************************/
void trial_blockingredonly(Data * pData)
{
  int ierr;
  ierr = MPI_Allreduce(pData->array_send,pData->array_recv,DATASIZE,MPI_MY_REAL,MPI_SUM,MPI_COMM_WORLD);CHKERR(ierr);
}

void trial_nonblockingredonly(Data * pData)
{
  int ierr;
  MPI_Request r;

  ierr = MPI_Iallreduce(pData->array_send,pData->array_recv,DATASIZE,MPI_MY_REAL,MPI_SUM,MPI_COMM_WORLD,&r);CHKERR(ierr);
  ierr = MPI_Wait(&r,MPI_STATUS_IGNORE);CHKERR(ierr);
}

/***************************************************************************************\
  Local work only control
\***************************************************************************************/
void trial_loconly(Data * pData)
{
  other_work(pData);
}