#include "Sundance.h"
#include "NewtonSolver.h"
#include "BlockBacksolver.h"

/** \example inlinePoissonBoltzmann1D.cpp
 * Solve the Poisson-Boltzmann equation \f$\nabla^2 u = e^-u$ on the unit
 * line with boundary conditions
 * Left:  Dirichlet   u(0) = 1.0
 * Right: Dirichlet   u(1) = uRight = 2 log(cosh(1/sqrt(2)))
 * The solution is 2 log(cosh(x/sqrt(2))).
 *
 * We set up this problem in a contrived way for use as a model control problem. 
 * The boundary conditions are specified in terms a control variable alpha
 * u(0) = alpha(0)
 * u(1) = alpha(1)
 *
 * In an actual control problem, we would find alpha through
 * optimization. Here, we will supply an auxiliary algebraic equation
 * for alpha such that the boundary conditions alpha(0)=1,
 * alpha(1)=uRight are obtained. We can use
 *
 * alpha = uRight * x
 *
 * which has the correct solutions alpha(x=0)=0, alpha(x=1)=uRight.
 *
 * We'll group the variables into blocks [ u | alpha ], so that the resulting linear
 * system is a 2x2 block matrix problem

 *
 * [ K00   K01 ] [  du  ]   [ f0 ]
 * [           ] [      ] = [    ]
 * [ 0     K11 ] [dAlpha]   [ f1 ]
 *
 * Eliminating dAlpha, we have the system 
 *
 * K00 du = f0 - K01 K11^-1 f1
 *
 * which we can solve using the block backsolver.
 * */



int main(int argc, void** argv)
{
  try
    {
			
      Sundance::init(&argc, &argv);

      /* create a simple mesh on the unit line */
      double L=1.0;
      int nx = 10;
      int ny = 10;
      MeshGenerator mesher = new RectangleMesher(0.0, L, nx, 0.0, L, ny);
      Mesh mesh = mesher.getMesh();

      double uRight = 2.0*log(cosh(L/sqrt(2.0)));	

      /* define an expression representing the x-coordinate function */
      Expr x = new CoordExpr(0);

      /* create a cell set representing the right boundary */
      CellSet boundary = new BoundaryCellSet();
      CellSet left = boundary.subset( x == 0.0 );
      CellSet right = boundary.subset( x == L );

      /* create a discrete state space on the entire mesh */
      TSFVectorSpace discreteStateSpace
        = new SundanceVectorSpace(mesh, new Lagrange(1));
      /* create a discrete control space on the control points only */
      TSFVectorSpace discreteControlSpace 
        = new SundanceVectorSpace(mesh, new Lagrange(1), left+right);

      /* create an expression for the initial state. This will be reused as the
       * starting point for each newton step. Assume u(x)=x as an initial 
       * guess, and discretize it.
       */
      Expr u0 = new DiscreteFunction(discreteStateSpace, 
                                     x, "u0");

      /* create an expression for the initial control. This will be reused as the
       * starting point for each newton step. Assume alpha(x)=x as an initial 
       * guess, and discretize it.
       */
      Expr alpha0 = new DiscreteFunction(discreteControlSpace, 1.0, "alpha0");

      /* state test and unknown */
      Expr v = new TestFunction(new Lagrange(1), "v");
      Expr du = new UnknownFunction(new Lagrange(1), "du");

      /* control test and unknown */
      Expr beta = new TestFunction(new Lagrange(1), "beta");
      Expr dAlpha = new UnknownFunction(new Lagrange(1), "dAlpha");

      /* arrange test and unknowns into [state, control] blocks */
      TSFVectorType vectorType = new PetraVectorType();
      TSFArray<Block> unkBlocks = tuple(Block(du, vectorType), 
                                        Block(dAlpha, vectorType));
      TSFArray<Block> varBlocks = tuple(Block(v, vectorType), 
                                        Block(beta, vectorType));

      /* create a differential operator representing the x-derivative. */
      Expr dx = new Derivative(0);
      Expr dy = new Derivative(1);
      Expr grad = List(dx, dy);
			
      /* linearized weak equation for the step du */
      Expr linearizedStateEqn = (grad*(u0+du))*(grad*v) + v*exp(-u0) - v*exp(-u0)*du;

      /* linearized algebraic equation for the step in the control variable */
      Expr linearizedControlEqn 
        = beta*(alpha0 + dAlpha - uRight*x);

      /* integrate */
      Expr eqn = Integral(linearizedStateEqn) 
        + Integral(left+right, linearizedControlEqn);

      /* Dirichlet boundary condition */

      EssentialBC bc = EssentialBC(left+right, (u0+du-(alpha0 + dAlpha))*v) ;
			
      /* linear problem for the step [du, dAlpha] */
      StaticLinearProblem prob(mesh, eqn, bc, varBlocks, unkBlocks); 

      /* create linear solver for the block system */
      TSFPreconditionerFactory ilu1 = new ILUKPreconditionerFactory(1); 
      TSFLinearSolver stateSolver = new BICGSTABSolver(ilu1, 1.0e-13, 300);
      TSFLinearSolver controlSolver = new BICGSTABSolver(1.0e-14, 300);
			
      TSFLinearSolver blockSolver = new BlockBacksolver(stateSolver, 
                                                        controlSolver);
			
      /* create nonlinear solver and solve */
      NewtonLinearization newton(prob, List(u0, alpha0), blockSolver);
      Expr soln = newton.solve(NewtonSolver(blockSolver, 50, 1.0e-12, 1.0e-12));
			
      /* write the control variables */
      //			cerr << "control vars " << endl;
      FieldWriter writer = new MatlabWriter("u.dat");
      writer.writeField(soln[0]);

      // Added by RAB: 
      std::ofstream u_dump("u_dump.dat");
      std::ofstream alpha_dump("alpha_dump.dat");
      soln[0].matlabDump(u_dump);
      soln[1].matlabDump(alpha_dump);
			
      /* compare states to exact solution */
      Expr exactSoln = 2.0*log(cosh(x/sqrt(2.0)));

      /* compute the norm of the error */
      double errorNorm = (exactSoln - soln[0]).norm(2);
      double tolerance = 1.0e-4;
      TSFOut::printf("error = %g\n", errorNorm);

      Testing::passFailCheck(__FILE__, errorNorm, tolerance);
    }
  catch(exception& e)
    {
      Sundance::handleError(e, __FILE__);
    }
  Sundance::finalize();

}