/* ///////////////////////////////////////////////////////////////////// */
/* Authors:  A. Mignone (mignone@to.infn.it)                             */
/*           V. Cesare  (valentina.cesare@inaf.it)                       */
/*           D. Goz     (david.goz@inaf.it)                              */
/*                                                                       */
/* Date   : June 2024                                                    */
/*                                                                       */
/* ///////////////////////////////////////////////////////////////////// */

#include "allvars.h"
#include "tools.h"
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>

/* #define DEBUG */

typedef struct MyGrid
{
  int start[NDIM];
  int end[NDIM];
} myDomain;

/* function prototypes */
void BoundaryConditions(MyData **const restrict,
			MyData  *const restrict,
			MyData  *const restrict,
			const int,
			const int);

void JacobiAlgorithm(MyData **const restrict, MyData **const restrict, MyData *const restrict,
		     const MyData *const restrict, const int, const int, const int, const int);

void get_domains(MyData  **const buffer,
		 myDomain *const domain,
		 const MPI_Comm  comm);

void mpi_exchange_cartesian(MyData **const  buffer,
			    const int       nbrtop,
			    const int       nbrbottom,
			    const int       nbrleft,
			    const int       nbrright,
			    myDomain *const domain,
			    const int       grid_y,
			    const MPI_Comm  comm2d);

void WriteSolution(MyData **const phi, const int nx, const int ny);

void copy_grids(MyData **const restrict A,
                MyData **const restrict B,
                const int               xbeg,
                const int               xend,
                const int               ybeg,
		const int               yend);

int main(int argc, char **argv)
{
  int rank, Nranks;
  const MPI_Comm comm = MPI_COMM_WORLD;
  MPI_Init(&argc, &argv);
  MPI_Comm_rank(comm, &rank);
  MPI_Comm_size(comm, &Nranks);

  if ((argc <= 2) && (!rank))
    {
      printf("\n\t Usage: <executable> <grid_size> <cartesian communicator size_X>\n\n");
      MPI_Abort(comm, EXIT_FAILURE);
      exit(EXIT_FAILURE);
    }

  /* global X and Y grid size (square matrix supposed) */
  const int NX_GLOB = (int) strtol(argv[1], NULL, 10);
  const int NY_GLOB = NX_GLOB;

  /********************************** Cartesin topology *******************************************************/
  const int cartesian_grid_x = (int)strtol(argv[2], NULL, 10);
  const int cartesian_grid_y = ((Nranks % cartesian_grid_x == 0) ? (Nranks / cartesian_grid_x) : -1);
  if (cartesian_grid_y == -1)
    {
      if (!rank)
        {
          printf("\n\t Nranks mod cartesian_grid_x != 0 ... aborting ...\n");
          MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
          exit(EXIT_FAILURE);
        }
    }
  
  const int dims[NDIM] = {cartesian_grid_x, cartesian_grid_y};
  const int periods[NDIM] = {0 ,0}; /* not periodic */
  const int reorder = 0;            /* setting the logical value reorder to TRUE allows MPI to */
                                    /* reorder the rank of the processes figuring out the      */
                                    /* neighbors in the actual hardware for better performance */

  /* make a new communicator to which topology information has been attached */
  MPI_Comm comm2d = MPI_COMM_NULL;
  MPI_Cart_create(MPI_COMM_WORLD, NDIM, dims, periods, reorder, &comm2d);
  if (comm2d == MPI_COMM_NULL)
    {
      printf("\n\t Process %d is not part of the new communicator \n", rank);
      fflush(stdout);
      MPI_Abort(comm, EXIT_FAILURE);
      exit(EXIT_FAILURE);
    }

  /* determine the process coords in cartesian topology given rank in group */
  int coords[NDIM];
  MPI_Cart_coords(comm2d, rank, NDIM, coords);

  /* get bottom and top neighbors (X direction) */
  int nbrtop, nbrbottom;
  MPI_Cart_shift(comm2d, X, 1, &nbrbottom, &nbrtop);

  /* get left and right neighbors (Y direction) */
  int nbrright, nbrleft;
  MPI_Cart_shift(comm2d, Y, 1, &nbrleft, &nbrright);
  
  /************************************************************************************************************/
  /* 2D MPI-cartesian decomposition:
     the grids are replicated across the MPI processes.
     This approach is not the most efficient in terms of memory usage,
     because the arrays are replicated across MPI process instead of to be distributed */

  /* get the reminder, i.e. take into account uneven
     decomposition of points among the processes    */
  const int rem[NDIM] = {(NX_GLOB + 2) % dims[X],
                         (NY_GLOB + 2) % dims[Y]};

  /* get the amount of data for each MPI process:
     - chunk is supposed to be >= 1 */
  const int chunk[NDIM] = {(NX_GLOB + 2 - rem[X]) / dims[X],
                           (NY_GLOB + 2 - rem[Y]) / dims[Y]};
  if ((chunk[X] < 1) || (chunk[Y] < 1))
    {
      printf("\n\t chunk[X] < 1 || chunk[Y] < 1 ... aborting ...[rank: %d]\n", rank);
      MPI_Comm_free(&comm2d);
      MPI_Abort(comm, EXIT_FAILURE);
      exit(EXIT_FAILURE);
    }
  
  /* get the subgrid dimension along X and Y directions: */
  int incr[NDIM], offset[NDIM];
  for (int dim=0 ; dim<NDIM ; dim++)
    {
      if (coords[dim] < rem[dim])
	{
	  incr[dim] = chunk[dim] + 1;
	  offset[dim] = 0;
	}
      else
	{
	  incr[dim] = chunk[dim];
	  offset[dim] = rem[dim];
	}
    }

  /* subdomain managed by the task */
  myDomain domain;
  for (int dim=0 ; dim<NDIM ; dim++)
    {
      domain.start[dim] = ((coords[dim] * incr[dim]) + offset[dim]);
      domain.end[dim]   = (domain.start[dim] + incr[dim]) - 1;

      /* boundaries */
      domain.start[dim] = ((domain.start[dim] == 0)         ? NGHOST  : domain.start[dim]);
      domain.end[dim]   = ((domain.end[dim] == NX_GLOB + 1) ? NX_GLOB : domain.end[dim]);
    }

#if defined(DEBUG)
  for (int task=0 ; task<Nranks ; task++)
    {
      MPI_Barrier(comm);
      
      if (task == rank)
	{
	  printf("\n\t rank: %d", rank);
	  printf("\n\t\t coords = [%d, %d]", coords[X], coords[Y]);
	  printf("\n\t\t domain.start[X] = %d - domain.end[X] = %d", domain.start[X], domain.end[X]);
	  printf("\n\t\t domain.start[Y] = %d - domain.end[Y] = %d", domain.start[Y], domain.end[Y]);
	  printf("\n\t\t nbrtop  : %d - nbrbottom: %d", nbrtop, nbrbottom);
	  printf("\n\t\t nbrright: %d - nbrleft  : %d \n", nbrright, nbrleft);
	  fflush(stdout);
	}
    }

  MPI_Barrier(comm);
#endif /* DEBUG */
  
  /******************************************************************************************************/
  
  const MyData xbeg = 0.0;
  const MyData xend = 1.0;
  const MyData ybeg = 0.0;
  const MyData yend = 1.0;

  const MyData delta[NDIM] = {(xend - xbeg)/(NX_GLOB + 1),
			      (yend - ybeg)/(NY_GLOB + 1)};

  /* --------------------------------------------------------
     1. Set grid indices
     -------------------------------------------------------- */
    
  const int ibeg   = NGHOST;
  const int iend   = ibeg + NX_GLOB - 1;
  const int nx     = iend - ibeg + 1;
  const int nx_tot = nx + 2 * NGHOST;
    
  const int jbeg   = NGHOST;
  const int jend   = jbeg + NY_GLOB - 1;
  const int ny     = jend - jbeg + 1;
  const int ny_tot = ny + 2 * NGHOST;

  if (rank == MASTER)
    {
      printf("\n\t Grid indices:");
      printf("\n\t\t ibeg, iend = %d, %d; nx_tot = %d"    ,ibeg, iend, nx_tot);
      printf("\n\t\t jbeg, jend = %d, %d; ny_tot = %d\n\n",jbeg, jend, ny_tot);
    }
  
  /* --------------------------------------------------------
     2. Generate grid, allocate memory
        Not optimized because the grids are (unnecessarily)
	replicated across MPI processes
     -------------------------------------------------------- */

  /* memory allocation */
  MyData *xg = (MyData *) malloc((NX_GLOB + 2*NGHOST) * sizeof(MyData));
  MyData *yg = (MyData *) malloc((NY_GLOB + 2*NGHOST) * sizeof(MyData));
  assert((xg != NULL) && (yg != NULL));

  /* initial conditions */
  for (int i=0 ; i<(NX_GLOB + 2*NGHOST) ; i++) xg[i] = xbeg + (i - ibeg + 1) * delta[X];
  for (int j=0 ; j<(NY_GLOB + 2*NGHOST) ; j++) yg[j] = ybeg + (j - jbeg + 1) * delta[Y];
  MyData *x = xg; /* Global and local grids are the same  */
  MyData *y = yg; /* for serial version of the code       */

  /* grids memory allocation */
  MyData **phi  = Allocate_2DdblArray(ny_tot, nx_tot);
  MyData **phi0 = Allocate_2DdblArray(ny_tot, nx_tot);
    
  /* --------------------------------------------------------
     3. Initialize solution array to 0
     -------------------------------------------------------- */
    
  for (int j=jbeg ; j<=jend ; j++)
    for (int i=ibeg ; i<=iend ; i++)
      {
	phi0[j][i] = 0.0;
	phi[j][i]  = 0.0;
      }
  
  /* --------------------------------------------------------
     4. Main iteration cycle
     -------------------------------------------------------- */

  const double time_start = MPI_Wtime();
  
  /* -- 4a. Set boundary conditions first -- */  
  BoundaryConditions(phi0, x, y, nx, ny);
  BoundaryConditions(phi, x, y, nx, ny);

  MyData err = 1.0;
  /* iterations */
  int k = 0;
  while (1)
    {
      /* -- 4c. Jacobi's method and residual (interior points) -- */
      /*       core algorithm                                     */

      err = 0.0;
      JacobiAlgorithm(phi, phi0, &err, delta,
		      domain.start[Y], domain.end[Y],
		      domain.start[X], domain.end[X]);

      if (!rank)
	printf("\n\t Iteration = %d - err = %lg\n",k, err);

      /* increase the counter of loop iterations */
      k++;

      /* get the total error */
      MyData toterr;
      /* combines values from all processes and distributes the result back to all processes */
      MPI_Allreduce(&err, &toterr, 1, MPI_MyDatatype, MPI_SUM, comm2d);
      
      /* check convergence */
      if (toterr <= TOL)
	{
	  /* master task gathers all the domains */
	  get_domains(phi, &domain, comm2d);

	  break;
	}

      /* -- 4b. MPI communications */
      mpi_exchange_cartesian(phi,
			     nbrtop, nbrbottom, nbrleft, nbrright,
			     &domain, NY_GLOB,
			     comm2d);

      /* swap the pointers */
      MyData **tmp = phi;
      phi = phi0;
      phi0 = tmp;
    }

  /* master rank writes the solution */
  if (!rank)
    {
      WriteSolution(phi, nx, ny);
      
      printf("\n\t NX_GLOB x NY_GLOB = %d x %d\n", NX_GLOB, NY_GLOB);
      printf("\n\t Time = %lf [s]\n\n", MPI_Wtime() - time_start);
    }

  // free memory
  if (phi0)
    {
      free(phi0[0]);
      free(phi0);
    }

  if (phi)
    {
      free(phi[0]);
      free(phi);
    }

  if (yg)
    free(yg);
  
  if (xg)
    free(xg);

  MPI_Comm_free(&comm2d);
  
  MPI_Finalize();
  
  return 0;
}

/* ********************************************************************* */
void BoundaryConditions(MyData **const restrict phi,
			MyData  *const restrict x,
			MyData  *const restrict y,
                        const int               nx,
			const int               ny)
/*
*********************************************************************** */
{
  const int ibeg = NGHOST;
  const int iend = ibeg + nx - 1;
    
  const int jbeg = NGHOST;
  const int jend = jbeg + ny - 1;

  int i,j;
  
  /* -- Left -- */
  i = ibeg - 1;
  for (int j=jbeg ; j<=jend ; j++)
    phi[j][i] = (1.0 - y[j]);
    
  /* -- Right -- */
  i = jend + 1;
  for (int j=jbeg ; j<=jend ; j++)
    phi[j][i] = (y[j] * y[j]);
    
  /* -- Bottom -- */    
  j = jbeg - 1;
  for (int i=ibeg ; i<=iend ; i++)
    phi[j][i] = (1.0 - x[i]);
    
  /* -- Top -- */
  j = jend + 1;
  for (int i=ibeg ; i<=iend ; i++)
    phi[j][i] = x[i];

  return;
}

/* ********************************************************************* */

void JacobiAlgorithm(MyData      **const restrict Phi,
		     MyData      **const restrict Phi0,
		     MyData       *const restrict error,
		     const MyData *const restrict delta,
		     const int                    ibeg,
		     const int                    iend,
		     const int                    jbeg,
		     const int                    jend)
{
  *error = 0.0;
  for (int j=jbeg ; j<=jend ; j++)
    {
      for (int i=ibeg ; i<=iend ; i++)
	{
	  Phi[j][i] = 0.25 * (Phi0[j][i-1] + Phi0[j][i+1] +
			      Phi0[j-1][i] + Phi0[j+1][i]);
                
	  *error += delta[X] * delta[Y] * fabs(Phi[j][i] - Phi0[j][i]);
	} /* loop over columns */
    } /* loop over rows */
  
  return;
}

/* ********************************************************************* */
void mpi_exchange_cartesian(MyData **const  buffer,
			    const int       nbrtop,
			    const int       nbrbottom,
			    const int       nbrleft,
			    const int       nbrright,
			    myDomain *const domain,
			    const int       grid_y,
			    const MPI_Comm  comm2d)
{
 /* In the data exchange 1D routines seen so far, buffers are contiguous areas in memory,
    so that, for example, the datatypes MPI_DOUBLE_PRECISION coupled with the count of
    occurrences are sufficient to describe the buffer to be sent. 
    MPI derived datatypes allow the programmer to specify non-contiguous areas in memory, 
    such as a column of an array stored in a row-major order like in C programming language */
  
  /* Creates a vector (strided) datatype:                                              */
  /* - 1st argument (count)      : number of blocks;                                   */
  /* - 2nd argument (blocklength): number of elements in each block;                   */
  /* - 3rd argument (stride)     : number of elements between the start of each block; */
  /* - 4th argument (oldtype)    : old datatype;                                       */
  /* - 5th argument (newtype)    : new datatype                                        */

  /* count = (x_end - x_start + 1), i.e. the size of each column                       */
  /* blocklength = 1, i.e. one double (old datatype) in each block                     */
  /* stride = NY_GLOB + 2, i.e. ghost points are discarded (grid is replicated)        */
  /* oldtype = MPI_MyDatatype                                                          */
  
  MPI_Datatype column;
  MPI_Type_vector((domain->end[X] - domain->start[X] + 1), 1, grid_y + 2, MPI_MyDatatype, &column);
  /* commits the datatype */
  MPI_Type_commit(&column);

  /* communication with nbrtop and nbrbottom */
  const int data_row_size = (domain->end[Y] - domain->start[Y] + 1);
  MPI_Sendrecv(&buffer[domain->end[X]][domain->start[Y]]    , data_row_size, MPI_MyDatatype, nbrtop,    0,
	       &buffer[domain->start[X]-1][domain->start[Y]], data_row_size, MPI_MyDatatype, nbrbottom, 0,
	       comm2d, MPI_STATUS_IGNORE);

  MPI_Sendrecv(&buffer[domain->start[X]][domain->start[Y]], data_row_size, MPI_MyDatatype, nbrbottom, 1,
	       &buffer[domain->end[X]+1][domain->start[Y]], data_row_size, MPI_MyDatatype, nbrtop,    1,
	       comm2d, MPI_STATUS_IGNORE);

  /* communication with nbrleft, nbrright - vector datatype is used */
  MPI_Sendrecv(&buffer[domain->start[X]][domain->end[Y]],     1, column, nbrright, 0,
	       &buffer[domain->start[X]][domain->start[Y]-1], 1, column, nbrleft,  0,
	       comm2d, MPI_STATUS_IGNORE);

  MPI_Sendrecv(&buffer[domain->start[X]][domain->start[Y]], 1, column, nbrleft, 1,
	       &buffer[domain->start[X]][domain->end[Y]+1], 1, column, nbrright, 1,
	       comm2d, MPI_STATUS_IGNORE);  
  
  /* when the datatype is no longer needed, it should be freed, */
  /* i.e. datatype is set to MPI_TYPE_NULL                      */
  MPI_Type_free(&column);
  
  return;
}

/* ********************************************************************* */

void get_domains(MyData  **const buffer,
		 myDomain *const domain,
		 const MPI_Comm  comm)
{
  int rank, nranks;
  MPI_Comm_size(comm, &nranks);
  MPI_Comm_rank(comm, &rank);
  
  /* Master process gathers all the domains */
  
  /***************************** get the domain boundaries from each process *************************************/
  myDomain *boundaries = NULL;
  if (rank == MASTER)
    {
      boundaries = (myDomain *)malloc(nranks * sizeof(*boundaries));
      assert(boundaries != NULL);
    }

  MPI_Gather(domain,     sizeof(*domain), MPI_BYTE,
	     boundaries, sizeof(*domain), MPI_BYTE,
	     MASTER, comm);

#if defined(DEBUG)
  if (rank == MASTER)
    {
      for (int task=0 ; task<nranks ; task++)
	{
	  printf("\n\t Task: %d", task);
	  printf("\n\t\t boundaries[%d].start[X] = %d - boundaries[%d].end[X] = %d",
		 task, boundaries[task].start[X], task, boundaries[task].end[X]);
	  printf("\n\t\t boundaries[%d].start[Y] = %d - boundaries[%d].end[Y] = %d\n",
		 task, boundaries[task].start[Y], task, boundaries[task].end[Y]);
	}
      printf("\n");
      fflush(stdout);
    }
#endif /* DEBUG */

  /**************************************************************************************************/
  
  /***************************************** get the domain from each process *************************/
  /* MPI_Gatherv
     Gathers into specified locations from all processes in a group

     Synopsis
     int MPI_Gatherv(const void *sendbuf, int sendcount, MPI_Datatype sendtype,
                     void *recvbuf, const int recvcounts[], const int displs[],
                     MPI_Datatype recvtype, int root, MPI_Comm comm) */

  int *recvcounts = NULL;
  int *displs = NULL;
  MyData *recvbuf = NULL;
  MyData *sendbuf = NULL;
  
  if (rank == MASTER)
    {
      recvcounts = (int *)malloc(nranks * sizeof(*recvcounts));
      assert(recvcounts != NULL);

      displs = (int *)malloc(nranks * sizeof(*displs));
      assert(displs != NULL);

      for (int task=0 ; task<nranks ; task++)
	{
	  const int nrows = (boundaries[task].end[X] - boundaries[task].start[X] + 1);
	  const int ncols = (boundaries[task].end[Y] - boundaries[task].start[Y] + 1);
	  recvcounts[task] = (nrows * ncols);
	}

      displs[0] = 0;
      for (int task=1 ; task<nranks ; task++)
	displs[task] = (recvcounts[task - 1] + displs[task - 1]);

      /* allocate the recvbuffer */
      recvbuf = (MyData *)malloc((displs[nranks - 1] + recvcounts[nranks - 1]) * sizeof(*recvbuf));
      assert(recvbuf != NULL);
      
#if defined(DEBUG)
      printf("\n\t********************** recvcounts - displs *************************");
      for (int task=0 ; task<nranks ; task++)
	{
	  printf("\n\t recvcounts[%d] = %d - displs[%d] = %d",
		 task, recvcounts[task], task, displs[task]);
	}
      printf("\n\t********************************************************************\n");
      fflush(stdout);
#endif /* DEBUG */

    } /* MASTER */

  /* every process sets up the sendbuf */
  const int nrows = (domain->end[X] - domain->start[X] + 1);
  const int ncols = (domain->end[Y] - domain->start[Y] + 1);
  sendbuf = (MyData *)malloc(nrows * ncols * sizeof(*sendbuf));
  assert(sendbuf != NULL);
  /* store the actual grid owned by each process into sendbuf */
  for (int row=0 ; row<nrows ; row++)
    memcpy(&sendbuf[row * ncols], &buffer[domain->start[X]+row][domain->start[Y]], (ncols * sizeof(MyData)));
  
  /* perform the MPI collective */
  MPI_Gatherv(sendbuf, (nrows * ncols), MPI_MyDatatype,
	      recvbuf, recvcounts, displs, MPI_MyDatatype,
	      MASTER, comm);

  /* store the recvbuf into the actual grid owned by the MASTER task */
  if (rank == MASTER)
    {
      for (int task=1 ; task<nranks ; task++)
	{
	  const int nrows = (boundaries[task].end[X] - boundaries[task].start[X] + 1);
	  const int ncols = (boundaries[task].end[Y] - boundaries[task].start[Y] + 1);
	  for (int row=0 ; row<nrows ; row++)
	    memcpy(&buffer[boundaries[task].start[X] + row][boundaries[task].start[Y]],
		   &recvbuf[displs[task] + (row * ncols)], (ncols * sizeof(MyData)));
	}
    }

  free(sendbuf);
  if (rank == MASTER)
    {
      free(recvbuf);
      free(displs);
      free(recvcounts);
      free(boundaries);
    }
  
  return;
}

/* ********************************************************************* */
void WriteSolution(MyData **const phi,
		   const int      nx,
		   const int      ny)
/*
*********************************************************************** */
{
  const int ibeg = NGHOST;    
  const int jbeg = NGHOST;
  const int jend = jbeg + ny - 1;
    
  static int nfile = 0;  /* File counter */

  char fname[128];
  sprintf(fname,"jacobi2D_mpi_cartesian_%02d.bin", nfile);
    
  FILE *fp;
  printf ("> Writing %s\n",fname);
  fp = fopen(fname, "wb");

  /* discard boundaies */
  for (int j=jbeg ; j<=jend ; j++)
    {
      fwrite (phi[j] + ibeg, sizeof(MyData), nx, fp);
    }
    
  nfile++;
  fclose(fp);
}
