#include "Sundance.h"

/** \example vorticityLDC.cpp
 * Solution of the lid-driven cavity problem using a vorticity-streamfunction
 * formulation of the Navier-Stokes equations. 
 *
 */

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 vorticity and stream function */
      Expr varOmega = new TestFunction(new Lagrange(1));
      Expr dOmega = new UnknownFunction(new Lagrange(1));
      Expr varPsi = new TestFunction(new Lagrange(1));
      Expr dPsi = new UnknownFunction(new Lagrange(1));
			
      /* set up discrete function space. This problem is big enough that
       * we will use Petra */
      BasisFamily lagr2 = new Lagrange(1);
      TSFVectorSpace discreteSpace 
        = new SundanceVectorSpace(mesh, tuple(lagr2, lagr2), 
                                  new PetraVectorType());
			
      Expr zero = List(Expr(0.0), Expr(0.0));
      Expr u0 = DiscreteFunction::discretize(discreteSpace, zero);

      Expr psi0 = u0[0];
      Expr omega0 = u0[1];

      Expr psi = psi0 + dPsi;
      Expr omega = omega0 + dOmega;
      Expr du = List(dPsi, dOmega);
      Expr varU = List(varPsi, varOmega);
			

      /* 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);

      Expr velocity = List(dy*psi, -dx*psi);
      Expr velocity0 = List(dy*psi0, -dx*psi0);


      /* 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 vorticityEqn = -(grad*varOmega)*(grad*omega)
        - reynolds*varOmega*((velocity*grad)*omega0 + (velocity0*grad)*dOmega);

      Expr streamfunctionEqn = -(grad*varPsi)*(grad*psi) - varPsi*omega;

      Expr eqn = Integral(vorticityEqn) + Integral(streamfunctionEqn)
        + Integral(top, -varPsi);

      /*
       * Boundary conditions: 
       * psi=0 on all surfaces
       * d(psi)/dn = 0 on left, right, bottom
       * d(psi)/dn = 1 on top
       * Note that we apply the BCs on psi in the row space of Omega.
       */
      EssentialBC bc = EssentialBC(top, varOmega*psi)
        && EssentialBC(bottom, varOmega*psi)
        && EssentialBC(left, varOmega*psi)
        && EssentialBC(right, varOmega*psi);


      /* Create a stablized biconjugate gradient solver */
      TSFPreconditionerFactory prec = new ILUKPreconditionerFactory(1);
      TSFLinearSolver solver = new BICGSTABSolver(prec, 1.0e-14, 1000);

      /* Package the linearized problem into an object */
      StaticLinearProblem linearizedProblem(mesh, eqn, bc, varU, du,
                                            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 = 6;
      for (int r=0; r<nr; r++)
        {
          /* set the next value of the reynolds number */
          double Re = pow(1000.0, (double) r / ((double) nr-1));
          reynolds.setParameterValue(Re);
					
          NewtonLinearization newton(linearizedProblem, u0, solver);
          NewtonSolver newtonSolver(solver, 150, 1.0e-12, 1.0e-12);
          Expr soln = newton.solve(newtonSolver);
					
          /*
            write the solution in a form readable by matlab
          */
          FieldWriter psiWriter = new MatlabWriter("psi.dat");
          psiWriter.writeField(soln[0]);
					
          FieldWriter omegaWriter = new MatlabWriter("omega.dat");
          omegaWriter.writeField(soln[1]);
					
        }
    }
  catch(exception& e)
    {
      Sundance::handleError(e, __FILE__);
    }
  Sundance::finalize();
			
}