#ifdef __MPI__
#else
  #include <mpi-common.h>
  #ifdef __CIVLC__
  #else
    #include <civlc.h>
  #endif
#endif

//TODO make a Datatype struct, which has a field "int size;" Define one of these objects for MPI_INT, MPI_DOUBLE, etc. 
//TODO Then provide methods like MPI provides for creating new ones.
//TODO then support MPI_Type_contig(datatype, int n).

/****************************** Helper Functions **********************************/
int sizeofDatatype(MPI_Datatype datatype) {
  switch (datatype) {
  case MPI_INT:
    return sizeof(int);
  case MPI_FLOAT:
    return sizeof(float);
  case MPI_DOUBLE:
    return sizeof(double);
  case MPI_CHAR:
    return sizeof(char);
  default:
    $assert(0, "Unreachable");
  }
}

/* Helpers for MPI_Reduce and MPI_Allreduce */
void sumInt(int result[], int buf[], int real_count, int nprocs){
  //init result array
  $atomic{    
    for(int i=0; i<real_count; i++){
      result[i] = 0;
    }
    //sum up
    for(int i=0; i<nprocs; i++)
      for(int j=0; j<real_count; j++)
	result[j] = result[j] + buf[i * real_count + j];
  }
}

void sumFloat(float result[], float buf[], int real_count, int nprocs){
  $atomic{    
    for(int i=0; i<real_count; i++){
      result[i] = 0;
    }
    //sum up
    for(int i=0; i<nprocs; i++)
      for(int j=0; j<real_count; j++)
	result[j] = result[j] + buf[i * real_count + j];
  }
}

void sumDouble(double result[], double buf[], int real_count, int nprocs){
  $atomic{    
    for(int i=0; i<real_count; i++){
      result[i] = 0;
    }
    //sum up
    for(int i=0; i<nprocs; i++)
      for(int j=0; j<real_count; j++)
	result[j] = result[j] + buf[i * real_count + j];
  }
}


void maxInt(int result[], int buf[], int real_count, int nprocs){
  //init result arra
  $atomic{
    for(int i=0; i<real_count; i++){
      result[i] = buf[i * real_count];
    }
    
    for(int i=0; i<nprocs; i++)
      for(int j=0; j<real_count; j++){
	if(buf[i * real_count + j] > result[j])
	  result[j] = buf[i * real_count + j];
      }
  }
}

void maxFloat(float result[], float buf[], int real_count, int nprocs){
  //init result array
  $atom{
    for(int i=0; i<real_count; i++){
      result[i] = buf[i * real_count];
  }

    for(int i=0; i<nprocs; i++)
      for(int j=0; j<real_count; j++){
	if(buf[i * real_count + j] > result[j])
	  result[j] = buf[i * real_count + j];
      }
  }
}

void maxDouble(double result[], double buf[], int real_count, int nprocs){
  //init result array
  $atom{
    for(int i=0; i<real_count; i++){
      result[i] = buf[i * real_count];
    }
    
    for(int i=0; i<nprocs; i++)
      for(int j=0; j<real_count; j++){
	if(buf[i * real_count + j] > result[j])
	  result[j] = buf[i * real_count + j];
      }
  }
}

void minInt(int result[], int buf[], int real_count, int nprocs){
  //init result array
  $atom{
    for(int i=0; i<real_count; i++){
    result[i] = buf[i * real_count];
  }

    for(int i=0; i<nprocs; i++)
      for(int j=0; j<real_count; j++){
	if(buf[i * real_count + j] < result[j])
	  result[j] = buf[i * real_count + j];
      }
  }
}
  
void minFloat(float result[], float buf[], int real_count, int nprocs){
  //init result array
  $atom{
    for(int i=0; i<real_count; i++){
      result[i] = buf[i * real_count];
    }

    for(int i=0; i<nprocs; i++)
      for(int j=0; j<real_count; j++){
	if(buf[i * real_count + j] < result[j])
	  result[j] = buf[i * real_count + j];
      }
  }
}

void minDouble(double result[], double buf[], int real_count, int nprocs){
  //init result array
  $atom{
    for(int i=0; i<real_count; i++){
      result[i] = buf[i * real_count];
    }

    for(int i=0; i<nprocs; i++)
      for(int j=0; j<real_count; j++){
	if(buf[i * real_count + j] < result[j])
	  result[j] = buf[i * real_count + j];
      }
  }
}

/************************** MPI LIB Implementations *******************************/
CMPI_Gcomm CMPI_Gcomm_create($scope scope, int size) {
  CMPI_Gcomm result;

  result.p2p = $gcomm_create(scope, size);
  result.col = $gcomm_create(scope, size);
  return result;
}

void CMPI_Gcomm_destroy(CMPI_Gcomm gc) {
  $gcomm_destroy(gc.p2p);
  $gcomm_destroy(gc.col);
}

MPI_Comm MPI_Comm_create($scope scope, CMPI_Gcomm gc, int rank) {
  MPI_Comm result;
  
  result.p2p = $comm_create(scope, gc.p2p, rank);
  result.col = $comm_create(scope, gc.col, rank);
  result.status = __UNINIT;
  return result;
}

void MPI_Comm_destroy(MPI_Comm comm) {
  $comm_destroy(comm.p2p);
  $comm_destroy(comm.col);
}

int __MPI_Init(MPI_Comm *comm) {
  comm->status = __INIT;
  return 0;
}

int __MPI_Finalize(MPI_Comm *comm) {
  comm->status = __FINALIZED;
  return 0;
}

int MPI_Comm_size(MPI_Comm comm, int *size) {
  $assert(comm.status == __INIT, "MPI_Comm_size() cannot be invoked without MPI_Init() being called before.\n");
  *size = $comm_size(comm.p2p);
  return 0;
}

int MPI_Comm_rank(MPI_Comm comm, int *rank) {
  $assert(comm.status == __INIT, "MPI_Comm_rank() cannot be invoked without MPI_Init() being called before.\n");
  *rank = $comm_place(comm.p2p);
  return 0;
}


int CMPI_Send(void *buf, int count, MPI_Datatype datatype, int dest,
	      int tag, $comm comm) {
  if (dest >= 0) {
    int size = count*sizeofDatatype(datatype);
    int place = $comm_place(comm);
    $message out = $message_pack(place, dest, tag, buf, size);

    $comm_enqueue(comm, out);
  }
  return 0;
}

int MPI_Send(void *buf, int count, MPI_Datatype datatype, int dest,
	     int tag, MPI_Comm comm) {
  $assert(comm.status == __INIT, "MPI_Send() cannot be invoked without MPI_Init() being called before.\n");
  return CMPI_Send(buf, count, datatype, dest, tag, comm.p2p);
}


int CMPI_Recv(void *buf, int count, MPI_Datatype datatype, int source,
	      int tag, $comm comm, MPI_Status *status) {
  if (source >= 0) {
    $message in = $comm_dequeue(comm, source, tag);
    int size = count*sizeofDatatype(datatype);

    $message_unpack(in, buf, size);
    if (status != MPI_STATUS_IGNORE) {
      status->size = $message_size(in);
      status->MPI_SOURCE = $message_source(in);
      status->MPI_TAG = $message_tag(in);
      status->MPI_ERROR = 0;
    }
  }
  return 0;
}

int MPI_Recv(void *buf, int count, MPI_Datatype datatype, int source,
	     int tag, MPI_Comm comm, MPI_Status *status) {
  $assert(comm.status == __INIT, "MPI_Recv() cannot be invoked without MPI_Init() being called before.\n");
  return CMPI_Recv(buf, count, datatype, source, tag, comm.p2p, status);
}

int MPI_Get_count(MPI_Status *status, MPI_Datatype datatype, int *count) {
  //$assert(__my_status == __INIT, "MPI status is not INIT.\n");
  *count = status->size/sizeofDatatype(datatype);
  return 0;
}

int MPI_Sendrecv(void *sendbuf, int sendcount, MPI_Datatype sendtype,
		 int dest, int sendtag,
		 void *recvbuf, int recvcount, MPI_Datatype recvtype,
		 int source, int recvtag,
		 MPI_Comm comm, MPI_Status *status) {
  $assert(comm.status == __INIT, "MPI_Sendrecv() cannot be invoked without MPI_Init() being called before.\n");
  MPI_Send(sendbuf, sendcount, sendtype, dest, sendtag, comm);
  MPI_Recv(recvbuf, recvcount, recvtype, source, recvtag, comm, status);
  return 0;
}

/* Broadcasts a message from root to everyone else.
 * Need to use a differnt comm.
 */
int MPI_Bcast(void *buf, int count, MPI_Datatype datatype, int root,
	      MPI_Comm comm) {
  int place = $comm_place(comm.col);

  $assert(comm.status == __INIT, "MPI_Bcast() cannot be invoked without MPI_Init() being called before.\n");
  place = $comm_place(comm.col);
  if (place == root) {
    int nprocs = $comm_size(comm.col);

    for (int i=0; i<nprocs; i++)
      if (i != root)
	CMPI_Send(buf, count, datatype, i, BCAST_TAG, comm.col);
  } else {
    CMPI_Recv(buf, count, datatype, root, BCAST_TAG, comm.col,
	     MPI_STATUS_IGNORE);
  }
  return 0;
}

/* Reduces values on all processes to a single value */
int MPI_Reduce(void* sendbuf, void* recvbuf, int count, MPI_Datatype datatype,
               MPI_Op op, int root, MPI_Comm comm){
  int place;
  MPI_Status status;
  
  $assert(comm.status == __INIT, "MPI_Reduce() cannot be invoked without MPI_Init() being called before.\n");
  place = $comm_place(comm.col);
  //non-root
  CMPI_Send(sendbuf, count, datatype, root, REDUCE_TAG, comm.col);
  if(place != root)
    return 0;
  else{
    //root
    int nprocs = $comm_size(comm.col);
    int real_count = -1;               //the real count gotton from MPI_Status
    //void ** buf;                     //Buffer stores data from other processes
    void * recv;                       //Buffer used in MPI_Recv()
    void * results;                    //Array of results
    $scope here = $root;
    struct Bundle_buf{
      int I[nprocs * count];
      float F[nprocs * count];
      double D[nprocs * count];
    } buf;

    switch(datatype){
    case MPI_INT :
      recv = (int *)$malloc(here, count * sizeof(int));
      results = (int *)$malloc(here, count * sizeof(int));
      break;
    case MPI_FLOAT :
      recv = (float *)$malloc(here, count * sizeof(float));
      results = (float *)$malloc(here, count * sizeof(float));
      break;
    case MPI_DOUBLE :
      recv = (double *)$malloc(here, count * sizeof(double));
      results = (double *)$malloc(here, count * sizeof(double));
      break;
    }
    for(int i=0; i<nprocs; i++){
      CMPI_Recv(recv, count, datatype, i, REDUCE_TAG, comm.col, 
		&status);
      MPI_Get_count(&status, datatype, &real_count);
      $assert(real_count == count);

      switch(datatype){
      case MPI_INT :
	for(int j=0; j<real_count; j++)
	  buf.I[i * real_count + j] = ((int *)recv)[j];
	break;
      case MPI_FLOAT : 
	for(int j=0; j<real_count; j++)
	  buf.F[i * real_count + j] = ((float *)recv)[j];
	break;
      case MPI_DOUBLE : 
	for(int j=0; j<real_count; j++)
	  buf.D[i * real_count + j] = ((double *)recv)[j];
	break;
      }
    }
    
    //operations
    switch(datatype){
    case MPI_INT: 
      if(op == MPI_SUM)
	sumInt(results, buf.I, real_count, nprocs);
      if(op == MPI_MAX)
	maxInt(results, buf.I, real_count, nprocs);
      if(op == MPI_MIN)
	minInt(results, buf.I, real_count, nprocs);
      break;
    case MPI_FLOAT:
      if(op == MPI_SUM)
	sumFloat(results, buf.F, real_count, nprocs);
      if(op == MPI_MAX)
	maxFloat(results, buf.F, real_count, nprocs);
      if(op == MPI_MIN)
	minFloat(results, buf.F, real_count, nprocs);
      break;
    case MPI_DOUBLE:
      if(op == MPI_SUM)
	sumDouble(results, buf.D, real_count, nprocs);
      if(op == MPI_MAX)
	maxDouble(results, buf.D, real_count, nprocs);
      if(op == MPI_MIN)
	minDouble(results, buf.D, real_count, nprocs);
      break;
    } 
    MPI_Sendrecv(results, real_count, datatype,
		     place, 0,
		     recvbuf, real_count, datatype,
		     place, 0,
		     comm, &status);
    $free(recv);
    $free(results);
  }
  return 0;
}

/* Combines values from all processes and distributes the result back to all processes */
/* default root is 0 */
int MPI_Allreduce(void* sendbuf, void* recvbuf, int count, MPI_Datatype datatype,
               MPI_Op op, MPI_Comm comm){
  int place;
  int root;
  MPI_Status status;
  
  $assert(comm.status == __INIT, "MPI_Allreduce() cannot be invoked without MPI_Init() being called before.\n");
  place = $comm_place(comm.col);
  root = 0;
  //non-root
  CMPI_Send(sendbuf, count, datatype, root, REDUCE_TAG, comm.col);
  if(place != root){
    MPI_Bcast(recvbuf, count, datatype, root,
	      comm);
    return 0;
  }
  else{
    //root
    int nprocs = $comm_size(comm.col);
    int real_count = -1;               //the real count gotton from MPI_Status
    //void ** buf;                     //Buffer stores data from other processes
    void * recv;                       //Buffer used in MPI_Recv()
    void * results;                    //Array of results
    $scope here = $root;
    struct Bundle_buf{
      int I[nprocs * count];
      float F[nprocs * count];
      double D[nprocs * count];
    } buf;

    switch(datatype){
    case MPI_INT :
      recv = (int *)$malloc(here, count * sizeof(int));
      results = (int *)$malloc(here, count * sizeof(int));
      break;
    case MPI_FLOAT :
      recv = (float *)$malloc(here, count * sizeof(float));
      results = (float *)$malloc(here, count * sizeof(float));
      break;
    case MPI_DOUBLE :
      recv = (double *)$malloc(here, count * sizeof(double));
      results = (double *)$malloc(here, count * sizeof(double));
      break;
    }
    //recv = (void *)$malloc(here, count * sizeofDatatype(datatype));
    //results = (void *)$malloc(here, count * sizeofDatatype(datatype));

    for(int i=0; i<nprocs; i++){
      CMPI_Recv(recv, count, datatype, i, REDUCE_TAG, comm.col, 
		&status);
      MPI_Get_count(&status, datatype, &real_count);
      $assert(real_count == count);

      switch(datatype){
      case MPI_INT :
	for(int j=0; j<real_count; j++)
	  buf.I[i * real_count + j] = ((int *)recv)[j];
	break;
      case MPI_FLOAT : 
	for(int j=0; j<real_count; j++)
	  buf.F[i * real_count + j] = ((float *)recv)[j];
	break;
      case MPI_DOUBLE : 
	for(int j=0; j<real_count; j++)
	  buf.D[i * real_count + j] = ((double *)recv)[j];
	break;
      }
    }

    //operations
    switch(datatype){
    case MPI_INT: 
      if(op == MPI_SUM)
	sumInt(results, buf.I, real_count, nprocs);
      if(op == MPI_MAX)
	maxInt(results, buf.I, real_count, nprocs);
      if(op == MPI_MIN)
	minInt(results, buf.I, real_count, nprocs);
      break;
    case MPI_FLOAT:
      if(op == MPI_SUM)
	sumFloat(results, buf.F, real_count, nprocs);
      if(op == MPI_MAX)
	maxFloat(results, buf.F, real_count, nprocs);
      if(op == MPI_MIN)
	minFloat(results, buf.F, real_count, nprocs);
      break;
    case MPI_DOUBLE:
      if(op == MPI_SUM)
	sumDouble(results, buf.D, real_count, nprocs);
      if(op == MPI_MAX)
	maxDouble(results, buf.D, real_count, nprocs);
      if(op == MPI_MIN)
	minDouble(results, buf.D, real_count, nprocs);
      break;
    }
    MPI_Sendrecv(results, real_count, datatype,
		     place, 0,
		     recvbuf, real_count, datatype,
		     place, 0,
		     comm, &status);
    MPI_Bcast(results, real_count, datatype, root,
	      comm);
    $free(recv);
    $free(results);
  }
  return 0;
}