#include "jacobi.h"

int migHandle;

CtvDeclare(jacobi *, jacobiPtr);
CtvDeclare(int, numMigrateCalls);
extern "C" main_(void);
static Array1D *jacobiArray;

jacobi::jacobi(ArrayElementCreateMessage *msg) : TempoArray(msg)
{
  jacobiArray = thisArray;
  delete msg;
}

jacobi::jacobi(ArrayElementMigrateMessage *msg) : TempoArray(msg)
{
  jacobiArray = thisArray;
  thread_id = CthUnpackThread(msg->packData);
  CthAwaken(thread_id);
  finishMigration(); 
  delete msg;
}

int
jacobi::packsize(void)
{
  return CthPackBufSize(thread_id);
}

void
jacobi::pack(void *buf)
{
  TempoMessage *msg;
  int itags[2];
  itags[0] = TEMPO_ANY; itags[1] = TEMPO_ANY;
  // resend pending messages in table
  while((msg=(TempoMessage *)CmmGet(tempoMessages, 2, itags, 0))) {
    ckTempoSendElem(msg->tag1, msg->tag2, msg->data, msg->length, thisIndex);
  }
  CthPackThread(thread_id, buf);
}

void
jacobi::run(void)
{
  static int initCtv = 0;
  int myIdx = getIndex();
  jacobi *myThis;

  if(!initCtv) {
    CtvInitialize(jacobi *, jacobiPtr);
    CtvInitialize(int, numMigrateCalls);
    initCtv = 1;
  }
  CtvAccess(jacobiPtr) = this;
  CtvAccess(numMigrateCalls) = 0;

  main_();
  myThis = (jacobi*) jacobiArray->getElement(myIdx);
  CkPrintf("[%d] main_ finished\n", myThis->getIndex());

  itersDone();

  /*
  myThis->ckTempoBarrier();
  if(myThis->getIndex()==0)
    CmiAbort("");
  else
    CthSuspend();
  */
  CthSuspend();
}

extern "C" void migrate_ready_(void)
{
  CtvAccess(numMigrateCalls)++;
  // migrate to next processor every 2 iterations
  if(CtvAccess(numMigrateCalls)%2 == 0) {
    int index = CtvAccess(jacobiPtr)->getIndex();
    int where = (CkMyPe()+1) % CkNumPes();
    if(where == CkMyPe())
      return;
    CProxy_migrator pmg(migHandle);
    pmg.migrateElement(new MigrateInfo(CtvAccess(jacobiPtr), where), CkMyPe());
    CkPrintf("[%d] Migrating from %d to %d\n", index, CkMyPe(), where);
    CthSuspend();
    CtvAccess(jacobiPtr) = (jacobi*) jacobiArray->getElement(index);
    CkPrintf("[%d] Migrated to %d\n", index, CkMyPe());
  }
}

extern "C" void mpi_init_(int *ierr)
{
  *ierr = 0;
}

extern "C" void mpi_comm_rank_(int *comm, int *rank, int *ierr)
{
  *ierr = 0;
  *rank = CtvAccess(jacobiPtr)->getIndex();
}

extern "C" void mpi_comm_size_(int *comm, int *size, int *ierr)
{
  *ierr = 0;
  *size = CtvAccess(jacobiPtr)->getSize();
}

extern "C" void mpi_finalize_(int *ierr)
{
  *ierr = 0;
}

// these values have to match values in mpi.f90

#define CMPI_DOUBLE_PRECISION 0
#define CMPI_INTEGER 1
#define CMPI_REAL 2
#define CMPI_COMPLEX 3
#define CMPI_LOGICAL 4
#define CMPI_CHARACTER 5
#define CMPI_BYTE 6
#define CMPI_PACKED 7

#define CMPI_MAX 1
#define CMPI_MIN 2
#define CMPI_SUM 3
#define CMPI_PROD 4

static int typesize(int type, int count)
{
  switch(type) {
    case CMPI_DOUBLE_PRECISION : return count*sizeof(double);
    case CMPI_INTEGER : return count*sizeof(int);
    case CMPI_REAL : return count*sizeof(float);
    case CMPI_COMPLEX: return 2*count*sizeof(double);
    case CMPI_LOGICAL: return 2*count*sizeof(int);
    case CMPI_CHARACTER:
    case CMPI_BYTE:
    case CMPI_PACKED:
    default:
      return 2*count;
  }
}

extern "C" void mpi_send_(void *msg, int *count, int *type, int *dest, 
                          int *tag, int *comm, int *ierr)
{
  int size = typesize(*type, *count);
  jacobi *ptr = CtvAccess(jacobiPtr);
  CkPrintf("[%d] sending %d bytes to %d tagged %d\n", ptr->getIndex(), 
           size, *dest, *tag);
  ptr->ckTempoSendElem(*tag, msg, size, *dest);
  *ierr = 0;
}

extern "C" void mpi_recv_(void *msg, int *count, int *type, int *src,
                          int *tag, int *comm, int *status, int *ierr)
{
  int size = typesize(*type, *count);
  jacobi *ptr = CtvAccess(jacobiPtr);
  CkPrintf("[%d] waits for %d bytes tagged %d\n", ptr->getIndex(), size, *tag);
  ptr->ckTempoRecv(*tag, msg, size);
  CkPrintf("[%d] received %d bytes tagged %d\n", ptr->getIndex(), size, *tag);
  *ierr = 0;
}

extern "C" void mpi_barrier_(int *comm, int *ierr)
{
  jacobi *ptr = CtvAccess(jacobiPtr);
  ptr->ckTempoBarrier();
  CkPrintf("[%d] Barrier finished\n", ptr->getIndex());
  *ierr = 0;
}

#define CMPI_BCAST_TAG 999

extern "C" void mpi_bcast_(void *buf, int *count, int *type, int *root, 
                           int *comm, int *ierr)
{
  int size = typesize(*type, *count);
  jacobi *ptr = CtvAccess(jacobiPtr);
  ptr->ckTempoBcast(*root==ptr->getIndex(), CMPI_BCAST_TAG, buf, size);
  CkPrintf("[%d] Broadcast finished\n", ptr->getIndex());
  *ierr = 0;
}

static void optype(int inop, int intype, int *outop, int *outtype)
{
  switch(inop) {
    case CMPI_MAX : *outop = TEMPO_MAX; break;
    case CMPI_MIN : *outop = TEMPO_MIN; break;
    case CMPI_SUM : *outop = TEMPO_SUM; break;
    case CMPI_PROD : *outop = TEMPO_PROD; break;
    default:
      ckerr << "Op " << inop << " not supported." << endl;
      CmiAbort("exiting");
  }
  switch(intype) {
    case CMPI_REAL : *outtype = TEMPO_FLOAT; break;
    case CMPI_INTEGER : *outtype = TEMPO_INT; break;
    case CMPI_DOUBLE_PRECISION : *outtype = TEMPO_DOUBLE; break;
    default:
      ckerr << "Type " << intype << " not supported." << endl;
      CmiAbort("exiting");
  }
}

extern "C" void mpi_reduce_(void *inbuf, void *outbuf, int *count, int *type,
                            int *op, int *root, int *comm, int *ierr)
{
  int myop, mytype;
  optype(*op, *type, &myop, &mytype);
  jacobi *ptr = CtvAccess(jacobiPtr);
  ptr->ckTempoReduce(*root,myop,inbuf,outbuf,*count,mytype);
  CkPrintf("[%d] reduction finished\n", ptr->getIndex());
}

extern "C" void mpi_allreduce_(void *inbuf,void *outbuf,int *count,int *type,
                               int *op, int *comm, int *ierr)
{
  int myop, mytype;
  optype(*op, *type, &myop, &mytype);
  jacobi *ptr = CtvAccess(jacobiPtr);
  ptr->ckTempoAllReduce(myop,inbuf,outbuf,*count,mytype);
  CkPrintf("[%d] Allreduce finished\n", ptr->getIndex());
}

#include "jacobi.def.h"
