# include <cmath>
# include <cstdlib>
# include <iomanip>
# include <iostream>

using namespace std;

# include "poisson_2d.hpp"

//****************************************************************************80

void poisson_2d ( int nx, int ny, void rhs ( int nx, int ny, double **f ), 
  double u_exact ( double x, double y ) )

//****************************************************************************80
//
//  Purpose:
//
//    poisson_2d() solves the Poisson problem using Jacobi iteration.
//
//  Discussion:
//
//    The Poisson equation
//
//      - DEL^2 U(x,y) = F(x,y)
//
//    is solved on the unit square [0,1] x [0,1] using a grid of NX by
//    NX evenly spaced points.  The first and last points in each direction
//    are boundary points.
//
//    The boundary conditions and F are set so that the exact solution is
//
//      U(x,y) = sin ( pi * x * y )
//
//    so that
//
//      - DEL^2 U(x,y) = pi^2 * ( x^2 + y^2 ) * sin ( pi * x * y )
//
//    The Jacobi iteration is repeatedly applied until convergence is detected.
//
//  Licensing:
//
//    This code is distributed under the MIT license. 
//
//  Modified:
//
//    04 October 2024
//
//  Author:
//
//    John Burkardt
//
//  Input:
//
//    int NX, NY, the X and Y grid dimensions.
//
//    RHS(X,Y,F): the routine that evaluates the right hand side.
//
//    U_EXACT(X,Y): the function that evaluates the exact solution.
//
{
  bool converged;
  double diff;
  double dx;
  double dy;
  double error;
  double **f;
  int i;
  int it;
  int it_max = 30000;
  int j;
  double tolerance = 0.000001;
  double **u;
  double u_norm;
  double **udiff;
  double **uexact;
  double **unew;
  double unew_norm;
  double x;
  double y;

  dx = 1.0 / ( double ) ( nx - 1 );
  dy = 1.0 / ( double ) ( ny - 1 );
//
//  Print a message.
//
  cout << "\n";
  cout << "poisson_2d():\n";
  cout << "  C++ version\n";
  cout << "  A program for solving the Poisson equation.\n";
  cout << "\n";
  cout << "  -DEL^2 U = F(X,Y)\n";
  cout << "\n";
  cout << "  on the rectangle 0 <= X <= 1, 0 <= Y <= 1.\n";
  cout << "\n";
  cout << "  F(X,Y) = pi^2 * ( x^2 + y^2 ) * sin ( pi * x * y )\n";
  cout << "\n";
  cout << "  The number of interior X grid points is " << nx << "\n";
  cout << "  The number of interior Y grid points is " << ny << "\n";
  cout << "  The X grid spacing is " << dx << "\n";
  cout << "  The Y grid spacing is " << dy << "\n";
//
//  Allocate the array data.
//
  f = new double* [nx];
  for ( i = 0; i < nx; i++)
  {
    f[i] = new double[ny];
  }

  u = new double* [nx];
  for ( i = 0; i < nx; i++)
  {
    u[i] = new double[ny];
  }

  udiff = new double* [nx];
  for ( i = 0; i < nx; i++)
  {
    udiff[i] = new double[ny];
  }

  uexact = new double* [nx];
  for ( i = 0; i < nx; i++)
  {
    uexact[i] = new double[ny];
  }

  unew = new double* [nx];
  for ( i = 0; i < nx; i++)
  {
    unew[i] = new double[ny];
  }
//
//  Initialize the right hand side.
//
  rhs ( nx, ny, f );
//
//  Set the initial solution estimate.
//  We are "allowed" to pick up the boundary conditions exactly.
//
  for ( j = 0; j < ny; j++ )
  {
    for ( i = 0; i < nx; i++ )
    {
      if ( i == 0 || i == nx - 1 || j == 0 || j == ny - 1 )
      {
        unew[i][j] = f[i][j];
      }
      else
      {
        unew[i][j] = 0.0;
      }
    }
  }
  unew_norm = rms_norm ( nx, ny, unew );
//
//  Set up the exact solution.
//
  for ( j = 0; j < ny; j++ )
  {
    y = ( double ) ( j ) / ( double ) ( ny - 1 );
    for ( i = 0; i < nx; i++ )
    {
      x = ( double ) ( i ) / ( double ) ( nx - 1 );
      uexact[i][j] = u_exact ( x, y );
    }
  }
  u_norm = rms_norm ( nx, ny, uexact );
  cout << "  RMS of exact solution = " << u_norm << "\n";
//
//  Do the iteration.
//
  converged = false;

  cout << "\n";
  cout << "  Step    ||Unew||     ||Unew-U||     ||Unew-Exact||\n";
  cout << "\n";

  for ( j = 0; j < ny; j++ )
  {
    for ( i = 0; i < nx; i++ )
    {
      udiff[i][j] = unew[i][j] - uexact[i][j];
    }
  }
  error = rms_norm ( nx, ny, udiff );
  cout << "  " << setw(4) << 0
       << "  " << setw(14) << unew_norm
       << "  " << "              "
       << "  " << setw(14) << error << "\n";

  for ( it = 1; it <= it_max; it++ )
  {
    for ( j = 0; j < ny; j++ )
    {
      for ( i = 0; i < nx; i++ )
      {
        u[i][j] = unew[i][j];
      }
    }
//
//  UNEW is derived from U by one Jacobi step.
//
    jacobi ( nx, ny, dx, dy, f, u, unew );
//
//  Check for convergence.
//
    u_norm = unew_norm;
    unew_norm = rms_norm ( nx, ny, unew );

    for ( j = 0; j < ny; j++ )
    {
      for ( i = 0; i < nx; i++ )
      {
        udiff[i][j] = unew[i][j] - u[i][j];
      }
    }
    diff = rms_norm ( nx, ny, udiff );

    for ( j = 0; j < ny; j++ )
    {
      for ( i = 0; i < nx; i++ )
      {
        udiff[i][j] = unew[i][j] - uexact[i][j];
      }
    }
    error = rms_norm ( nx, ny, udiff );

    if ( it % 500 == 0 || diff <= tolerance )
    {
      cout << "  " << setw(4)  << it
           << "  " << setw(14) << unew_norm
           << "  " << setw(14) << diff
           << "  " << setw(14) << error << "\n";
    }

    if ( diff <= tolerance )
    {
      converged = true;
      break;
    }

  }

  if ( converged )
  {
    cout << "  The iteration has converged.\n";
  }
  else
  {
    cout << "  The iteration has NOT converged.\n";
  }
//
//  Free memory.
//
  for ( i = 0; i < nx; i++)
  {
    delete [] f[i];
  }
  delete [] f;

  for ( i = 0; i < nx; i++)
  {
    delete [] u[i];
  }
  delete [] u;

  for ( i = 0; i < nx; i++)
  {
    delete [] udiff[i];
  }
  delete [] udiff;

  for ( i = 0; i < nx; i++)
  {
    delete [] uexact[i];
  }
  delete [] uexact;

  for ( i = 0; i < nx; i++)
  {
    delete [] unew[i];
  }
  delete [] unew;

  return;
}
//****************************************************************************80

void jacobi ( int nx, int ny, double dx, double dy, double **f, 
  double **u, double **unew )

//****************************************************************************80
//
//  Purpose:
//
//   jacobi() carries out one step of the Jacobi iteration.
//
//  Licensing:
//
//    This code is distributed under the MIT license. 
//
//  Modified:
//
//    26 October 2011
//
//  Author:
//
//    John Burkardt
//
//  Input:
//
//    int NX, NY, the X and Y grid dimensions.
//
//    double DX, DY, the spacing between grid points.
//
//    double F[NX][NY], the right hand side data.
//
//    double U[NX][NY], the previous solution estimate.
//
//  Output:
//
//    double UNEW[NX][NY], the updated solution estimate.
//
{
  int i;
  int j;

  for ( j = 0; j < ny; j++ )
  {
    for ( i = 0; i < nx; i++ )
    {
      if ( i == 0 || j == 0 || i == nx - 1 || j == ny - 1 )
      {
        unew[i][j] = f[i][j];
      }
      else
      { 
        unew[i][j] = 0.25 * ( 
          u[i-1][j] + u[i][j+1] + u[i][j-1] + u[i+1][j] + f[i][j] * dx * dy );
      }
    }
  }
  return;
}
//****************************************************************************80

double rms_norm ( int nx, int ny, double **a )

//****************************************************************************80
//
//  Purpose:
//
//    rms_norm() returns the RMS norm of a vector stored as a matrix.
//
//  Licensing:
//
//    This code is distributed under the MIT license.
//
//  Modified:
//
//    01 March 2003
//
//  Author:
//
//    John Burkardt
//
//  Input:
//
//    int NX, NY, the number of rows and columns in A.
//
//    double **A, the vector.
//
//  Output:
//
//    double rms_norm, the root mean square of the entries of A.
//
{
  int i;
  int j;
  double v;

  v = 0.0;

  for ( j = 0; j < ny; j++ )
  {
    for ( i = 0; i < nx; i++ )
    {
      v = v + a[i][j] * a[i][j];
    }
  }
  v = sqrt ( v / ( double ) ( nx * ny )  );

  return v;
}