/* TASS MPI library.  Currently supporting a subset of MPI 2.2...
 *
 * Types defined: 
 *   MPI_Comm
 *   MPI_Request
 *   MPI_Status : this is a record with at least fields MPI_SOURCE and MPI_TAG
 *   MPI_Datatype
 *   MPI_Op
 *
 * Constants defined:
 *   MPI_COMM_WORLD    : MPI_Comm
 *   MPI_ANY_SOURCE    : int
 *   MPI_ANY_TAG       : int
 *   MPI_STATUS_IGNORE : MPI_Status*
 *   MPI_REQUEST_NULL  : MPI_Request
 *   MPI_INT, MPI_FLOAT, MPI_DOUBLE, MPI_CHAR, MPI_BYTE : MPI_Datatype
 *   MPI_SUM           : MPI_Op
 *   MPI_MAX           : MPI_Op
 *
 * Functions defined:
 *
 * Startup/shutdown:        MPI_Init, MPI_Finalize
 * Information:             MPI_Comm_rank, MPI_Comm_size
 * Nonblocking p2p:         MPI_Isend, MPI_Irecv, MPI_Wait
 * Blocking p2p:            MPI_Send, MPI_Recv, MPI_Sendrecv
 * Other point-to-point:    MPI_Get_count
 * Collective:              MPI_Barrier, MPI_Bcast, MPI_Reduce, MPI_Allreduce
 */

/* stdlib needed for "NULL" */
#include<stdlib.h>
 
/* MPI State (type int) */
#define MPIX_UNINITIALIZED 0
#define MPIX_INITIALIZED 1
#define MPIX_FINALIZED 2
 
/* Communicators (type MPI_Comm): */
#define MPI_COMM_WORLD 1

/* Ranks and tags (type int): */

#define MPI_ANY_SOURCE -1
#define MPI_ANY_TAG -2
#define MPI_PROC_NULL -3
#define MPIX_ALLREDUCE_TAG -1001

/* Operations (type MPI_Op): */
#define MPI_SUM 1
#define MPI_MAX 2

/* Datatypes (type MPI_Datatype): */
#define MPI_INT 1
#define MPI_FLOAT 2
#define MPI_DOUBLE 3
#define MPI_CHAR 4
#define MPI_BYTE 5


/********************************* Types **********************************************/

typedef int MPI_Comm;

typedef int MPI_Op;

typedef int MPI_Datatype;

typedef  struct MPIX_Comm_Record {
  int id;
} MPIX_Comm_Record;

typedef MPIX_Comm_Record *MPI_Request;

typedef struct MPI_Status {
  int MPI_SOURCE;
  int MPI_TAG;
  int SIZE;
} MPI_Status;

/********************************* Constants *******************************************/


/* Communication requests (type MPI_Request): */

MPI_Request MPI_REQUEST_NULL = (MPIX_Comm_Record*)NULL;

/* Status objects (type MPI_Status and MPI_Status*): */

MPI_Status *MPI_STATUS_IGNORE = (MPI_Status*)NULL;


/************************************* State *********************************************/

int MPIX_State = 0; /* 0=pre-init, 1=initalized, 2=finalized */


/*************************************** Functions ***************************************/



int MPI_Init(int *argc, char ***argv) {
  #pragma TASS assert MPIX_State == MPIX_UNINITIALIZED, "Attempt to call MPI_Init when MPI is already initialized";
  MPIX_State = MPIX_INITIALIZED;
  return 0;
}

int MPI_Finalize() {
  #pragma TASS assert MPIX_State == MPIX_INITIALIZED, "Attempt to call MPI_Finalize when MPI is not initialized";
  MPIX_State = MPIX_FINALIZED;
  return 0;
}

int MPI_Comm_rank(MPI_Comm _comm, int *_rank) {
  #pragma TASS assert MPIX_State == MPIX_INITIALIZED, "Attempt to call MPI_Comm_rank when MPI is not initialized";
  /* for now only support MPI_COMM_WORLD */
  #pragma TASS assert _comm == MPI_COMM_WORLD;
  *_rank = PID;
  return 0;
}

int MPI_Comm_size(MPI_Comm _comm, int *_size) {
  #pragma TASS assert MPIX_State == MPIX_INITIALIZED, "Attempt to call MPI_Comm_size when MPI is not initialized";
  /* for now only support MPI_COMM_WORLD */
  #pragma TASS assert _comm == MPI_COMM_WORLD;
  *_size = NPROCS;
  return 0;
}

int MPI_Isend(void* buf, int count, MPI_Datatype datatype, int dest, int tag, 
              MPI_Comm comm, MPI_Request *request) {
  #pragma TASS assert MPIX_State == MPIX_INITIALIZED;
  #pragma TASS assert false, "Not yet implemented.";
}

int MPI_Irecv(void* buf, int count, MPI_Datatype datatype, int source, int tag, 
              MPI_Comm comm, MPI_Request *request) {
  #pragma TASS assert MPIX_State == MPIX_INITIALIZED;
  #pragma TASS assert false, "Not yet implemented.";
}

int MPI_Wait(MPI_Request *request, MPI_Status *status) {
  #pragma TASS assert MPIX_State == MPIX_INITIALIZED;
  #pragma TASS assert false, "Not yet implemented.";
}

/* for now only support MPI_COMM_WORLD, MPI_DOUBLE */
int MPI_Send(void* _buf, int _count, MPI_Datatype _datatype, int _dest, int _tag,
             MPI_Comm _comm) {
  int _result;
  
  #pragma TASS assert MPIX_State == MPIX_INITIALIZED, "Attempt to call MPI_Send when MPI is not initialized";
  if (_dest == MPI_PROC_NULL)
    return 0;
  #pragma TASS assert _comm == MPI_COMM_WORLD;
  if (_datatype == MPI_DOUBLE)
    _result = MPIX_Send_double((double*)_buf, _count, _dest, _tag, _comm);
  else if (_datatype == MPI_INT)
    _result = MPIX_Send_int((int*)_buf, _count, _dest, _tag, _comm);
  else
    #pragma TASS assert false, "Currently only MPI_DOUBLE and MPI_INT supported";
  return _result;
}

int MPIX_Send_double(double* _buf, int _count, int _dest, int _tag, MPI_Comm _comm) {
  double _tmp[_count];
  int _i;
    
  for (_i=0; _i<_count; _i++) _tmp[_i] = _buf[_i];
  send(_tmp, _dest, _tag);
  return 0;
}

int MPIX_Send_int(int* _buf, int _count, int _dest, int _tag, MPI_Comm _comm) {
  int _tmp[_count];
  int _i;
    
  for (_i=0; _i<_count; _i++) _tmp[_i] = _buf[_i];
  send(_tmp, _dest, _tag);
  return 0;
}

/* for now only support MPI_COMM_WORLD, MPI_DOUBLE */
int MPI_Recv(void* _buf, int _count, MPI_Datatype _datatype, int _source, int _tag,
             MPI_Comm _comm, MPI_Status *_status) {
  int _result;
  
  #pragma TASS assert MPIX_State == MPIX_INITIALIZED, "Attempt to call MPI_Recv when MPI is not initialized";
   if (_source == MPI_PROC_NULL)
    return 0;
  #pragma TASS assert _comm == MPI_COMM_WORLD;
  if (_datatype == MPI_DOUBLE)
    _result = MPIX_Recv_double((double*)_buf, _count, _source, _tag, _comm, _status);
  else if (_datatype == MPI_INT)
    _result = MPIX_Recv_int((int*)_buf, _count, _source, _tag, _comm, _status); 
  else
    #pragma TASS assert false;
  return _result;
}

int MPIX_Recv_double(double *_buf, int _count, int _source, int _tag, MPI_Comm _comm,
                     MPI_Status *_status) {
  double _tmp[_count];
  int _i;
  int thesource;
  int thetag;
  int thesize;
  int thecount;
  
  if (_source != MPI_ANY_SOURCE && _tag != MPI_ANY_TAG) {
    recv(_tmp, _source, _tag, thesize);
    thesource = _source;
    thetag = _tag;
  } else if (_source != MPI_ANY_SOURCE && _tag == MPI_ANY_TAG) {
    recv(_tmp, _source, any(thetag), thesize);
    thesource = _source;
  } else if (_source == MPI_ANY_SOURCE && _tag != MPI_ANY_TAG) {
    recv(_tmp, any(thesource), _tag, thesize);
    thetag = _tag;
  } else {
    recv(_tmp, any(thesource), any(thetag), thesize);
  }
  if (_status != MPI_STATUS_IGNORE) {
    _status->MPI_SOURCE = thesource;
    _status->MPI_TAG = thetag;
    _status->SIZE = thesize;
  }
  thecount = thesize/sizeof(double);
  for (_i=0; _i<thecount; _i++) _buf[_i] = _tmp[_i];
  return 0;
}

int MPIX_Recv_int(int *_buf, int _count, int _source, int _tag, MPI_Comm _comm,
                     MPI_Status *_status) {
  int _tmp[_count];
  int _i;
  int thesource;
  int thetag;
  int thesize;
  int thecount;
  
  if (_source != MPI_ANY_SOURCE && _tag != MPI_ANY_TAG) {
    recv(_tmp, _source, _tag, thesize);
    thesource = _source;
    thetag = _tag;
  } else if (_source != MPI_ANY_SOURCE && _tag == MPI_ANY_TAG) {
    recv(_tmp, _source, any(thetag), thesize);
    thesource = _source;
  } else if (_source == MPI_ANY_SOURCE && _tag != MPI_ANY_TAG) {
    recv(_tmp, any(thesource), _tag, thesize);
    thetag = _tag;
  } else {
    recv(_tmp, any(thesource), any(thetag), thesize);
  }
  if (_status != MPI_STATUS_IGNORE) {
    _status->MPI_SOURCE = thesource;
    _status->MPI_TAG = thetag;
    _status->SIZE = thesize;
  }
  thecount = thesize/sizeof(int);
  for (_i=0; _i<thecount; _i++) _buf[_i] = _tmp[_i];
  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) {
  int result;
  
  #pragma TASS assert MPIX_State == MPIX_INITIALIZED, "Attempt to call MPI_Sendrecv when MPI is not initialized";
  MPI_Send(_sendbuf, _sendcount, _sendtype, _dest, _sendtag, _comm);
  result = MPI_Recv(_recvbuf, _recvcount, _recvtype, _source, _recvtag, _comm, _status);
  return result;
}

int MPI_Get_count(MPI_Status *status, MPI_Datatype datatype, int *count) {
  #pragma TASS assert MPIX_State == MPIX_INITIALIZED, "Attempt to call MPI_Get_count when MPI is not initialized";
  
  if (datatype == MPI_DOUBLE)
    *count = status->SIZE/sizeof(double);
  else if (datatype == MPI_FLOAT)
    *count = status->SIZE/sizeof(float);
  else if (datatype == MPI_INT)
    *count = status->SIZE/sizeof(int);
  else
    #pragma TASS assert false;
  return 0;
}

int MPI_Barrier(MPI_Comm comm) {
  #pragma TASS assert MPIX_State == MPIX_INITIALIZED, "Attempt to call MPI_Barrier when MPI is not initialized";
}

int MPI_Bcast(void* buffer, int count, MPI_Datatype datatype, int root, MPI_Comm comm ) {
  #pragma TASS assert MPIX_State == MPIX_INITIALIZED, "Attempt to call MPI_Bcast when MPI is not initialized";
}

int MPI_Reduce(void* sendbuf, void* recvbuf, int count, MPI_Datatype datatype,
               MPI_Op op, int root, MPI_Comm comm) {
  #pragma TASS assert MPIX_State == MPIX_INITIALIZED, "Attempt to call MPI_Reduce when MPI is not initialized";
}

/* For now: just doubles and sum */
int MPI_Allreduce(void* _sendbuf, void* _recvbuf, int _count, MPI_Datatype _datatype,
                  MPI_Op _op, MPI_Comm _comm) {
                  
  int _i;
  int _j;
  double _tempBuffer[_count];
  
  #pragma TASS assert MPIX_State == MPIX_INITIALIZED, "Attempt to call MPI_Allreduce when MPI is not initialized";
  #pragma TASS assert _op == MPI_SUM;
  #pragma TASS assert _datatype == MPI_DOUBLE;
  if (PID == 0) {
    for (_j=0; _j<_count; _j++) _tempBuffer[_j] = ((double*)_sendbuf)[_j];
  	for (_i=1; _i<NPROCS; _i++) {
  	  MPI_Recv(_recvbuf, _count, _datatype, _i, MPIX_ALLREDUCE_TAG, 
  	           _comm, MPI_STATUS_IGNORE);
  	  for(_j=0; _j<_count; _j++) _tempBuffer[_j] += ((double*)_recvbuf)[_j];
  	}
  	for (_j=0; _j<_count; _j++) ((double*)_recvbuf)[_j]=_tempBuffer[_j];
  	for (_i=1; _i<NPROCS; _i++) {
  	  MPI_Send(_recvbuf, _count, _datatype, _i, MPIX_ALLREDUCE_TAG, _comm);
  	}
  } else {
    MPI_Send(_sendbuf, _count, _datatype, 0, MPIX_ALLREDUCE_TAG, _comm);
    MPI_Recv(_recvbuf, _count, _datatype, 0, MPIX_ALLREDUCE_TAG, _comm,
             MPI_STATUS_IGNORE);
  }
  return 0;
}
 