#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

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

#include "pio_comm.h"
#include "cdi.h"
#include "dmemory.h"
#include "pio_util.h"

#include "cdipio.h"

typedef struct
{
  int IOMode;
  int isProcIO;

  MPI_Comm commGlob;
  int sizeGlob;
  int rankGlob;
  int root;

  MPI_Comm commPio; /* intracommunicator of PIO servers */
  int sizePio;
  int rankPio;

  MPI_Comm commColl; /* intracommunicator of collector
                      * subset of PIO servers only */
  int sizeColl;
  int rankColl;

  MPI_Comm commCalc;
  int rankCalc, sizeCalc;

  /**
   * communicator connecting group of clients and group of collectors
   * notice: it is always implicit that the group of collectors is no
   * larger than the group of clients
   */
  MPI_Comm pioInterComm;
  /**
   * group containing only the active remote
   * processes, i.e.
   * - on collector processes contains all processes that are a client of
   *   the server
   * - on client processes contains only the server process used for
   *   data/RPC communication with this client process
   */
  MPI_Group remoteGroup;
  /**
   * Communicator containing both clients and collectors, needed since
   * MPI_Win_create only accepts intracommunicators.
   * Constructed such that collectors occupy the top ranks in the same
   * sequence they have on their own communicator.
   */
  MPI_Comm collClientIntraComm;
} pioInfo_t;

static pioInfo_t *info = NULL;

static void setupCollectorInterComm(pioInfo_t *p, int isCollector, MPI_Comm localGroupComm, int globRankOfGroupLeader[2]);

static void
pioInfoInit(pioInfo_t *p, MPI_Comm commGlob, int clientServerRole)
{
  xassert(info != NULL && commGlob != MPI_COMM_NULL);
  p->IOMode = -1;

  xmpi(MPI_Comm_dup(commGlob, &p->commGlob));
  xmpi(MPI_Comm_size(p->commGlob, &p->sizeGlob));
  xmpi(MPI_Comm_rank(p->commGlob, &p->rankGlob));
  p->root = 0;

  p->commColl = MPI_COMM_NULL;
  p->rankColl = -1;

  int isProcIO = (clientServerRole != PIO_ROLE_CLIENT), isCollector = (clientServerRole & PIO_ROLE_COLLECTOR);
  p->isProcIO = isProcIO;

  {
    MPI_Comm temp;
    /* key is 0 on all clients
     *        0 on all servers that are also collectors
     *        1 on all servers that are not collectors
     */
    int key = isProcIO - isCollector;
    xmpi(MPI_Comm_split(p->commGlob, isProcIO, key, &temp));
    enum
    {
      PIO = 0,
      Client = 1,
      numGroups
    };
    int globRankOfGroupLeader[numGroups] = { 0, 0 };
    if (isProcIO)
      {
        p->commPio = temp;
        xmpi(MPI_Comm_rank(temp, &p->rankPio));
        p->commCalc = MPI_COMM_NULL;
        p->rankCalc = -1;
        if (p->rankPio == 0) globRankOfGroupLeader[PIO] = p->rankGlob;
      }
    else
      {
        p->commPio = MPI_COMM_NULL;
        p->rankPio = -1;
        p->commCalc = temp;
        xmpi(MPI_Comm_rank(temp, &p->rankCalc));
        xmpi(MPI_Comm_size(temp, &p->sizeCalc));
        if (p->rankCalc == 0) globRankOfGroupLeader[Client] = p->rankGlob;
      }
    int nProcsIO;
    {
      enum
      {
        numGlobSums = 4
      };
      int iBuf[numGlobSums]
          = { isProcIO, (clientServerRole & PIO_ROLE_COLLECTOR), globRankOfGroupLeader[0], globRankOfGroupLeader[1] };
      xmpi(MPI_Allreduce(MPI_IN_PLACE, iBuf, numGlobSums, MPI_INT, MPI_SUM, commGlob));
      p->sizePio = nProcsIO = iBuf[0];
      p->sizeCalc = p->sizeGlob - nProcsIO;
      p->sizeColl = iBuf[1];
      globRankOfGroupLeader[0] = iBuf[2];
      globRankOfGroupLeader[1] = iBuf[3];
    }
    if (isProcIO)
      {
        MPI_Group pioGroup, collGroup;
        int collRange[1][3] = { { 0, p->sizeColl - 1, 1 } };
        xmpi(MPI_Comm_group(p->commPio, &pioGroup));
        xmpi(MPI_Group_range_incl(pioGroup, 1, collRange, &collGroup));
        xmpi(MPI_Comm_create(p->commPio, collGroup, &temp));
        if (isCollector) xmpi(MPI_Comm_rank(temp, &p->rankColl));
        xmpi(MPI_Group_free(&pioGroup));
        xmpi(MPI_Group_free(&collGroup));
        p->commColl = temp;
      }
    if (isCollector || !isProcIO)
      setupCollectorInterComm(p, isCollector, temp, globRankOfGroupLeader);
    else /* isProcIO && !isCollector */
      {
        p->pioInterComm = MPI_COMM_NULL;
        p->collClientIntraComm = MPI_COMM_NULL;
        p->remoteGroup = MPI_GROUP_NULL;
      }
  }
}

/* create client <-> collector intercomm and fill-in related information */
static void
setupCollectorInterComm(pioInfo_t *p, int isCollector, MPI_Comm localGroupComm, int globRankOfGroupLeader[2])
{
  xmpi(MPI_Intercomm_create(localGroupComm, 0, p->commGlob, globRankOfGroupLeader[isCollector], 1, &p->pioInterComm));
  xmpi(MPI_Intercomm_merge(p->pioInterComm, isCollector, &p->collClientIntraComm));
  MPI_Group collClientGroup;
  int remoteRanksRange[3];
  xmpi(MPI_Comm_group(p->collClientIntraComm, &collClientGroup));
  int numClients = p->sizeCalc, numColl = p->sizeColl;
  if (isCollector)
    {
      int collRank = p->rankColl;
      remoteRanksRange[0] = cdiPioClientRangeStart(collRank, numClients, numColl);
      remoteRanksRange[1] = cdiPioClientRangeStart(collRank + 1, numClients, numColl) - 1;
    }
  else /* !isProcIO i.e. is client */
    {
      int clientRank = p->rankCalc;
      int collRank = cdiPioCollRank(clientRank, numClients, numColl);
      remoteRanksRange[1] = remoteRanksRange[0] = numClients + collRank;
    }
  remoteRanksRange[2] = 1;
  xmpi(MPI_Group_range_incl(collClientGroup, 1, &remoteRanksRange, &p->remoteGroup));
  xmpi(MPI_Group_free(&collClientGroup));
}

int
cdiPioCommInit(MPI_Comm commGlob, int IOMode, int clientServerRole)
{
  xassert(info == NULL && IOMode >= PIO_MINIOMODE && IOMode <= PIO_MAXIOMODE);
  info = Malloc(sizeof(*info));
  pioInfoInit(info, commGlob, clientServerRole);
  info->IOMode = IOMode;
  return info->sizePio;
}

void
cdiPioCommFinalize(void)
{
  xassert(info != NULL);

  MPI_Comm *comms[]
      = { &info->commGlob, &info->commPio, &info->commColl, &info->commCalc, &info->pioInterComm, &info->collClientIntraComm };

  for (size_t i = 0; i < sizeof(comms) / sizeof(comms[0]); ++i)
    if (*comms[i] != MPI_COMM_NULL) xmpi(MPI_Comm_free(comms[i]));

  MPI_Group *groups[] = { &info->remoteGroup };
  for (size_t i = 0; i < sizeof(groups) / sizeof(groups[0]); ++i)
    if (*groups[i] != MPI_GROUP_NULL) xmpi(MPI_Group_free(groups[i]));

  Free(info);
  info = NULL;
}

MPI_Comm
commInqCommGlob(void)
{
  xassert(info != NULL && info->commGlob != MPI_COMM_NULL);
  return info->commGlob;
}

int
commInqSizeGlob(void)
{
  xassert(info != NULL && info->sizeGlob != CDI_UNDEFID);
  return info->sizeGlob;
}

int
commInqRankGlob(void)
{
  xassert(info != NULL && info->rankGlob != CDI_UNDEFID);
  return info->rankGlob;
}

int
commInqRootGlob(void)
{
  xassert(info != NULL && info->root != CDI_UNDEFID);
  return info->root;
}

int
cdiPioInqCollectorOfClient(void)
{
  xassert(info != NULL);
  int clientRank = commInqRankModel(), numClients = cdiPioCommInqSizeClients(), numServers = commInqSizeColl(),
      numClientsPerServer = (numClients + numServers - 1) / numServers, collRank = clientRank / numClientsPerServer;
  return (!info->isProcIO) ? collRank : -1;
}

MPI_Group
cdiPioInqRemoteGroup(void)
{
  xassert(info != NULL);
  return info->remoteGroup;
}

int
commInqIsProcIO(void)
{
  xassert(info != NULL && info->isProcIO != CDI_UNDEFID);
  return info->isProcIO;
}

int
commInqIOMode(void)
{
  return info != NULL ? info->IOMode : PIO_NONE;
}

MPI_Comm
commInqCommPio(void)
{
  xassert(info != NULL && info->commPio != MPI_COMM_NULL && info->isProcIO == 1);
  return info->commPio;
}

int
commInqSizePio(void)
{
  xassert(info != NULL && info->sizePio > 0);
  return info->sizePio;
}

MPI_Comm
commInqCommModel(void)
{
  xassert(info != NULL && info->commCalc != MPI_COMM_NULL && info->isProcIO == 0);
  return info->commCalc;
}

int
commInqRankPio(void)
{
  xassert(info != NULL && info->rankPio >= 0 && info->isProcIO == 1);
  return info->rankPio;
}

int
commInqRankModel(void)
{
  xassert(info != NULL && info->rankCalc >= 0 && info->isProcIO == 0);
  return info->rankCalc;
}

int
cdiPioCommInqSizeClients(void)
{
  xassert(info != NULL && info->sizeCalc >= 0);
  return info->sizeCalc;
}

MPI_Comm
commInqCommColl(void)
{
  xassert(info != NULL && info->commColl != MPI_COMM_NULL);
  return info->commColl;
}

int
commInqSizeColl(void)
{
  xassert(info != NULL && info->sizeColl >= 0);
  return info->sizeColl;
}

int
commInqRankColl(void)
{
  xassert(info != NULL);
  return info->rankColl;
}

MPI_Comm
cdiPioInqInterComm(void)
{
  xassert(info != NULL);
  return info->pioInterComm;
}

MPI_Comm
cdiPioInqCollClientIntraComm(void)
{
  xassert(info != NULL);
  return info->collClientIntraComm;
}

int
commInqSpecialRank(void)
{
  xassert(info != NULL);
  return info->sizeColl;
}

MPI_Comm
commInqCommCalc(void)
{
  xassert(info != NULL && info->commCalc != MPI_COMM_NULL);
  return info->commCalc;
}

int
commInqNProcsColl(void)
{
  xassert(info != NULL && info->sizeColl != CDI_UNDEFID);
  return info->sizeColl;
}

void
commPrint(FILE *fp)
{
  if (ddebug == 0) return;

  xassert(info != NULL);

  fprintf(fp, "\n");
  fprintf(fp, "######## pioinfo PE%d ###########\n", info->rankGlob);
  fprintf(fp, "#\n");
  fprintf(fp, "# IOMode      = %d\n", info->IOMode);
  fprintf(fp, "# isProcIO    = %d\n", info->isProcIO);
  fprintf(fp, "#\n");
  fprintf(fp, "# commGlob    = %d\n", (int) MPI_Comm_c2f(info->commGlob));
  fprintf(fp, "# sizeGlob    = %d\n", info->sizeGlob);
  fprintf(fp, "# rankGlob    = %d\n", info->rankGlob);
  fprintf(fp, "# root        = %d\n", info->root);
  fprintf(fp, "#\n");
  fprintf(fp, "# commPio     = %d\n", (int) MPI_Comm_c2f(info->commPio));
  fprintf(fp, "# sizePio     = %d\n", info->sizePio);
  fprintf(fp, "# rankPio     = %d\n", info->rankPio);
  fprintf(fp, "#\n");
  fprintf(fp, "# commColl    = %d\n", (int) MPI_Comm_c2f(info->commColl));
  fprintf(fp, "# sizeColl    = %d\n", info->sizeColl);
  fprintf(fp, "# rankColl    = %d\n", info->rankColl);
  fprintf(fp, "#\n");
  fprintf(fp, "# commCalc    = %d\n", (int) MPI_Comm_c2f(info->commCalc));
  fprintf(fp, "# sizeCalc  = %d\n", info->sizeCalc);
  fprintf(fp, "# rankCalc  = %d\n", info->rankCalc);
  fprintf(fp, "#\n");
  fprintf(fp, "############################\n");
  fprintf(fp, "\n");
}

/************************************************************************/
/*
 * Local Variables:
 * c-file-style: "Java"
 * c-basic-offset: 2
 * indent-tabs-mode: nil
 * show-trailing-whitespace: t
 * require-trailing-newline: t
 * End:
 */
