#include "Sundance.h"

/** \example heat2D.cpp
 * Solve Poisson's equation with a unit source term on the 
 * rectangle [0,1] x [0, 2] with the following boundary conditions:
 *
 * Left:   Natural, du/dx = 0
 * Bottom: Dirichlet, u= 0.5 x^2
 * Right:  Robin, u + du/dx = 3/2 + y/3
 * Top:    Neumann, du/dy = 1/3
 *
 * The solution is u(x,y) = 0.5*x^2 + y/3.
 *
 * This problem can be solved exactly in the space of second-order polynomials.
 */

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

      /* create a simple mesh on the rectangle */

      int nx = 32;
      int ny = 32;
      int npx = 1;
      int npy = 1;
      
      int fill=1;
      int overlap=1;
      
      double tol = 1.0e-12;
      int maxiter = 2000;
      
      TSFCommandLine::findInt("-npx", npx);
      TSFCommandLine::findInt("-npy", npy);
      TSFCommandLine::findInt("-nx", nx);
      TSFCommandLine::findInt("-ny", ny);
      TSFCommandLine::findInt("-maxiter", maxiter);
      TSFCommandLine::findInt("-fill", fill);
      TSFCommandLine::findInt("-overlap", overlap);

      TSFOut::printf("npx=%d npy=%d\n", npx, npy);

      bool useAztec = TSFCommandLine::find("-aztec");
      TSFCommandLine::findDouble("-tol", tol);

      MeshGenerator mesher 
        = new PartitionedRectangleMesher(0.0, 1.0, npx*nx, npx, 
                                         0.0, 2.0, npy*ny, npy);
      Mesh mesh = mesher.getMesh();
      

      /* define coordinate functions for x and y coordinates */
      Expr x = new CoordExpr(0);
      Expr y = new CoordExpr(1);

      /* define cells sets for each of the four sides of the rectangle */
      CellSet boundary = new BoundaryCellSet();
      CellSet left = boundary.subset( x == 0.0 );
      CellSet right = boundary.subset( x == 1.0 );
      CellSet bottom = boundary.subset( y == 0.0 );
      CellSet top = boundary.subset( y == 2.0 );
      
      /* Create a vector space factory, used to 
       * specify the low-level linear algebra representation */
      TSFVectorType petra = new PetraVectorType();
      
      /* create a discrete space on the mesh */
      TSFVectorSpace discreteSpace = new SundanceVectorSpace(mesh, new Lagrange(2), petra);


      /* We'll use a discrete function to represent the 
       * source term, providing a test
       * of our ability to evaluate discrete functions on maximal cells */
      Expr f = new DiscreteFunction(discreteSpace, 1.0);

      /* We'll use a discrete function to represent the imposed
       * boundary value on the right-hand boundary.  This provides a
       * test of our ability to evaluate discrete functions on
       * lower-dimensional cells. */
      Expr rightBCExpr = new DiscreteFunction(discreteSpace, 1.5 + y/3.0);


      /* create symbolic objects for test and unknown functions */
      Expr v = new TestFunction(new Lagrange(2));
      Expr u = new UnknownFunction(new Lagrange(2));

      /* create symbolic differential operators */
      Expr dx = new Derivative(0,1);
      Expr dy = new Derivative(1,1);
      Expr grad = List(dx, dy);

      /* Write symbolic weak equation and Neumann and Robin BCs */
      Expr poisson = Integral(-(grad*v)*(grad*u)  - f*v, new GaussianQuadrature(2))
        + Integral(top, v/3.0) + Integral(right, v*(rightBCExpr - u));


      /* Write essential BCs: 
       * Bottom: u=x^2
       */
      EssentialBC bc = EssentialBC(bottom, v*(u - 0.5*x*x), 
                                   new GaussianQuadrature(4));

      /* Assemble everything into a problem object, with a specification that
       * Petra be used as the low-level linear algebra representation */
      StaticLinearProblem prob(mesh, poisson, bc, v, u, petra);
      
      /* create a preconditioner and solver */
      TSFHashtable<int, int> azOptions;
      TSFHashtable<int, double> azParams;
      TSFLinearSolver solver;

      if (useAztec)
        {
          azOptions.put(AZ_solver, AZ_gmres);
          azOptions.put(AZ_precond, AZ_dom_decomp);
          azOptions.put(AZ_subdomain_solve, AZ_ilu);
          azOptions.put(AZ_graph_fill, fill);
          azParams.put(AZ_tol, tol);
          azOptions.put(AZ_max_iter, maxiter);
          solver = new AZTECSolver(azOptions, azParams);
        }
      else
        {
          TSFPreconditionerFactory precond 
            = new ILUKPreconditionerFactory(fill, overlap);
          solver = new BICGSTABSolver(precond, tol, maxiter);
        }

      /* solve the problem and return the solution as a symbolic object */
      Expr soln = prob.solve(solver);
      Expr exactSoln = 0.5*x*x + y/3.0;

      /* write to matlab */
      FieldWriter writer = new MatlabWriter("heat2D.dat");
      writer.writeField(soln);



      /* compare to exact solution */
      Expr error = new DiscreteFunction(discreteSpace, soln-exactSoln);
      FieldWriter writer2 = new MatlabWriter("error.dat");
      writer2.writeField(error);  
      /* compare to known solution */


      // compute the norm of the error
      double errorNorm = (soln-exactSoln).norm(2);
      double tolerance = 1.0e-9;
        
      Testing::passFailCheck(__FILE__, errorNorm, tolerance);
    }
  catch(exception& e)
    {
      Sundance::handleError(e, __FILE__);
    }
  Sundance::finalize();
}