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

/**
 * \example timeStepHeat1D.cpp 
 *
 * This example shows how to do timestepping in Sundance. We solve the
 * transient heat equation in one dimension using Crank-Nicolson time
 * discretization. The time discretization is done at the symbolic level.
 * Spatial discretization is done via StaticLinearProblem, yielding system
 * matrices and vectors that can be used to march the problem in time. In this
 * example we march with a simple constant timestep loop; for more difficult
 * problems one would use a high-quality DAE/ODE integration code such 
 * as DASSLQ or PVODE.
 *
 * We solve the heat equation u_xx = u_t with boundary conditions
 * u(0)=u(1)=0 and initial conditions u(x,t=0)=sin(pi x). The solution
 * is u(x,t)=exp(-pi^2 t) sin(pi x).
 */

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

      /* create a simple mesh on the unit line */
      int n = 100;
      MeshGenerator mesher = new LineMesher(0.0, 0.5, n);
      Mesh mesh = mesher.getMesh();

      /* create unknown and variational functions */
      Expr delU = new TestFunction(new Lagrange(1));
      Expr U = new UnknownFunction(new Lagrange(1));

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

      /* the initial conditions will be u0(x,t=0) = sin(pi*x).
       * create a coordinate expression to represent x, then
       * create sin(pi*x), and then project it onto a discrete function. */
      Expr x = new CoordExpr(0);

      double pi = 4.0*atan(1.0);

      TSFVectorSpace discreteSpace 
        = new SundanceVectorSpace(mesh, new Lagrange(1));

      Expr u0 = new DiscreteFunction(discreteSpace, sin(pi*x));


      /* 
         set up crank-nicolson stepping with timestep = 0.02. The time
         discretization is done at the symbolic level, yielding
         an elliptic problem that we solve repeatedly for the updated
         solution at each time level. 
      */

      double deltaT = 0.02;
      Expr cnStep = delU*(U - u0) + deltaT*(dx*delU)*(dx*(U + u0)/2.0);
      Expr eqn = Integral(cnStep);

      /* Define BCs to be zero at both ends */
      CellSet boundary = new BoundaryCellSet();
      CellSet left = boundary.subset( fabs(x - 0.0) < 1.0e-10 );
      CellSet right = boundary.subset( fabs(x - 1.0) < 1.0e-10 );
      EssentialBC bc = EssentialBC(left, delU*U);
			
      /* create a solver object */
      TSFPreconditionerFactory prec = new ILUKPreconditionerFactory(1);
      TSFLinearSolver solver = new BICGSTABSolver(prec, 1.0e-14, 300);
      solver.setVerbosityLevel(0);

      /* 
         put the time-discretized eqn into a StaticLinearProblem object
         which will do the spatial discretization.
      */
      StaticLinearProblem prob(mesh, eqn, bc, delU, U);

      /* 
         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 nSteps = 100;
      for (int i=0; i<nSteps; i++)
        {
          /* solve the problem */
          Expr soln = prob.solve(solver);
          TSFVector solnVec;
          soln.getVector(solnVec);
          u0.setVector(solnVec);
          cerr << "eqn = " << cnStep << endl;
          cerr << "u0 = " << u0 << endl;
          /* write the solution at step i to a file */
          char fName[20];
          sprintf(fName, "timeStepHeat%d.dat", i);
          ofstream of(fName);
          FieldWriter writer = new MatlabWriter(fName);
          writer.writeField(u0);
          cerr << "[" << i << "]";
          /* flush the matrix and RHS values */
          prob.flushMatrixValues();
        }
      cerr << endl;

      /* compute the exact solution and the error */
      double tFinal = nSteps * deltaT;
      Expr exactSoln = exp(-pi*pi*tFinal) * sin(pi*x);
			
      /*
        compute the norm of the error
      */
      double errorNorm = (exactSoln-u0).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();
}