// @(#)root/minuit2:$Id$
// Author: A. Lazzaro 2009
/***************************************************************************
 * Package: Minuit2                                                        *
 *    File: $Id$         *
 *  Author: Alfio Lazzaro, alfio.lazzaro@mi.infn.it                        *
 *                                                                         *
 * Copyright: (C) 2008 by Universita' and INFN, Milan                      *
 ***************************************************************************/

#include "Minuit2/MPIProcess.h"

#include <iostream>

namespace ROOT {

namespace Minuit2 {

unsigned int MPIProcess::fgGlobalSize = 1;
unsigned int MPIProcess::fgGlobalRank = 0;

// By default all procs are for X
unsigned int MPIProcess::fgCartSizeX = 0;
unsigned int MPIProcess::fgCartSizeY = 0;
unsigned int MPIProcess::fgCartDimension = 0;
bool MPIProcess::fgNewCart = true;

#ifdef MPIPROC
MPI::Intracomm *MPIProcess::fgCommunicator = 0;
int MPIProcess::fgIndexComm = -1; // -1 for no-initialization
MPI::Intracomm *MPIProcess::fgCommunicators[2] = {0};
unsigned int MPIProcess::fgIndicesComm[2] = {0};
#endif

MPIProcess::MPIProcess(unsigned int nelements, unsigned int indexComm) : fNelements(nelements), fSize(1), fRank(0)
{

   // check local requested index for communicator, valid values are 0 and 1
   indexComm = (indexComm == 0) ? 0 : 1;

#ifdef MPIPROC

   StartMPI();

   if (fgGlobalSize == fgCartDimension && fgCartSizeX != fgCartDimension && fgCartSizeY != fgCartDimension) {
      // declare the cartesian topology

      if (fgCommunicator == 0 && fgIndexComm < 0 && fgNewCart) {
         // first call, declare the topology
         std::cout << "Info --> MPIProcess::MPIProcess: Declare cartesian Topology (" << fgCartSizeX << "x"
                   << fgCartSizeY << ")" << std::endl;

         int color = fgGlobalRank / fgCartSizeY;
         int key = fgGlobalRank % fgCartSizeY;

         fgCommunicators[0] = new MPI::Intracomm(MPI::COMM_WORLD.Split(key, color)); // rows for Minuit
         fgCommunicators[1] = new MPI::Intracomm(MPI::COMM_WORLD.Split(color, key)); // columns for NLL

         fgNewCart = false;
      }

      fgIndexComm++;

      if (fgIndexComm > 1 ||
          fgCommunicator == (&(MPI::COMM_WORLD))) { // Remember, no more than 2 dimensions in the topology!
         std::cerr << "Error --> MPIProcess::MPIProcess: Requiring more than 2 dimensions in the topology!"
                   << std::endl;
         MPI::COMM_WORLD.Abort(-1);
      }

      // requiring columns as first call. In this case use all nodes
      if (((unsigned int)fgIndexComm) < indexComm)
         fgCommunicator = &(MPI::COMM_WORLD);
      else {
         fgIndicesComm[fgIndexComm] = indexComm;
         fgCommunicator = fgCommunicators[fgIndicesComm[fgIndexComm]];
      }

   } else {
      // no cartesian topology
      if (fgCartDimension != 0 && fgGlobalSize != fgCartDimension) {
         std::cout << "Warning --> MPIProcess::MPIProcess: Cartesian dimension doesn't correspond to # total procs!"
                   << std::endl;
         std::cout << "Warning --> MPIProcess::MPIProcess: Ignoring topology, use all procs for X." << std::endl;
         std::cout << "Warning --> MPIProcess::MPIProcess: Resetting topology..." << std::endl;
         fgCartSizeX = fgGlobalSize;
         fgCartSizeY = 1;
         fgCartDimension = fgGlobalSize;
      }

      if (fgIndexComm < 0) {
         if (fgCartSizeX == fgCartDimension) {
            fgCommunicators[0] = &(MPI::COMM_WORLD);
            fgCommunicators[1] = 0;
         } else {
            fgCommunicators[0] = 0;
            fgCommunicators[1] = &(MPI::COMM_WORLD);
         }
      }

      fgIndexComm++;

      if (fgIndexComm > 1) { // Remember, no more than 2 nested MPI calls!
         std::cerr << "Error --> MPIProcess::MPIProcess: More than 2 nested MPI calls!" << std::endl;
         MPI::COMM_WORLD.Abort(-1);
      }

      fgIndicesComm[fgIndexComm] = indexComm;

      // require 2 nested communicators
      if (fgCommunicator != 0 && fgCommunicators[indexComm] != 0) {
         std::cout << "Warning --> MPIProcess::MPIProcess: Requiring 2 nested MPI calls!" << std::endl;
         std::cout << "Warning --> MPIProcess::MPIProcess: Ignoring second call." << std::endl;
         fgIndicesComm[fgIndexComm] = (indexComm == 0) ? 1 : 0;
      }

      fgCommunicator = fgCommunicators[fgIndicesComm[fgIndexComm]];
   }

   // set size and rank
   if (fgCommunicator != 0) {
      fSize = fgCommunicator->Get_size();
      fRank = fgCommunicator->Get_rank();
   } else {
      // no MPI calls
      fSize = 1;
      fRank = 0;
   }

   if (fSize > fNelements) {
      std::cerr << "Error --> MPIProcess::MPIProcess: more processors than elements!" << std::endl;
      MPI::COMM_WORLD.Abort(-1);
   }

#endif

   fNumElements4JobIn = fNelements / fSize;
   fNumElements4JobOut = fNelements % fSize;
}

MPIProcess::~MPIProcess()
{
   // destructor
#ifdef MPIPROC
   fgCommunicator = 0;
   fgIndexComm--;
   if (fgIndexComm == 0)
      fgCommunicator = fgCommunicators[fgIndicesComm[fgIndexComm]];

#endif
}

bool MPIProcess::SyncVector(ROOT::Minuit2::MnAlgebraicVector &mnvector)
{

   // In case of just one job, don't need sync, just go
   if (fSize < 2)
      return false;

   if (mnvector.size() != fNelements) {
      std::cerr << "Error --> MPIProcess::SyncVector: # defined elements different from # requested elements!"
                << std::endl;
      std::cerr << "Error --> MPIProcess::SyncVector: no MPI synchronization is possible!" << std::endl;
      exit(-1);
   }

#ifdef MPIPROC
   unsigned int numElements4ThisJob = NumElements4Job(fRank);
   unsigned int startElementIndex = StartElementIndex();
   unsigned int endElementIndex = EndElementIndex();

   double dvectorJob[numElements4ThisJob];
   for (unsigned int i = startElementIndex; i < endElementIndex; i++)
      dvectorJob[i - startElementIndex] = mnvector(i);

   double dvector[fNelements];
   MPISyncVector(dvectorJob, numElements4ThisJob, dvector);

   for (unsigned int i = 0; i < fNelements; i++) {
      mnvector(i) = dvector[i];
   }

   return true;

#else

   std::cerr << "Error --> MPIProcess::SyncVector: no MPI synchronization is possible!" << std::endl;
   exit(-1);

#endif
}

bool MPIProcess::SyncSymMatrixOffDiagonal(ROOT::Minuit2::MnAlgebraicSymMatrix &mnmatrix)
{

   // In case of just one job, don't need sync, just go
   if (fSize < 2)
      return false;

   if (mnmatrix.size() - mnmatrix.Nrow() != fNelements) {
      std::cerr
         << "Error --> MPIProcess::SyncSymMatrixOffDiagonal: # defined elements different from # requested elements!"
         << std::endl;
      std::cerr << "Error --> MPIProcess::SyncSymMatrixOffDiagonal: no MPI synchronization is possible!" << std::endl;
      exit(-1);
   }

#ifdef MPIPROC
   unsigned int numElements4ThisJob = NumElements4Job(fRank);
   unsigned int startElementIndex = StartElementIndex();
   unsigned int endElementIndex = EndElementIndex();
   unsigned int nrow = mnmatrix.Nrow();

   unsigned int offsetVect = 0;
   for (unsigned int i = 0; i < startElementIndex; i++)
      if ((i + offsetVect) % (nrow - 1) == 0)
         offsetVect += (i + offsetVect) / (nrow - 1);

   double dvectorJob[numElements4ThisJob];
   for (unsigned int i = startElementIndex; i < endElementIndex; i++) {

      int x = (i + offsetVect) / (nrow - 1);
      if ((i + offsetVect) % (nrow - 1) == 0)
         offsetVect += x;
      int y = (i + offsetVect) % (nrow - 1) + 1;

      dvectorJob[i - startElementIndex] = mnmatrix(x, y);
   }

   double dvector[fNelements];
   MPISyncVector(dvectorJob, numElements4ThisJob, dvector);

   offsetVect = 0;
   for (unsigned int i = 0; i < fNelements; i++) {

      int x = (i + offsetVect) / (nrow - 1);
      if ((i + offsetVect) % (nrow - 1) == 0)
         offsetVect += x;
      int y = (i + offsetVect) % (nrow - 1) + 1;

      mnmatrix(x, y) = dvector[i];
   }

   return true;

#else

   std::cerr << "Error --> MPIProcess::SyncMatrix: no MPI synchronization is possible!" << std::endl;
   exit(-1);

#endif
}

#ifdef MPIPROC
void MPIProcess::MPISyncVector(double *ivector, int svector, double *ovector)
{
   int offsets[fSize];
   int nconts[fSize];
   nconts[0] = NumElements4Job(0);
   offsets[0] = 0;
   for (unsigned int i = 1; i < fSize; i++) {
      nconts[i] = NumElements4Job(i);
      offsets[i] = nconts[i - 1] + offsets[i - 1];
   }

   fgCommunicator->Allgatherv(ivector, svector, MPI::DOUBLE, ovector, nconts, offsets, MPI::DOUBLE);
}

bool MPIProcess::SetCartDimension(unsigned int dimX, unsigned int dimY)
{
   if (fgCommunicator != 0 || fgIndexComm >= 0) {
      std::cout << "Warning --> MPIProcess::SetCartDimension: MPIProcess already declared! Ignoring command..."
                << std::endl;
      return false;
   }
   if (dimX * dimY <= 0) {
      std::cout << "Warning --> MPIProcess::SetCartDimension: Invalid topology! Ignoring command..." << std::endl;
      return false;
   }

   StartMPI();

   if (fgGlobalSize != dimX * dimY) {
      std::cout << "Warning --> MPIProcess::SetCartDimension: Cartesian dimension doesn't correspond to # total procs!"
                << std::endl;
      std::cout << "Warning --> MPIProcess::SetCartDimension: Ignoring command..." << std::endl;
      return false;
   }

   if (fgCartSizeX != dimX || fgCartSizeY != dimY) {
      fgCartSizeX = dimX;
      fgCartSizeY = dimY;
      fgCartDimension = fgCartSizeX * fgCartSizeY;
      fgNewCart = true;

      if (fgCommunicators[0] != 0 && fgCommunicators[1] != 0) {
         delete fgCommunicators[0];
         fgCommunicators[0] = 0;
         fgIndicesComm[0] = 0;
         delete fgCommunicators[1];
         fgCommunicators[1] = 0;
         fgIndicesComm[1] = 0;
      }
   }

   return true;
}

bool MPIProcess::SetDoFirstMPICall(bool doFirstMPICall)
{

   StartMPI();

   bool ret;
   if (doFirstMPICall)
      ret = SetCartDimension(fgGlobalSize, 1);
   else
      ret = SetCartDimension(1, fgGlobalSize);

   return ret;
}

#endif

#ifdef MPIPROC
MPITerminate dummyMPITerminate = MPITerminate();
#endif

} // namespace Minuit2

} // namespace ROOT
