#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "tools.c"

#define NX_GLOB     128  /* Global number of interior points */
#define NY_GLOB     128  /* Global number of interior points */
#define NGHOST   1
#define NDIM     2

void BoundaryConditions(double **, double *, double *, int, int, int procL[], int procR[]);
void WriteSolution (double **phi, int nx, int ny);

int main(int argc, char ** argv)
{
  int    nx_tot, nx, i, ibeg, iend; 
  int    ny_tot, ny, j, jbeg, jend; 
  int    k, rank=0, size=1;
  double xbeg = 0.0, xend = 1.0;
  double ybeg = 0.0, yend = 1.0;
  double dx = (xend - xbeg)/(NX_GLOB + 1);
  double dy = (yend - ybeg)/(NY_GLOB + 1);
  double *xg, *yg, *x, *y, **phi, **phi0;
  double err, tol;
  int procL[NDIM] = {-1,-1};
  int procR[NDIM] = {-1,-1};

  nx = NX_GLOB; 
  ny = NY_GLOB;

/* --------------------------------------------------------
   1. Set grid indices
   -------------------------------------------------------- */
  
  ibeg   = NGHOST;
  iend   = ibeg + nx - 1;
  nx     = iend - ibeg + 1;
  nx_tot = nx + 2*NGHOST;

  jbeg   = NGHOST;
  jend   = jbeg + ny - 1;
  ny     = jend - jbeg + 1;
  ny_tot = ny + 2*NGHOST;

  printf ("ibeg, iend = %d, %d; nx_tot = %d\n",ibeg, iend, nx_tot);
  printf ("jbeg, jend = %d, %d; ny_tot = %d\n",jbeg, jend, ny_tot);

/* --------------------------------------------------------
   2. Generate grid, allocate memory
   -------------------------------------------------------- */

  xg = (double *) malloc ( (NX_GLOB+2*NGHOST)*sizeof(double)); 
  yg = (double *) malloc ( (NY_GLOB+2*NGHOST)*sizeof(double));
  for (i = 0; i < (NX_GLOB+2*NGHOST); i++) xg[i] = xbeg + (i-ibeg+1)*dx;
  for (j = 0; j < (NY_GLOB+2*NGHOST); j++) yg[j] = ybeg + (j-jbeg+1)*dy;
  x = xg; /* Global and local grid are the same  */ 
  y = yg; /* for serial version of the code      */

  phi  = Allocate_2DdblArray(ny_tot, nx_tot);
  phi0 = Allocate_2DdblArray(ny_tot, nx_tot);

/* --------------------------------------------------------
   3. Initialize solution array to 0 
   -------------------------------------------------------- */

  for (j = jbeg; j <= jend; j++){
  for (i = ibeg; i <= iend; i++){
    phi0[j][i] = 0.0;
  }}

/* --------------------------------------------------------
   4. Main iteration cycle
   -------------------------------------------------------- */

  tol = 1.e-5;
  err = 1.0;
  k   = 0;
  while  (err > tol){

  /* -- 4a. Set boundary conditions first -- */

    BoundaryConditions(phi0, x, y, nx, ny, procL, procR);

  /* -- 4b. Jacobi's method and residual (interior points) -- */

    err = 0.0;
    for (j = jbeg; j <= jend; j++){
    for (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 += dx*dy*fabs(phi[j][i] - phi0[j][i]);
    }}

    for (j = jbeg; j <= jend; j++){
    for (i = ibeg; i <= iend; i++){
      phi0[j][i] = phi[j][i];
    }}

    if (rank == 0){
      printf ("k = %d; err = %8.3e\n",k, err);
    }
    k++;
  }

  WriteSolution (phi, nx, ny);

  return 0;
}

/* ********************************************************************* */
void BoundaryConditions(double **phi, double *x, double *y,
                        int nx, int ny, int procL[], int procR[])
/*
 *********************************************************************** */
{
  int i,j;
  int ibeg   = NGHOST;
  int iend   = ibeg + nx - 1;

  int jbeg   = NGHOST;
  int jend   = jbeg + ny - 1;
  
/* -- Left -- */

  if (procL[0] < 0){
    i = ibeg-1;
    for (j = jbeg; j <= jend; j++) phi[j][i] = 1.0-y[j];
  }

/* -- Right -- */

  if (procR[0] < 0){
    i = jend+1;
    for (j = jbeg; j <= jend; j++) phi[j][i] = y[j]*y[j];
  }

/* -- Bottom -- */

  if (procL[1] < 0){
    j = jbeg-1;
    for (i = ibeg; i <= iend; i++) phi[j][i] = 1.0-x[i];
  }

/* -- Top -- */

  if (procR[1] < 0){
    j = jend+1;
    for (i = ibeg; i <= iend; i++) phi[j][i] = x[i];
  }

  return;
}

/* ********************************************************************* */
void WriteSolution (double **phi, int nx, int ny)
/*
 *********************************************************************** */
{
  int i,j;
  int ibeg   = NGHOST;
  int iend   = ibeg + nx - 1;

  int jbeg   = NGHOST;
  int jend   = jbeg + ny - 1;

  static int nfile = 0;  /* File counter */
  char fname[32];
  
  sprintf (fname,"laplace2D_%02d.bin",nfile);

  FILE *fp;
  printf ("> Writing %s\n",fname);
  fp = fopen(fname, "wb");

  for (j = jbeg; j <= jend; j++){
    fwrite (phi[j] + ibeg, sizeof(double), nx, fp);
  }

  nfile++;
  fclose(fp);
}