//===========================================================================
//  CVS Information:                                                         
//                                                                           
//     $RCSfile: bratu.cpp,v $  $Revision: 1.3 $  $State: Exp $ 
//     $Author: llee $  $Date: 2001/10/26 16:33:55 $ 
//     $Locker:  $ 
//---------------------------------------------------------------------------
//                                                                           
// DESCRIPTION                                                               
//  Bratu problem, using Newton-Krylov method
//
//  -Laplacian(U) - lambda * exp (U) = 0
//  u = 0 at boundary. lambda = 6.0
//
//
//  nonlinear equations: f(i, j) = 0 where
//  f(i,j) =  U(i,j)                                       if i,j = 0, n
//  f(i,j) =  4U(i,j)-U(i-1,j)-U(i+1,j)-U(i,j-1)-U(i,j+1)  otherwise
//                              -hx*hy*lambda*exp(U(i,j))
//
//                                                                           
//---------------------------------------------------------------------------
//                                                                           
// LICENSE AGREEMENT                                                         
//     $COPYRIGHT$
//---------------------------------------------------------------------------
//                                                                           
// REVISION HISTORY:                                                         
//                                                                           
// $Log: bratu.cpp,v $
// Revision 1.3  2001/10/26 16:33:55  llee
// *** empty log message ***
//
// Revision 1.2  2001/10/18 21:37:05  llee
// *** empty log message ***
//
// Revision 1.1  2000/07/26 21:49:30  llee1
// change file extension from .cc to .cpp
//
// Revision 1.3  2000/07/26 17:14:44  llee1
// *** empty log message ***
//
//                                                                           
//===========================================================================

#include <fstream>

#include <iostream>
#include <string>
#include "blitz/array.h"
#include "blitz/array/where.h"
#include "blitz/array/stencil-et.h"

using namespace blitz;

#include "itl/interface/blitz.h"

#include "itl/matrix_free_operator.h"
#include "itl/krylov/gmres.h"
//#include "itl/krylov/bicg.h"  //require trans_mult()
#include "itl/krylov/bicgstab.h"
#include "itl/krylov/cgs.h"
//#include "itl/krylov/qmr.h"   //require trans_mult()
#include "itl/krylov/tfqmr.h"

#include "parser.h"

using namespace itl;

double lamhxhy;

BZ_DECLARE_STENCIL2(Bratu_stencil, U, F)
  F = -Laplacian2D(U) - exp(U) * lamhxhy;
BZ_END_STENCIL

#define MIN(X,Y) where((X) < (Y), (X), (Y))

namespace itl {
// Wrap up applyStencil
struct apply_op {
  void
  operator()(Array<double, 2>& U, Array<double, 2>& F) const {
    applyStencil(Bratu_stencil(), U, F);
  }
};
}


int 
main(int argc, char *argv[])
{
  using std::cout;
  using std::endl;
  int monitor = 0;
  int kmax;
  int iter_max;
  int restart;

  std::string ksp_type;

  double ksp_rtol, ksp_atol, newton_rtol, newton_atol, newton_stol;

  int nx, ny;

  if ( argc == 1 ) {
    std::cout << argv[0] << " --help will get your the usage of flags." 
	      << std::endl;
  }

  parser myparser(argc, argv);

  myparser.register_flag("--mx", 1, "number of points in x axis in the mesh");
  myparser.register_flag("--my", 1, "number of points in y axis in the mesh");
  myparser.register_flag("--ksp_max_it", 1, "maxmal iterations allowed in KSP");  
  myparser.register_flag("--ksp_atol", 1, "Absolute tolerance in KSP iteration");  
  myparser.register_flag("--ksp_rtol", 1, "Relative tolerance in KSP iteration");  
  myparser.register_flag("--newton_rtol", 1, "Relative tolerance in Newton iteration");
  myparser.register_flag("--newton_stol", 1, "Relative tolerance in two continuous Newton iteration");
  myparser.register_flag("--newton_atol", 1, "Absolute tolerance in Newton iteration");
  myparser.register_flag("--gmres_restart", 1, "GMRES restart");
  myparser.register_flag("--newton_max_it", 1, "Maximum iterations in Newton");
  myparser.register_flag("--ksp_monitor", 0, "Monitoring KSP tolerance.");
  myparser.register_flag("--ksp_type", 1, "KSP type.");

  myparser.help();
 
  myparser.get("--mx", nx, 16);
  myparser.get("--my", ny, nx);
  myparser.get("--ksp_max_it", iter_max, 256);
  myparser.get("--ksp_atol", ksp_atol, 1.0e-50);
  myparser.get("--ksp_rtol", ksp_rtol, 1.0e-7);
  myparser.get("--newton_rtol", newton_rtol, 1.0e-7);
  myparser.get("--newton_stol", newton_stol, 1.0e-7);
  myparser.get("--newton_atol", newton_atol, 1.0e-50);
  myparser.get("--gmres_restart", restart, 60);
  myparser.get("--newton_max_it", kmax, 60);
  myparser.get("--ksp_monitor", monitor, 0);
  myparser.get("--ksp_type", ksp_type, std::string("gmres"));
  
  cout << "Relative tolerance for Newton iteration: " << newton_rtol << endl
       << "Absolute tolerance for Newton iteration: " << newton_atol << endl
       << "Secussive tolerance for Newton iteration: " << newton_stol << endl
       << "Relative tolerance for KSP iteraion: " << ksp_rtol << endl
       << "Absolute tolerance for KSP iteraion: " << ksp_atol << endl;

  const double hx = 1./(nx-1), hy = 1./(ny-1);
  const double lambda = 6.0;

  lamhxhy = lambda * hx * hy;

  // Create arrays U and F
  Array<double, 2> U(shape(nx, ny)), F(shape(nx, ny)), DX(shape(nx, ny));

  firstIndex ix;
  secondIndex jx;
  Range Irange(0, nx-2), Jrange(0, ny-2);
  
  // Initial guess
  U = 0.0;
  U(Irange, Jrange) = (lambda / (lambda+1.0)) * 
              sqrt(MIN(hy * MIN(jx, ny-jx-1), hx * MIN(ix,nx-ix-1)));
  
  identity_preconditioner p;
  apply_op f;			// Wraps up applyStencil

  f(U, F);
  double norm0 = itl::two_norm(F);
  double two_norm_F = 0.0, old_two_norm_F = 0.0;

  modified_gram_schmidt< Array<double, 2> > orth(restart, size(U));
  
  // Start Newton solver
  int k;
  for (k = 0; k < kmax; k++) {
    
    matrix_free_operator<Array<double, 2>, apply_op> A(f, U, F);
    noisy_iteration<double> iter(F, iter_max, ksp_rtol, ksp_atol);
    basic_iteration<double>& biter = static_cast<basic_iteration<double>&>(iter);
    //  Solve J dx = f  with matrix-free GMRES (identity precond)
    DX = 0.0;

    if (ksp_type == "gmres") {
      if ( monitor )
	gmres(A, DX, F, p, restart, iter, orth);
      else
	gmres(A, DX, F, p, restart, biter, orth);
	      
    } else if (ksp_type == "bicgstab") {
      if ( monitor )
        bicgstab(A, DX, F, p, iter);
      else 
        bicgstab(A, DX, F, p, biter);

    } else if (ksp_type == "cgs") {
      if ( monitor )
        cgs(A, DX, F, p, iter);
      else 
        cgs(A, DX, F, p, biter);

    } else if (ksp_type == "tfqmr") {
      if ( monitor )
        tfqmr(A, DX, F, p.left(), p.right(), iter);
      else 
        tfqmr(A, DX, F, p.left(), p.right(), biter);

    } else {
      std::cerr << "No suitable Krylov subspace method is selected." 
		<< std::endl;
	return 0;
    }

    U -= DX;

    //  Compute F
    f(U, F);

    old_two_norm_F = two_norm_F;
    two_norm_F = two_norm(F);


    cout << k << " Newton Iteration:      " << "|F|: " << two_norm_F << endl
	 << "*****number of GMRES iteration used: " << iter.iterations() 
	 << "     residual is " << iter.resid() << endl;

    if ( two_norm_F < newton_rtol * norm0 
	 || std::abs( old_two_norm_F - two_norm_F ) < newton_stol
	 || two_norm_F < newton_atol )
      break;
  }

  cout << "Total Number of Newton iteration: " << k+1 << endl;

  return 0;
}