/*
 * solves Poisson-Boltzmann equation div(grad(u)) = exp(-u) on a unit square
 * domain with u = 0 on the boundary; uses own nonlinear solver, not
 * Sundance's. The iterations step can be computed either via steepest descent
 * or Newton's method (to easy understanding, the Newton Krylov implementation
 * has been simplified to Newton).
 *
 * The intent of this code is to show how we can implement our own nonlinear
 * solver rather than rely solely on Sundance's Newton solver.
 *
 * by Alex Cunha, nov/2004
 *
 */

#include <math.h>
#include "Sundance.h"

#ifdef HAVE_NEW_IOSTREAMS
#include <fstream>                        // file stream objects
#include <iomanip>                        // stream manipulator objects
#include <iostream>                       // standard stream objects
#else
#include <fstream.h>
#include <iomanip.h>
#include <iostream.h>
#define std
#endif

//
// Constants
//

// step computation methods
#define STEEPEST_DESCENT 1
#define CG               2
#define NEWTON           3
#define NEWTON_KRYLOV    4

#define MAX_BACKTRACK   20                // max bactrack steps

#define SIG_DIGITS 12                     // significant digits in output


//
// Global variables
//

int nx = 32, ny = 32;                     // x & y grid divisons
int lmaxiter = 500;                       // max number of linear iterations
int nmaxiter = 200;                       // max number of newton iterations
int degree = 1;                           // interpolation degree
int precfill = 1;                         // fill level for preconditioner
int verb_level = 0;                       // solver verbosity level
double h;                                 // mesh size

Expr slabwidth = new ParameterExpr(3.0,"w"); // slab width as multiple of 'h'
double confidence = 0.99;                 // confidence interval
Expr alpha;                               // alpha for chi & delta functions

double bt_alpha = 1e-4;                   // alpha in backtracking tolerance
double linear_tol = 1e-12;                // linear solver tolerance
double abs_tol = 1e-6;                    // absolute tolerance
double rel_tol = 1e-6;                    // relative tolerance

bool plotIterates;                        // iterates plotting flag
int opt_method = NEWTON_KRYLOV;           // primary optimization method to use
string out_name;                          // prefix of output files

bool useobj;                              // flag to use obj function as merit


/*--------------------------------------------------------------------------
  FUNCTIONS
  --------------------------------------------------------------------------*/

void parseSwitches( int argc, char * argv[] )
{
  TSFCommandLine::findInt("-nx", nx);
  TSFCommandLine::findInt("-ny", ny);

  // set mesh size
  double hx = 1.0/(double) nx;
  double hy = 1.0/(double) ny;
  h = 0.5*(hx + hy);

  // plotting flag
  plotIterates = TSFCommandLine::find("-plot");

  // prefix for output file names  
  out_name = "nls";
  TSFCommandLine::findString("-o", out_name);

  // tolerances and max number of iterations
  TSFCommandLine::findDouble( "-atol", abs_tol );
  TSFCommandLine::findDouble( "-rtol", rel_tol );
  TSFCommandLine::findDouble( "-ltol", linear_tol);
  TSFCommandLine::findInt( "-iter", lmaxiter);
  TSFCommandLine::findInt( "-ni", nmaxiter);

  // step computation method
  if ( !TSFCommandLine::findInt("-m", opt_method) )
    TSFCommandLine::findInt("-method", opt_method);
  if ( opt_method < 1 || opt_method > 4 )
    TSFError::raise("Error: Invalid method in -m or -method\n");

  // preconditioner filling
  TSFCommandLine::findInt("-fill", precfill);

  // element interpolation degree
  TSFCommandLine::findInt("-d", degree);
  
  // determine alpha value used in heaviside and delta functions
  alpha = std::log((1.0+confidence)/(1.0-confidence)) / (slabwidth * h);

  // determine if using objective function as merti function on line search
  // default is to use it.  
  useobj = !TSFCommandLine::find("-noobj");
}

void plotExpr(const Expr & E, const string & filename, int digits = SIG_DIGITS,
              const TSFVectorSpace * space = NULL )
{
  Expr discreteExpr = (space != NULL) ? new DiscreteFunction(*space, E) : E;

  ofstream outfile( filename.c_str() );
  //outfile.setf( ios::scientific );
  outfile.precision( digits );
  discreteExpr.matlabDump( outfile );
}

void plotGridValues( const Expr & expr, const string & filename,
                     int digits = SIG_DIGITS )
{
  Mesh mesh;
  TSFVector vector;
  TSFSmartPtr<DOFMapBase> map;
  
  expr.getDOFMap(map);
  expr.getMesh(mesh);
  expr.getVector(vector);

  ofstream os( filename.c_str() );
  os.precision( digits );
  
  int index = expr.getReducedIndex();
  int nCell = mesh.numCells(0);
  
  for (int j = 0; j < nCell; j++) {
    int dof_index = map->guaranteedPointLookup(j, index);
    const Point & x = mesh.getCell(0,j).position();
    os << x[0] << " " << x[1] << " " << vector[dof_index] << endl;
  }
}


double computeTol( int nl_iter, double grad_norm, double previous_norm,
                   double nonlinear_tol )
{
  // ...for simplicity...
  return linear_tol;
}

void computeStep( TSFVector & step, StaticLinearProblem & prob, int nl_iter,
                  double grad_norm, double previous_norm, double non_tol )
{
  TSFLinearSolver solver;
  TSFPreconditionerFactory prec;
  Expr soln;
  
  //prec = new ILUKRightPreconditionerFactory( precfill );
  prec = new ILUKPreconditionerFactory( precfill );
  
  //solver = new GMRESSolver( prec, linear_tol, lmaxiter );
  solver = new BICGSTABSolver( prec, linear_tol, lmaxiter );
  solver.setVerbosityLevel( verb_level );

  //--------------------------------------------------------------------
  // solve H*p = -g system to determine step
  //--------------------------------------------------------------------
  
  if ( opt_method == STEEPEST_DESCENT ) {
    // descent step is negative gradient direction.
    // obs: getRHS() returns -g which is fine since in the
    //      update newpoint = oldpoint + step_length * (-g)

    step = prob.getRHS();                 // step = -g
    double step_len = 1.0;                // can be much smaller !
    step.scalarMult( step_len, step );
  }
  else if ( opt_method == NEWTON ) {
    // compute exact, full newton step
    soln = prob.solve( solver );
    soln.getVector( step );
  }
  else if ( opt_method == NEWTON_KRYLOV ) {
    // determine tolerance for linear solver
    double tolerance = computeTol( nl_iter, grad_norm,
                                   previous_norm, non_tol );

    //solver = new GMRESSolver( prec, tolerance, lmaxiter );
    solver = new BICGSTABSolver( prec, tolerance, lmaxiter );
    solver.setVerbosityLevel( verb_level );

    // determine inexact step
    soln = prob.solve( solver );
    soln.getVector( step );
  }
  else {
    TSFError::raise("Error: Invalid optimization method\n");
  }
}



int main( int argc, char ** argv )
{
  try {

    bool verbose = true;
    
    //-------------------------------------------------------------------
    // Initial settings, mesh, and boundary
    //-------------------------------------------------------------------
    
    Sundance::init( &argc, &(void**)argv );
    TSFCommandLine::init( argc, (void**)argv );
    TSFCommandLine::verbose() = false;

    // read command line switches
    parseSwitches( argc, argv );
    if ( verbose ) TSFOut::printf("\t >>> command line has been parsed\n");

    // x, y and gradient operator
    Expr x = new CoordExpr(0, "x");
    Expr y = new CoordExpr(1, "y");
    Expr dx = new Derivative(0,1);
    Expr dy = new Derivative(1,1);
    Expr grad = List( dx, dy );
    
    // create uniform mesh in rectangular domain
    MeshGenerator mesher = new RectangleMesher( 0.0, 1.0, nx, 0.0, 1.0, ny );
    Mesh mesh = mesher.getMesh();
    
    // get outer boundary cells
    CellSet boundary = new BoundaryCellSet();

    // quadrature integration order
    QuadratureFamily gaussQuad  = new GaussianQuadrature( 2 );
    Defaults::setQuadratureFamily( gaussQuad );

    //-------------------------------------------------------------------
    // Build discrete space and initial guess
    //-------------------------------------------------------------------
      
    // create a discrete space on the mesh for the inverse problem
    BasisFamily lagr = new Lagrange( degree );
    TSFVectorSpace space = new SundanceVectorSpace( mesh, lagr );

    // unknown (newton step) and its variation
    Expr uTil = new UnknownFunction( lagr, "uTil" );
    Expr uHat = new TestFunction( lagr, "uHat" );

    // initialize unknown iterate
    //Expr u0 = DiscreteFunction::discretize( space, x+y );
    Expr u0 = DiscreteFunction::discretize( space, Expr(1.0) );
    
    //-------------------------------------------------------------------
    // Define linear functional (hand derived) in weak form and linear
    // boundary conditions. Note that we only use differentials and
    // variations (test functions defined above).
    //-------------------------------------------------------------------

    Expr L_u  = (grad * u0)*(grad * uHat) + exp(-u0) * uHat;
    Expr L_uu = (grad * uTil)*(grad * uHat) - exp(-u0) * uHat * uTil;
    
    Expr linear_weak = Integral( L_u + L_uu );

    // BC: u = 0 on boundary
    EssentialBC bc = EssentialBC( boundary, uTil * uHat );
    
    //-------------------------------------------------------------------
    // Build linear problem object
    //-------------------------------------------------------------------
    
    StaticLinearProblem prob( mesh, linear_weak, bc, uHat, uTil );

    //-------------------------------------------------------------------
    // Solve the problem
    //-------------------------------------------------------------------
    
    //
    // set an inexact newton solver by doing repetitive calls to 
    // Sundance linear solver and performing at most 'max_backtrack'
    // backtracking iterations
    //
    
    if ( verbose ) TSFOut::printf("\t >>> solving nonlinear problem\n");
        
    TSFVector newPoint;                   // current iterate
    TSFVector oldPoint;                   // previous iterate
    TSFVector step;                       // newton step
    
    int nl_iter = 0;                      // nonlinear iterations counter
    int bt_iter = 0;                      // backtracking iterations counter
    
    bool converged = true;                // flag for nonlinear convergence
    double new_grad_norm;                 // new grad norm in backtracking
    double step_length;                   // length of newton step
    double nl_tol;                        // nonlinear convergence tolerance
    double grad_norm;                     // L2 norm of lagrangian gradient
    double previous_norm;                 // dito

    double obj_val, new_obj_val;          // objective functional values
    
    // the expression below is the objective function we use in this code.
    // note that it is the first variation applied at u0; in your segmentation
    // problem obj would be the functional expression at the current iterate
    Expr obj = (grad * u0)*(grad * u0) + exp(-u0) * u0; // NEW
    
    //------------------------------------------------------
    // prepare to start nonlinear iterations
    //------------------------------------------------------
    
    u0.getVector( newPoint );             // get vector we will be updating
    grad_norm = prob.getRHS().norm2();    // L2 norm of gradient
    previous_norm = grad_norm;            // save gradient norm

    // nonlinear tolerance
    nl_tol = abs_tol + rel_tol * grad_norm;

    // current objective function value
    obj_val = obj.integral( mesh ); // NEW

    //------------------------------------------------------
    // start nonlinear iterations
    //------------------------------------------------------

    while ( grad_norm > nl_tol && nl_iter < nmaxiter ) {

      // we changed problem parameters so we MUST flush the system
      prob.flushMatrixValues();
      
      // save current iterates
      oldPoint = newPoint.copy();

      // compute Newton step
      computeStep( step, prob, nl_iter, grad_norm, previous_norm, nl_tol );
      
      if ( plotIterates ) {
        // save current iterate on file
        // char sufix[8];
        // sprintf( sufix, "%03d", nl_iter );
        // plotGridValues( phi0, out_name + ".u." + sufix + ".2D.dat" );
      }
      
      //------------------------------------------------------
      // start backtracking
      //------------------------------------------------------
      
      bt_iter = 0;                        // counter
      step_length = 1.0;                  // step length
      bool decreased;                     // NEW - backtracking decrease flag
      
      do {

        // propose a new point
        newPoint = oldPoint + step_length * step;
        prob.flushMatrixValues();

        // get the new gradient vector and compute its norm
        TSFVector grad_vector = prob.getRHS();
        new_grad_norm = grad_vector.norm2();
        
        if ( useobj ) {
          // we are using the objective function as a descent criterion

          // check if we are indeed getting a descent direction
          double descent = - grad_vector.dot( step );
          if ( descent > 0.0 )
            TSFOut::printf("\t WARNING: direction is not descent!\n");

          // compute the new objective function value
          new_obj_val = obj.integral( mesh );

          // check for decrease
          decreased = ( new_obj_val - obj_val <
                        bt_alpha * step_length * descent );
          
          if ( verbose ) {
            if ( !decreased )
              TSFOut::printf( "\t\t>> backtrack %2d : objNew = %e\n",
                              bt_iter, new_obj_val );
            else if ( bt_iter )
              TSFOut::printf( "\t\t>> backtrack %2d : objNew = %e(< %e)\n",
                              bt_iter, new_obj_val, obj_val );
          }
        }
        else {
          // we are using the gradient norm as descent criterion
          
          double bt_tol = (1.0 - bt_alpha * step_length) * grad_norm;
          decreased = new_grad_norm < bt_tol;
          
          if ( verbose ) {
            if ( new_grad_norm > bt_tol )
              TSFOut::printf( "\t\t>> backtrack %2d : gNew = %e\n",
                              bt_iter, new_grad_norm );
            else if ( bt_iter )
              TSFOut::printf( "\t\t>> backtrack %2d : gNew = %e (< %e)\n",
                              bt_iter, new_grad_norm, bt_tol );
          }
        }
        
        step_length *= 0.5;
        bt_iter++;
        
      } while ( !decreased && bt_iter < MAX_BACKTRACK );
      
      //------------------------------------------------------
      // end backtracking iterations
      //------------------------------------------------------
      
      if ( bt_iter == MAX_BACKTRACK ) {
        // backtracking failed; halt nonlinear loop
        if ( verbose ) TSFOut::printf( "\t\t>> Backtracking FAILED\n" );
        
        break;
      }
      else {
        // backtracking succeded: a new iterate was found.
        // update norm values and counters

        previous_norm = grad_norm;
        grad_norm = new_grad_norm;
        obj_val = new_obj_val;            // NEW
        nl_iter++;

        // compute new values for the c1 and c2 variables.
        // c1 and c2 MUST be ParameterExpr expressions
        
        // double c1_value = ...
        // double c2_value = ...
        // c1.setParameterValue( c1_value );
        // c2.setParameterValue( c2_value );
      }
    }
    
    //------------------------------------------------------
    // end nonlinear iterations
    //------------------------------------------------------
    
    if ( grad_norm > nl_tol ) {
      if ( verbose ) TSFOut::printf( "\t >> Optimization FAILED!\n" );
    }
    else {
      if ( verbose ) {
        // convergence was achieved
        TSFOut::printf("\t >>> ------------------------------------\n");
        TSFOut::printf("\t >>> Optimization converged with |grad| %e ",
                       grad_norm );
        printf("after %d iterations\n", nl_iter );
        TSFOut::printf("\t >>> ------------------------------------\n");
      }
      converged = true;
    }
    
    if ( plotIterates ) {
      // save last iterate on file
      // char sufix[8];
      // sprintf( sufix, "%03d", nl_iter );
      // plotGridValues( u0, out_name + ".u." + sufix + ".2D.dat" );
    }

    // save results
    if ( converged ) {
      string filename = out_name + ".u.2D.dat";
      plotGridValues( u0, filename );
      TSFOut::printf("\t >>> results saved on %s file\n", filename.c_str() );
    }
    
  } // try...
  
  catch ( exception& e ) {
    TSFOut::printf( e.what() );
    Sundance::handleError(e, __FILE__);
  }
  
  //Sundance::finalize();
}