#include "Sundance.h"

/**
 * \example burgers1D.cpp
 * Solves the time-dependent Burgers equation
 *
 * u_t + u u_x = u_xx
 *
 * with Crank-Nicolson timestepping. Each timestep involves the solution
 * of a nonlinear PDE in space, which we solve using Newton's method.
 * 
 * The problem domain is [-pi/2, pi/2]. We choose initial conditions 
 * 
 * u(x, t=0) = -2.0*epsilon*cos(x) / (1.0 + epsilon*sin(x))
 *
 * and boundary conditions
 *
 *  u(-pi/2) = u(pi/2) = 0,
 * 
 * giving the solution
 *
 * u(x,t) = -2.0*epsilon*cos(x)*exp(-t) / (1.0 + epsilon*sin(x)*exp(-t))
 * 
 * See, e.g., Whitham's _Linear and Nonlinear Waves_ for methods of
 * obtaining exact solutions to Burgers equation.
 */

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

      /* mesh the problem domain */
      int n = 16;
      const double pi = 4.0*atan(1.0);
      MeshGenerator mesher = new LineMesher(-pi/2.0, pi/2.0, n);
      Mesh mesh = mesher.getMesh();

      /* create unknown and variational functions */
      Expr u = new UnknownFunction(new Lagrange(2),"u");
      Expr v = u.variation();

      /* create a parameter expression for the time. This must be a parameter
       * rather than a constant since it will change in the course of the
       * calculation */
      Expr t = new ParameterExpr(0.0);

      /* create a differentiation operator */
      Expr dx = new Derivative(0);

      /* create a coordinate expression */
      Expr x = new CoordExpr(0);

      /* create an expr for the exact solution. We've used the Cole-Hopf
       * transformation to obtain a solution of Burgers equation from
       * a solution of the heat equation. */
      Expr epsilon = 0.9;
      Expr phi0 = 1.0 + epsilon*sin(x)*exp(-t);
      Expr exactSoln = -2.0*(dx*phi0)/phi0;

      /* discretize the exact solution at t=0. Use it as a starting value
      * for both the time marching and the newton iteration. */
      TSFVectorSpace discreteSpace 
        = new SundanceVectorSpace(mesh, new Lagrange(2));
      /* value of u at the previous newton step */
      Expr uPrevNewton = new DiscreteFunction(discreteSpace, exactSoln);
      /* value of u at the previous time step */
      Expr uPrevT = new DiscreteFunction(discreteSpace, uPrevNewton);

      /* create a parameter for the timestep. In this problem, we use
       * constant timestepping so we could just as well use a constant. 
       * However, in general we will want to use a parameter for deltaT
       * to allow adaptive timestepping. */
      Expr deltaT = new ParameterExpr(0.01, "dt");

      /* define the average of the current and previous time levels.
       * This will be used in the RHS terms of the Crank-Nicolson 
       * formula */
      Expr uAvg = (u + uPrevT)/2.0;

      /* write out the Crank-Nicolson time discretization */
      Expr nonlinEqn = Integral(v*(u - uPrevT)/deltaT + v*uAvg*(dx*uAvg)
                                + (dx*v)*(dx*uAvg));
      
      /* linearize about the previous newton step */
      Expr linEqn = nonlinEqn.linearization(u, uPrevNewton);
      
      /* Define BCs to be zero at both ends */
      CellSet boundary = new BoundaryCellSet();
      CellSet left = boundary.subset( fabs(x + pi/2.0) < 1.0e-10 );
      CellSet right = boundary.subset( fabs(x - pi/2.0) < 1.0e-10 );

      EssentialBC nonLinBC = EssentialBC(left, u*v)
        && EssentialBC(right, u*v);

      /* linearize the BCs. Although the BC expressions are already linear,
       * we must go through the linearization process to convert the
       * BCs to equations involving the differential du */
      EssentialBC linBC = nonLinBC.linearization(u,uPrevNewton);


      /* create a solver object */
      TSFPreconditionerFactory prec = new ILUKPreconditionerFactory(1);
      TSFLinearSolver solver = new BICGSTABSolver(prec, 1.0e-14, 300);
      solver.setVerbosityLevel(0);
      NewtonSolver newtonSolver(solver, 100, 1.0e-10, 1.0e-10);

      /* 
         put the time-discretized eqn into a StaticLinearProblem object
         which will do the spatial discretization.
      */
      StaticLinearProblem prob(mesh, linEqn, linBC, v, u.differential());

      NewtonLinearization newton(prob, uPrevNewton, solver); 
      /* 
         Now, loop over timesteps, solving the elliptic problem for u at each
         step. At the end of each step, assign the solution solnU into u0.
         Because Exprs are stored by reference, the updating of u0 propagates
         to the copies of u0 in the equation set and in the 
         StaticLinearProblem. The same StaticLinearProblem can be reused
         at all timesteps.
      */
      int i=0;
      while (t.value() < 2.0)
        {
          /* solve the problem */
          Expr soln = newton.solve(newtonSolver);
          /* put the solution into the functions for the prev timestep and
           * newton step. IMPORTANT: we must use two different objects
           * for the prev timestep and prev newton step */
          TSFVector solnVec;
          soln.getVector(solnVec);
          uPrevT.setVector(solnVec.copy()); 
          uPrevNewton.setVector(solnVec); 
          /* update the time */
          t.setParameterValue(t.value() + deltaT.value());
          /* write the solution at step i to a file */
          char fName[20];
          sprintf(fName, "timeStep%d.dat", i);
          ofstream of(fName);
          FieldWriter writer = new MatlabWriter(fName);
          writer.writeField(uPrevT);
          i++;
        }

      /*
        compute the norm of the error
      */
      double errorNorm = (exactSoln-uPrevT).norm(2);
      double tolerance = 1.0e-4;

      /*
        decide if the error is within tolerance
      */
      Testing::passFailCheck(__FILE__, errorNorm, tolerance);
    }
  catch(exception& e)
    {
      Sundance::handleError(e, __FILE__);
    }
  Sundance::finalize();
}