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

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)
{
#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();

  MyData err = 1.0;
  /* iterations */
  int k = 0;
  while (err > TOL)
    {        
      /* -- 4a. Set boundary conditions first -- */
        
      BoundaryConditions(phi0, x, y, nx, ny);
        
      /* -- 4b. Jacobi's method and residual (interior points) -- */
      /*       core algorithm                                     */
      
      err = 0.0;
      JacobiAlgorithm(phi, phi0, &err, delta,
		      ibeg, iend, jbeg, jend, Nthread);
      
      /* -- 4c. Copy grids --*/
      copy_grids(phi0, phi,
		 jbeg, jend,
		 ibeg, jend);

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

      /* increase the counter of loop iterations */
      k++;
    }
    
  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,
		           int                    ibeg,
		           int                    iend,
		           int                    jbeg,
		           int                    jend,
		     const int                    Nthreads)
{
  *error = 0.0;

#pragma omp parallel default(none) shared(error, delta, Phi, Phi0)      \
                                   firstprivate(ibeg, iend, jbeg, jend) \
                                   num_threads(Nthreads)
  {
#if defined(DEBUG)
    const int id  = omp_get_thread_num();
    const int thr = omp_get_num_threads();
    printf("\n\t I am the omp thread %d out of %d", id, thr);
    #pragma omp barrier
#endif /* DEBUG */

    #pragma omp for schedule(static)
    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]);

	    #pragma omp critical
	    {
	      *error += delta[X] * delta[Y] * fabs(Phi[j][i] - Phi0[j][i]);
	    }
	  } /* loop over columns */
      } /* loop over rows */
  } /* omp parallel */

  return;
}

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

/* ********************************************************************* */
void copy_grids(MyData **const restrict A,
		MyData **const restrict B,
		const int               xbeg,
		const int               xend,
	        const int               ybeg,
		const int               yend)
{
  for (int i=xbeg ; i<=xend ; i++)
    for (int j=ybeg ; j<=yend ; j++)
      A[i][j] = B[i][j];
  
  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_not_opt_%02d.bin", nfile);
#else
  sprintf(fname,"jacobi2D_serial_not_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);
}
