#include "Sundance.h"

/** \example vorticityLDC.cpp
 * Solution of the Navier-Stokes equation on the Sandia thunderbird
 */

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

      /* create a simple mesh on the rectangle */
      int nx = 32;
      int ny = 32;
      MeshGenerator mesher = new RectangleMesher(0.0, 1.0, nx, 0.0, 1.0, ny);
      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 == 1.0 );

      /* define variations and unknowns for U, V, and P
       * using the same basis functions for velocity and pressure */
      Expr vx = new TestFunction(new Lagrange(1));
      Expr dux = new UnknownFunction(new Lagrange(1));
      Expr vy = new TestFunction(new Lagrange(1));
      Expr duy = new UnknownFunction(new Lagrange(1));
      Expr q = new TestFunction(new Lagrange(1));
      Expr dp = new UnknownFunction(new Lagrange(1));

      /* set up discrete function space. This problem is big enough that
       * we will use Petra */
      BasisFamily lagr = new Lagrange(1);
      TSFVectorSpace discreteSpace 
        = new SundanceVectorSpace(mesh, tuple(lagr, lagr, lagr), 
                                  new PetraVectorType());
			
      Expr zero = List(Expr(0.0), Expr(0.0), Expr(0.0));
      Expr s0 = DiscreteFunction::discretize(discreteSpace, zero);
      Expr u0 = List(s0[0], s0[1]);
      Expr p0 = s0[2];

      Expr ux = u0[0] + dux;
      Expr uy = u0[1] + duy;
      Expr u = List(ux, uy);
      Expr v = List(vx, vy);
      Expr p = p0 + dp;

      /* create differential operators for x and y directions, and
       * then form gradient operator. */
      Expr dx = new Derivative(0);
      Expr dy = new Derivative(1);
      Expr grad = List(dx, dy);

      /* Represent the Reynolds number as a ParameterExpr. This lets
       * us change its value inside an existing problem, which we will
       * do in the course of stepping up from small reynolds numbers */
      Expr reynolds = new ParameterExpr(10.0);

      /* write the linearized weak equations */
      Expr momentumEqn = -(grad*ux)*(grad*vx) - (grad*uy)*(grad*vy)
        + reynolds*p*(dx*vx + dy*vy) - reynolds*v*(u0*grad)*u 
        - reynolds*v*(dux*dx + duy*dy)*u0; 


      Expr beta = new ParameterExpr(0.02);
      Expr h = new CellDiameterExpr();

      Expr continuityEqn = -q*(grad*u) - beta*h*h*(grad*p)*(grad*q);

      Expr eqn = Integral(momentumEqn) + Integral(continuityEqn);

      /*
       * Boundary conditions
       */
      EssentialBC bc = EssentialBC(left + right + bottom, u*v)
        && EssentialBC(top, vx*(ux-1.0) + vy*uy);

      /* Create a stablized biconjugate gradient solver */
      TSFPreconditionerFactory prec = new ILUKPreconditionerFactory(1);
      TSFLinearSolver solver = new BICGSTABSolver(prec, 1.0e-10, 5000);
      solver.setVerbosityLevel(4);

      /* Package the linearized problem into an object */
      StaticLinearProblem linearizedProblem(mesh, eqn, bc, 
                                            List(vx, vy, q),
                                            List(dux, duy, dp),
                                            new PetraVectorType());

			
      /* Continuation loop: to improve convergence, we solve a sequence of
       * problems at increasing Reynolds number, using the solution to
       * each as an initial guess to the next. */
      int nr = 4;
      for (int r=0; r<nr; r++)
        {
          /* set the next value of the reynolds number */
					
          double Re = 100.0 * ((double) r)/((double) nr-1);
          TSFOut::printf("doing Re=%g\n", Re);
          reynolds.setParameterValue(Re);
					
          NewtonLinearization newton(linearizedProblem, s0, solver);
          NewtonSolver newtonSolver(solver, 150, 1.0e-6, 1.0e-6);
          Expr soln = newton.solve(newtonSolver);
					
          /*
            write the solution in a form readable by matlab
          */
          FieldWriter uWriter = new MatlabVectorWriter("uVec.dat");
          uWriter.writeField("u", List(soln[0], soln[1]));
        }
    }
  catch(exception& e)
    {
      Sundance::handleError(e, __FILE__);
    }
  Sundance::finalize();
			
}