/* ///////////////////////////////////////////////////////////////////// */
/* 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>


/* 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, const int);

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

int main(int argc, char **argv)
{
#if defined(_OPENMP)  
  if (argc <= 3)
    
    {
      printf("\n\t Usage: <executable> <x_grid_size> <y_grid_size> <omp threads>\n\n");
      exit(EXIT_FAILURE);
    }
#else
  if (argc <= 2)    
    {
      printf("\n\t Usage: <executable> <x_grid_size> <y_grid_size>\n\n");
      exit(EXIT_FAILURE);
    }
#endif /* _OPENMP */

  /* global X and Y grid size */
  const int NX_GLOB = (int) strtol(argv[1], NULL, 10);
  const int NY_GLOB = (int) strtol(argv[2], NULL, 10);
#if defined(_OPENMP)
  const int Nthread = (int) strtol(argv[3], NULL, 10);
#else
  const int Nthread = 1;
#endif /* _OPENMP */

  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;

  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
     -------------------------------------------------------- */

  /* 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;
  
  /* --------------------------------------------------------
     4. Main iteration cycle
     -------------------------------------------------------- */

  const double time_start = seconds();

  /* -- 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)
    {        
      /* -- 4b. Jacobi's method and residual (interior points) -- */
      /*       core algorithm                                     */
      
      err = 0.0;
      JacobiAlgorithm(phi, phi0, &err, delta,
		      ibeg, iend, jbeg, jend, Nthread);

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

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

      /* check convergence */
      if (err <= TOL)
	break;
      
      /* swap the pointers */
      MyData **tmp = phi;
      phi = phi0;
      phi0 = tmp;
    }
    
  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", seconds() - 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);
  
  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,
		     const int                    Nthreads)
{
  double err = 0.0;

#pragma omp parallel for collapse(2) schedule(static) reduction(+: err) num_threads(Nthreads)
  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]);
                
	  err += delta[X] * delta[Y] * fabs(Phi[j][i] - Phi0[j][i]);
	} /* loop over columns */
    } /* loop over rows */

  *error = err;
  
  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[32];
#if defined(_OPENMP)
  sprintf(fname,"jacobi2D_omp_opt_%02d.bin", nfile);
#else
  sprintf(fname,"jacobi2D_serial_opt_%02d.bin", nfile);
#endif /* _OPENMP */
    
  FILE *fp;
  printf ("> Writing %s\n",fname);
  fp = fopen(fname, "wb");
    
  for (int j=jbeg ; j<=jend ; j++)
    {
      fwrite (phi[j] + ibeg, sizeof(MyData), nx, fp);
    }
    
  nfile++;
  fclose(fp);
}
