/****************************************************************************
**  SIONLIB     http://www.fz-juelich.de/jsc/sionlib                       **
*****************************************************************************
**  Copyright (c) 2008-2009                                                **
**  Forschungszentrum Juelich, Juelich Supercomputing Centre               **
**                                                                         **
**  See the file COPYRIGHT in the package base directory for details       **
****************************************************************************/
/*
 * \file sion_mpi.c
 *
 * \brief MPI Impl
 */

#include <stdlib.h>
#include <stdio.h>
#include <stdarg.h>
#include <string.h>
#include <time.h>

#define USE_PMPIno
#ifdef USE_PMPI
#define MPI_Comm_rank      PMPI_Comm_rank
#define MPI_Comm_size      PMPI_Comm_size	  
#define MPI_Gather	   PMPI_Gather	  
#define MPI_Scatter	   PMPI_Scatter	  
#define MPI_Bcast	   PMPI_Bcast	  
#define MPI_Barrier	   PMPI_Barrier	  
#define MPI_Comm_split     PMPI_Comm_split    
#define MPI_Send           PMPI_Send
#define MPI_Recv           PMPI_Recv
#endif

#if defined(_BGL)
#include <rts.h>
#ifndef __USE_FILE_OFFSET64
#define __USE_FILE_OFFSET64
#endif
#endif

#include <sys/time.h>

#include <sys/types.h>
#include <fcntl.h>

#include <unistd.h>

#include "sion.h"
#include "sion_fd.h"
#include "printts.h"

#ifdef SION_MPI

int _gather_cb_mpi(void *indata, void *outdata, void *commvp, int dtype)
{
  int       rc;
  int       size, rank;
  MPI_Comm *commp = (MPI_Comm *) commvp;

  MPI_Comm_rank(*commp, &rank);
  MPI_Comm_size(*commp, &size);
  /*   fprintf(stderr,"in gather_cb_mpi comm=%d %d of %d\n",*commp,rank,size); */

  /* Dummy selectes the type of gather */
  switch (dtype) {
  case _SION_INT32:
    rc = MPI_Gather((sion_int32 *) indata, 1, SION_MPI_INT32, (sion_int32 *) outdata, 1, SION_MPI_INT32, 0, *commp);
    break;
  case _SION_INT64:
    rc = MPI_Gather((sion_int64 *) indata, 1, SION_MPI_INT64, (sion_int64 *) outdata, 1, SION_MPI_INT64, 0, *commp);
    break;
  case _SION_CHAR:
    rc = MPI_Gather((char *) indata, 1, MPI_CHAR, (char *) outdata, 1, MPI_CHAR, 0, *commp);
    break;
  default:
    rc = MPI_Gather((sion_int64 *) indata, 1, SION_MPI_INT64, (sion_int64 *) outdata, 1, SION_MPI_INT64, 0, *commp);
    break;
  }


  return rc;
}

int _scatter_cb_mpi(void *indata, void *outdata, void *commvp, int dtype)
{
  int       rc;

  MPI_Comm *commp = (MPI_Comm *) commvp;

  switch (dtype) {
  case _SION_INT32:
    rc = MPI_Scatter((sion_int32 *) indata, 1, SION_MPI_INT32, (sion_int32 *) outdata, 1, SION_MPI_INT32, 0, *commp);
    break;
  case _SION_INT64:
    rc = MPI_Scatter((sion_int64 *) indata, 1, SION_MPI_INT64, (sion_int64 *) outdata, 1, SION_MPI_INT64, 0, *commp);
    break;
  case _SION_CHAR:
    rc = MPI_Scatter((char *) indata, 1, MPI_CHAR, (char *) outdata, 1, MPI_CHAR, 0, *commp);
    break;
  default:
    rc = MPI_Scatter((sion_int64 *) indata, 1, SION_MPI_INT64, (sion_int64 *) outdata, 1, SION_MPI_INT64, 0, *commp);
    break;
  }

  return rc;
}

int _bcast_cb_mpi(void *data, void *commvp, int dtype)
{
  int       rc;
  MPI_Comm *commp = (MPI_Comm *) commvp;
  /*  rc=MPI_Bcast(data,1,SION_MPI_INT64,0,*commp); */
  /* dtype = *((int*) dummy); */
  switch (dtype) {
  case _SION_INT32:
    rc = MPI_Bcast((sion_int32 *) data, 1, SION_MPI_INT32, 0, *commp);
    break;
  case _SION_INT64:
    rc = MPI_Bcast((sion_int64 *) data, 1, SION_MPI_INT64, 0, *commp);
    break;
  case _SION_CHAR:
    rc = MPI_Bcast((char *) data, 1, MPI_CHAR, 0, *commp);
    break;
  default:
    rc = MPI_Bcast((sion_int64 *) data, 1, SION_MPI_INT64, 0, *commp);
    break;
  }
  return rc;
}

int _barrier_cb_mpi(void *commvp, void *dummy)
{
  int       rc;
  MPI_Comm *commp = (MPI_Comm *) commvp;
  rc = MPI_Barrier(*commp);
  return rc;
}

int sion_paropen_mpi(char       *fname,
                     char       *file_mode,
                     int        *numFiles,
                     MPI_Comm    gComm,
                     MPI_Comm   *lComm, 
		     sion_int64 *chunksize, 
		     sion_int32 *fsblksize, 
		     int        *globalrank, 
		     FILE      **fileptr, 
		     char      **newfname)
{
  int       rc, sid = -1;
  int       filenumber, gtasks, gRank, ltasks, lRank;
  char     *nfname,*lprefix;
  
  MPI_Comm_size(gComm, &gtasks);
  MPI_Comm_rank(gComm, &gRank);

  /* check parameters */
  if (lComm == NULL) {
    return(_sion_errorprint_mpi(-1,_SION_ERROR_RETURN,"sion_paropen_mpi: No lComm variable given"));
  }
  if (numFiles == NULL) {
    return(_sion_errorprint_mpi(-1,_SION_ERROR_RETURN,"sion_paropen_mpi: No numFiles variable given"));
  }

  lprefix = calloc(SION_FILENAME_LENGTH,1);
  if (lprefix == NULL) {
    return(_sion_errorprint_mpi(-1,_SION_ERROR_RETURN,"sion_paropen_mpi: cannot allocate temporary memory of size %lu (lprefix), aborting ...\n", (unsigned long) SION_FILENAME_LENGTH));
  }

  if (!strcmp(file_mode, "bw")) {
    /* file mode WRITE */

    if (*numFiles <= 0) {
      /* lComm contains local communicator */
      rc = _sion_get_color_from_splitted_comm_mpi(gComm, *lComm, numFiles, &filenumber);
      DPRINTFP((1, "sion_paropen_mpi", gRank, "%d local communicators found\n", *numFiles));

    }
    else {
      /* number of files is given */
      rc = _sion_split_comm_mpi(*numFiles, gComm, lComm, &filenumber);
      if(!rc) return(_sion_errorprint_mpi(rc,_SION_ERROR_RETURN,"sion_paropen_mpi: error in _sion_split_comm_mpi"));
      DPRINTFP((1, "sion_paropen_mpi", gRank, "Global communicator divided in %d local communicators\n", *numFiles));
    }

    /* overwrite globalrank set by user, necessary for multi-file support */
    *globalrank = gRank;

    strcpy(lprefix, fname);

  }
  else {
    /* file mode READ */

    /* get number of files and filenumber */
    rc = _sion_get_filenumber_from_files_mpi(fname, gComm, numFiles, &filenumber, &lRank);
    if(!rc) return(_sion_errorprint_mpi(rc,_SION_ERROR_RETURN,"sion_paropen_mpi: error in _sion_get_filenumber_from_files_mpi"));
  
    /* split communicator */
    MPI_Comm_split(gComm, filenumber, lRank, lComm);

  }

  /* set up parameters of call to generic open  */
  MPI_Comm_size(*lComm, &ltasks);
  MPI_Comm_rank(*lComm, &lRank);
  nfname=_sion_get_multi_filename(fname,filenumber);
  if(newfname!=NULL) {
    *newfname=nfname;
  }

  DPRINTFP((1, "sion_paropen_mpi", gRank, "enter parallel open of %d files (current name %s) in %s mode\n", *numFiles, nfname, file_mode));
  sid =
    _sion_paropen_multi_generic(nfname, file_mode, lprefix, numFiles, &filenumber, chunksize, fsblksize, lRank, ltasks, globalrank, SION_DIFFSIZE,
                                fileptr, _gather_cb_mpi, _scatter_cb_mpi, _bcast_cb_mpi, _barrier_cb_mpi, lComm, &gComm);
  DPRINTFP((1, "sion_paropen_mpi", gRank, "leave parallel open of %d files in %s mode #tasks=%d sid=%d\n", *numFiles, file_mode, ltasks, sid));

  {                             /* test sid */
    _sion_filedesc *sion_filedesc;
    if ((sid<0) || (_sion_vcdtype(sid) != SION_FILEDESCRIPTOR) || !(sion_filedesc = _sion_vcdtovcon(sid))) {
      return(_sion_errorprint_mpi(-1,_SION_ERROR_RETURN,"sion_paropen_mpi: invalid sion_filedesc %d", sid));
    }
  }
  return (sid);
}


int sion_parclose_mpi(int sid)
{
  int       grank, lrank, mapping_size, rc, receiver;
  _sion_filedesc *sion_filedesc;
  MPI_Comm  gComm;
  sion_int32 lpos[2];
  sion_int32 *mapping;
  MPI_Status status;

  if ( (sid<0) || (_sion_vcdtype(sid) != SION_FILEDESCRIPTOR) || !(sion_filedesc = _sion_vcdtovcon(sid))) {
    return(_sion_errorprint_mpi(-1,_SION_ERROR_RETURN,"sion_parclose_mpi: invalid sion_filedesc %d", sid));
  }

  MPI_Comm_rank(sion_filedesc->gComm, &grank);
  MPI_Comm_rank(sion_filedesc->lComm, &lrank);

  DPRINTFP((1, "sion_parclose_mpi", grank, "enter parallel close of sid %d\n", sid));
  DPRINTFP((1, "sion_parclose_mpi", grank, "closing %d file(s)\n", sion_filedesc->nfiles));

  gComm = sion_filedesc->gComm;

  if ((sion_filedesc->mode == SION_FILEMODE_WRITE) && (sion_filedesc->nfiles > 1)) {
    /* collect mapping to files to task 0 */

    /* mapping data will be collected by master of first physical file */
    if((sion_filedesc->filenumber==0) && (lrank==0)) {
      /* allocate data */
      MPI_Comm_size(sion_filedesc->gComm, &mapping_size);
      mapping = (sion_int32 *) malloc(mapping_size * 2 * sizeof(sion_int32));
      if (mapping == NULL) {
	return(_sion_errorprint_mpi(-1,_SION_ERROR_ABORT,"sion_parclose_mpi: Cannot allocate memory for mapping"));
      }
    } else {
      mapping = NULL;
      mapping_size = 0;
    }

    /* send global rank of master of first file to grank 0 */
    if((sion_filedesc->filenumber==0) && (lrank==0)) {
      if(grank!=0) {
	MPI_Send(&grank, 1, MPI_INT, 0, 1330, gComm);
      }
    }

    /* receive global rank of master of first file on grank 0 */
    if(grank==0) {
      if((sion_filedesc->filenumber==0) && (lrank==0)) {
	receiver=0;
      } else {
	MPI_Recv(&receiver, 1, MPI_INT, MPI_ANY_SOURCE, 1330, gComm, &status);
      }
    }
    
    MPI_Bcast(&receiver, 1, MPI_INT, 0, gComm);

    lpos[0] = sion_filedesc->filenumber;
    lpos[1] = lrank;
    MPI_Gather(lpos, 2, SION_MPI_INT32, mapping, 2, SION_MPI_INT32, receiver, sion_filedesc->gComm);
  }
  else {
    mapping = NULL;
    mapping_size = 0;
  }

  rc = _sion_parclose_generic(sid, sion_filedesc->rank, sion_filedesc->ntasks, mapping_size, mapping,
                              _gather_cb_mpi, _scatter_cb_mpi, _bcast_cb_mpi, _barrier_cb_mpi, &(sion_filedesc->lComm), NULL);

  /*     if((sion_filedesc->nfiles > 1)||(sion_filedesc->nfiles == 0)) */
  MPI_Barrier(gComm);

  DPRINTFP((1, "sion_parclose_mpi", grank, "leave parallel close of sid %d\n", sid));

  return (rc);
}

int _sion_errorprint_mpi(int rc, int level, const char *format, ...)
{
  int       rank=-1;
  va_list ap;
  int rankselect=-1;
  const char *t;

#ifdef SION_MPI
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
#endif

  
  t = getenv("SION_ERROR_MSG_RANK");
  if (t) rankselect = atoi(t);

  switch (level) {
    case _SION_ERROR_RETURN: 
      if((rankselect<0) || (rankselect==rank)) {
	fprintf(stderr,"SION_ERROR_RETURN on rank %d, rc=%d: ",rank,rc);
	va_start(ap, format);
	vfprintf(stderr, format, ap);
	va_end(ap);
	fprintf(stderr,"\n");
      }
      return (rc);
      break;
    case _SION_ERROR_ABORT: 
      if((rankselect<0) || (rankselect==rank)) {
	fprintf(stderr,"SION_ERROR_ABORT on rank %d, rc=%d: ",rank,rc);
	va_start(ap, format);
	vfprintf(stderr, format, ap);
	va_end(ap);
	fprintf(stderr,"\n");
      }
      exit (rc);
      break;
  default:
    if((rankselect<0) || (rankselect==rank)) {
      
      fprintf(stderr,"SION_ERROR_UNKNOWN on rank %d, rc=%d: ",rank,rc);
      va_start(ap, format);
      vfprintf(stderr, format, ap);
      va_end(ap);
      fprintf(stderr,"\n");
    }
    return (rc);
  }
  return (rc);
}



/* Help Functions */

/*!\brief Splits a Communicator in numfiles different communicators
 *
 * @param *numComm              number of sub-communicators
 * @param *gComm                global communicator
 * @param *lComm                local communicator
 * @param *color                color of local communicator
 * 
 */

int _sion_split_comm_mpi(int numComm, MPI_Comm gComm, MPI_Comm * lComm, int *color)
{
  int       gtasks, gRank;
  int       rc = -1;
  int       task_per_comm;
  MPI_Comm_size(gComm, &gtasks);
  MPI_Comm_rank(gComm, &gRank);

  if (gtasks < numComm) {
    return(_sion_errorprint_mpi(-1,_SION_ERROR_RETURN,"_sion_split_comm_mpi: Number of files(%d) cannot be bigger the the number of tasks(%d), aborting...\n", 
			      numComm, gtasks));
  }

  task_per_comm = gtasks / numComm;

  /* remaining tasks are added to last communicator */
  if (gRank >= (numComm * task_per_comm)) {
    *color = numComm - 1;
  }
  else {
    *color = gRank / task_per_comm;
  }

  /* *color = gRank % (numComm); */

  rc=(MPI_Comm_split(gComm, *color, gRank, lComm)==MPI_SUCCESS);

  DPRINTFP((1, "_sion_split_comm_mpi", gRank, "Global communicator divided in %d local communicators\n", numComm));

  return (rc);
}

int _sion_get_color_from_splitted_comm_mpi(MPI_Comm gComm, MPI_Comm lComm, int *numComm, int *color)
{
  int       gSize, gRank, lSize, lRank;
  int      *allRanks, *allSizes, i, ntasks;
  int       ncomms;
  int       rc = -1;
  MPI_Status status;

  MPI_Comm_size(gComm, &gSize);
  MPI_Comm_rank(gComm, &gRank);

  if (lComm != MPI_COMM_NULL) {
    MPI_Comm_size(lComm, &lSize);
    MPI_Comm_rank(lComm, &lRank);
  }
  else {
    lSize = gSize;
    lRank = gRank;
  }

  DPRINTFP((32, "_sion_get_color_from_splitted_comm_mpi", gRank, "lRank: %d\n", lRank));

  if (gRank == 0) {

    allRanks = (int *) malloc(gSize * sizeof(int));
    allSizes = (int *) malloc(gSize * sizeof(int));
    if ((allRanks == NULL) || (allSizes == NULL)) {
      return(_sion_errorprint_mpi(-1,_SION_ERROR_RETURN,"_sion_get_color_from_splitted_comm_mpi: Cannot allocate temp arrays of size %lu\n", (unsigned long) gSize * sizeof(int)));
    }
  }

  /*  Gather all local ranks to task 0 for counting.. */
  MPI_Gather(&lRank, 1, MPI_INT, allRanks, 1, MPI_INT, 0, gComm);
  MPI_Gather(&lSize, 1, MPI_INT, allSizes, 1, MPI_INT, 0, gComm);

  ncomms = 0;
  ntasks = 0;
  if (gRank == 0) {
    for (i = 0; i < gSize; i++) {
      if (allRanks[i] == 0) {
        /*  Use the current number of the communicator for the file suffix!! */
        if (i != 0) {
          MPI_Send(&ncomms, 1, MPI_INT, i, i, gComm);
          DPRINTFP((32, "sion_paropen_comms_mpi", gRank, "Sent ncomms=%05d to %d with TAG %d\n", ncomms, i, i));
        }
        else {
          /* it's my self */
          *color = ncomms;
        }
        ncomms++;
        ntasks += allSizes[i];
      }
    }
  }
  /*  Recv the num of communicator from global 0 =>used for the suffix */
  if ((lRank == 0) && (gRank != 0)) {
    MPI_Recv(color, 1, MPI_INT, 0, gRank, gComm, &status);
    DPRINTFP((32, "sion_paropen_comms_mpi", gRank, "Received ncomms=%05d from %d with TAG %d\n", *color, status.MPI_SOURCE, status.MPI_TAG));

  }

  MPI_Bcast(&ncomms, 1, MPI_INT, 0, gComm);
  MPI_Bcast(&ntasks, 1, MPI_INT, 0, gComm);

  *numComm = ncomms;

  DPRINTFP((32, "_sion_get_color_from_splitted_comm_mpi", gRank, "#Comms=%05d #Tasks=%05d #TotalTasks=%05d\n", ncomms, ntasks, gSize));

  if (lComm != MPI_COMM_NULL) {
    MPI_Bcast(color, 1, MPI_INT, 0, lComm);
  }

  return (rc);
}

int _sion_get_numfiles_from_file_mpi(char *fname) {

  int       rank, sid;
  FILE     *fileptr;
  int       numfiles = -1;
  sion_int64 chunksize;
  sion_int32 fsblksize;

  rank = 0;
  sid = sion_open_rank(fname, "br", &chunksize, &fsblksize, &rank, &fileptr);

  numfiles = sion_get_number_of_files(sid);

  sion_close(sid);

  DPRINTFP((1, "_sion_get_numfiles_from_file_mpi", -1, "sion file %s has %d files\n", fname, numfiles));

  return (numfiles);
}

int _sion_get_filenumber_from_files_mpi(char *fname, MPI_Comm gComm, int *numfiles, int *filenumber, int *lRank) {

  int       sid, gSize, gRank, ntasks, nfiles;
  int       rc=1;
  FILE     *fileptr;
  sion_int64 *chunksizes;
  sion_int32 fsblksize;
  int      *globalranks;
  int         mapping_size;
  sion_int32 *mapping;
  sion_int32 lpos[2];

  MPI_Comm_size(gComm, &gSize);
  MPI_Comm_rank(gComm, &gRank);

  if(gRank == 0) {
    /* open and get mapping of sion file */
    DPRINTFP((1, "_sion_get_filenumber_from_files_mpi", gRank, "before open\n"));
    chunksizes=NULL;globalranks=NULL;
    sid = sion_open(fname, "br", &ntasks, &nfiles, &chunksizes, &fsblksize, &globalranks, &fileptr);
    DPRINTFP((1, "_sion_get_filenumber_from_files_mpi", gRank, "after open\n"));
    rc=sion_get_mapping(sid,&mapping_size,&mapping,numfiles);
    DPRINTFP((1, "_sion_get_filenumber_from_files_mpi", gRank, "sion file %d files\n", *numfiles));
  }
  
  /* each task has to know if more than file was used in sion file */
  MPI_Bcast(numfiles, 1, MPI_INT, 0, gComm);

  if((gRank == 0) && (*numfiles>1)) {
    if(mapping_size!=gSize) {
      return(_sion_errorprint_mpi(-1,_SION_ERROR_RETURN,"_sion_get_filenumber_from_files_mpi: Incorrect sum of ntasks of files %d <> %d\n", mapping_size, gSize));
    }
  }

  if(*numfiles>1) {
    DPRINTFP((1, "_sion_get_filenumber_from_files_mpi", gRank, "before scatter\n"));
    MPI_Scatter(mapping, 2, SION_MPI_INT32, lpos, 2, SION_MPI_INT32, 0, gComm);
    *filenumber=lpos[0];
    *lRank     =lpos[1];
    DPRINTFP((1, "_sion_get_filenumber_from_files_mpi", gRank, "after scatter\n"));
  } else {
    *filenumber=0;
    *lRank     =gRank;
    DPRINTFP((1, "_sion_get_filenumber_from_files_mpi", gRank, "only one file -> filenumber=%d lRank=%d\n",*filenumber,*lRank));
  }

  if(gRank == 0) {
    /* frees also mapping vector */
    sion_close(sid);
  }

  return(rc);
}

/* wrapper for old SIONlib MPI function */
int sion_paropen_comms_mpi(char *fname,
                           char *file_mode,
                           int *numFiles,
                           MPI_Comm gComm,
                           MPI_Comm lComm, sion_int64 *chunksize, sion_int32 *fsblksize, int *globalrank, FILE **fileptr, char *newfname)
{
  *numFiles=0;
  return(sion_paropen_mpi(fname,file_mode,numFiles,gComm,&lComm,chunksize,fsblksize,globalrank,fileptr,&newfname));
}

int sion_paropen_multi_mpi(char *fname,
                           char *file_mode,
                           int *numFiles,
                           MPI_Comm gComm,
                           MPI_Comm *lComm, sion_int64 *chunksize, sion_int32 *fsblksize, int *globalrank, FILE **fileptr, char *newfname)
{
  return(sion_paropen_mpi(fname,file_mode,numFiles,gComm,lComm,chunksize,fsblksize,globalrank,fileptr,&newfname));
}

/* end of ifdef MPI */
#endif
