#include "Sundance.h"

/*
 * 
 */

Expr List(const Expr& a, const Expr& b, const Expr& c, const Expr& d,
          const Expr& e)
{
  Expr rtn = List(a,b,c,d);
  rtn.append(e);
  return rtn;
}

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

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

      /* place 2 vents on the left and right walls */
      CellSet leftVents = boundary.subset( (x==0.0 && y>0.2*L && y<0.3*L) 
                                           || (x==0.0 && y>0.7*L && y<0.8*L));
      CellSet rightVents = boundary.subset( (x==L && y>0.2*L && y<0.3*L) 
                                            || (x==L && y>0.7*L && y<0.8*L));
      CellSet vents = leftVents + rightVents;
      CellSet walls = boundary - vents;

      /* diffusivity */
      Expr nu = new ParameterExpr(1.0);;

      /* test and unknown functions for the velocity potential */
      Expr varPhi = new TestFunction(new Lagrange(1));
      Expr dPhi = new UnknownFunction(new Lagrange(1));

      /* test and unknown functions for the concentration */
      Expr varC = new TestFunction(new Lagrange(1));
      Expr dc = new UnknownFunction(new Lagrange(1));

      /* test and unknown functions for the flux lagrange multiplier */
      Expr varLambda = new TestFunction(new Lagrange(1), "lambda");
      Expr dLambda = new UnknownFunction(new Lagrange(1), "mu");

      /* test and unknown functions for the concentration
       * lagrange multiplier */
      Expr varMu = new TestFunction(new Lagrange(1));
      Expr dmu = new UnknownFunction(new Lagrange(1));

      /* test and unknown function for the vent flux */
      Expr varG = new TestFunction(new Lagrange(1));
      Expr dg = new UnknownFunction(new Lagrange(1));

			
      /* set up discrete functions for initial guess */
      BasisFamily lagr = new Lagrange(1);
      CellSet everything = CellSet::maximalCells();
      TSFVectorSpace discreteSpace 
        = new SundanceVectorSpace(mesh, 
                                  tuple(lagr, lagr, lagr, lagr, lagr),
                                  tuple(everything, everything, everything,
                                        everything, vents),
                                  new DenseSerialVectorType());
      Expr zero = List(Expr(0.0), Expr(0.0), Expr(0.0), Expr(0.0),
                       Expr(0.0));
      Expr u0 = DiscreteFunction::discretize(discreteSpace, zero);
      Expr phi0 = u0[0];
      Expr c0 = u0[1];		
      Expr mu0 = u0[2];
      Expr lambda0 = u0[3];
      Expr g0 = u0[4];
			

      /* define derivatives and gradient operator */
      Expr dx = new Derivative(0,1);
      Expr dy = new Derivative(1,1);
      Expr grad = List(dx, dy);


      /* Define a Gaussian source centered at (x0,y0) */
      double x0 = 0.5;
      double y0 = 0.5;
      Expr deltaX = hold(x-x0);
      Expr deltaY = hold(y-y0);
      Expr d2 = deltaX*deltaX + deltaY*deltaY;
      Expr f = exp(-d2);

      /* weak form for potential flow equation */
      Expr flowEqn = Integral((grad*varPhi)*(grad*(phi0 + dPhi)))
        + Integral(vents, -varPhi*(g0 + dg));

      /* weak form for convection-diffusion equation */
      Expr cdEqn = Integral((grad*(c0 + dc))*(grad*varC) 
                            + varC*(grad*(phi0+dPhi))*(grad*c0)
                            + varC*(grad*phi0)*(grad*dc)
                            + varC*f);
																				

      /* weak form for velocity lagr multipliers */
      Expr lambdaEqn = Integral((grad*varLambda)*(grad*dLambda)
                                + grad*(c0+dc)*(grad*varLambda)*(mu0) 
                                + (grad*(c0))*(grad*varLambda)*(dmu));

      /* weak form for convection-diffusion adjoints (mu) */
      Expr muEqn = Integral( nu * (grad*varMu)*(grad*(mu0+dmu)) +
                             (grad*varMu)*(grad*(phi0+dPhi))*mu0 +
                             (grad*varMu)*(grad*phi0)*dmu +
                             (c0+dc)*varMu);

      /* add up all equations */
      Expr eqn = flowEqn + cdEqn + lambdaEqn + muEqn;

      /*
       * Boundary conditions:
       * - c is dirichlet zero on walls, natural on vents.
       * - phi is natural on walls, Neumann dphi/dn=g on vents.
       * - lambda is natural everywhere
       * - mu is zero on the walls, natural on vents
       * - g does not obey a PDE and thus has no BCs
       */

      EssentialBC bc = EssentialBC(walls, varC*(c0+dc) + varMu*(mu0+dmu))
        && EssentialBC(vents, varG*(g0+dg-lambda0-dLambda));
			
			
      Expr tests = List(varPhi, varC, varMu, varLambda, varG);
      Expr unks = List(dPhi, dc, dmu, dLambda, dg);

      StaticLinearProblem linearizedProb(mesh, eqn, bc, 
                                         tests, unks,
                                         new DenseSerialVectorType()); 
			
      //			TSFPreconditionerFactory prec = new ILUKPreconditionerFactory(2);
      //			TSFLinearSolver solver = new BICGSTABSolver(prec, 1.0e-6, 400);
      TSFLinearSolver solver = new DirectSolver();

      NewtonLinearization newton(linearizedProb, u0, solver);
      NewtonSolver newtonSolver(solver, 150, 1.0e-4, 1.0e-4);
      Expr soln = newton.solve(newtonSolver);
    }
  catch(exception& e)
    {
      TSFOut::println(e.what());
      Testing::crash(__FILE__);
      Testing::timeStamp(__FILE__, __DATE__, __TIME__);
    }
  MPIComm::finalize();
}