#include "Sundance.h"

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 delOmega = new TestFunction(new Lagrange(1), "delOmega");
      Expr omega = new UnknownFunction(new Lagrange(1), "omega");
      Expr delPsi = new TestFunction(new Lagrange(1), "delPsi");
      Expr psi = new UnknownFunction(new Lagrange(1), "psi");

      BasisFamily lagr2 = new Lagrange(1);
      TSFVectorSpace discreteSpace 
        = new SundanceVectorSpace(mesh, tuple(lagr2, lagr2));
			
      Expr zero = List(Expr(0.0), Expr(0.0));
      Expr u0 = DiscreteFunction::discretize(discreteSpace, zero);

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

      Expr velocity0 = List(dy*u0[0], -dx*u0[0]);


      double Re = 1000.0;

      Expr vorticityEqn = -(grad*delOmega)*(grad*omega) 
        - Re*delOmega*(velocity0*grad)*omega;

      Expr streamfunctionEqn = -(grad*delPsi)*(grad*psi) + delPsi*omega;

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

      /*
       * Boundary conditions: 
       * psi=0 on all surfaces
       * d(psi)/dn = 0 on left, right, bottom
       * d(psi)/dn = 1 on top
       */

      EssentialBC bc = EssentialBC(top, delOmega*psi)
        && EssentialBC(bottom, delOmega*psi)
        && EssentialBC(left, delOmega*psi)
        && EssentialBC(right, delOmega*psi);


      /*
        Create a solver object: stablized biconjugate gradient solver
      */
      TSFPreconditionerFactory prec = new ILUKPreconditionerFactory(1);
      TSFLinearSolver solver = new BICGSTABSolver(prec, 1.0e-14, 1000);
			
      StaticLinearProblem prob(mesh, eqn, bc,List(delPsi, delOmega),
                               List(psi, omega), 
                               new PetraVectorType());
			
      PicardLinearization picard(prob, u0, solver);
      Expr soln = picard.solve(PicardSolver(50, 1.0e-12));
			
      /*
        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();
			
}