#include "Sundance.h"

/** \example bdrySrcInversion.cpp
 * Solves the Neumann control problem for the convection-diffusion equation on 
 * the rectangle [-L, L] x [-H, H].
 * The state variable is c, the concentration of a passive contaminent.
 * The adjoint variable is lambda.
 * The control variable is p.
 * The velocity field is u. 
 * Concentrations have been measured at a number of sensor points.
 *
 * Boundary conditions are:
 * \frac{\partial }{\partial n}(-L) = p 
 * c(\pm H) = 0
 * \frac{\partial }{\partial n}(-L) = 0
 *
 * The problem is solved using a full-space all-at-once approach.
 * Data is generated by first solving a forward problem.
 */

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

      /* ------ create some objects common nto both forward and inverse probs ------- */

      /* create mesh */
      int nx = 40;
      int ny = 40;
      double L = 1.0;
      double H = 1.0;

      MeshGenerator mesher = new RectangleMesher(-L, L, nx, -H, H, ny);
      Mesh mesh = mesher.getMesh();

      /* define coordinate functions for x and y coordinates */
      Expr x = new CoordExpr(0);
      Expr y = new CoordExpr(1);
			
      /* velocity field */
      Expr ux = (H-y)*(H+y);
      Expr uy = 0.0;
      Expr u = List(ux, uy);

      /* diffusivity */
      Expr k = 1.0;

      /* define cells sets for each of the four sides of the rectangle */
      CellSet boundary = new BoundaryCellSet();
      CellSet left = boundary.subset( x == -L );
      CellSet right = boundary.subset( x == L );
      CellSet bottom = boundary.subset( y == -H );
      CellSet top = boundary.subset( y == H );
		
      /* Create a vector space factory, used to 
       * specify the low-level linear algebra representation */
      TSFVectorType petra = new PetraVectorType();

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

      /* create solver */
      TSFPreconditionerFactory prec = new ILUKPreconditionerFactory(1);
      TSFLinearSolver solver = new BICGSTABSolver(prec, 1.0e-12, 3000);


      /* ------------------ solve forward problem to create sensor data ----------- */

      /* use a Gaussian flux profile in the Neumann BC */
      Expr y0 = 0.0;
      Expr deltaY = 0.2;
      Expr pSoln = exp(-(y-y0)*(y-y0)/deltaY/deltaY);
			
      /* expr for unknown concentration in forward problem */
      Expr w = new UnknownFunction(new Lagrange(2));
			
      /* write forward convection-diffusion problem */
      Expr forwardEqn = Integral( -k*(grad*w)*(grad*w.variation()) - w.variation()*(u*grad)*w )
        + Integral(left, k*w.variation()*pSoln);

      EssentialBC forwardBC = EssentialBC(top+bottom, w*w.variation());

      StaticLinearProblem forwardProb(mesh, forwardEqn, forwardBC, w.variation(), w);
			
      /* solve the forward problem */
      Expr forwardSoln = forwardProb.solve(solver);

      FieldWriter writer1 = new MatlabWriter("uForward.dat");
      writer1.writeField(forwardSoln);

      /* ------------------ read the sensor data ----------------------------------- */

      /* state variable */
      Expr c = new UnknownFunction(new Lagrange(2));


      CellSet pointCells = new DimensionalCellSet(0);

      /* The sensors are laid out on a nsx*nsy grid */
      int nsx = 3;
      int nsy = 3;

      /* form the squared-residual part of the objective function */
      Expr objectiveFunction = 0.0;
      for (int i=0; i<nsx; i++) 
        {
          double sx = -L + 2.0*(i+1)*L/((double) nsx+1);
          for (int j=0; j<nsy; j++) 
            {
              double sy = -H + 2.0*(j+1)*H/((double) nsy+1);
              Cell sensor = mesh.cellNearestToPoint(Point(sx, sy));
              string label = "sensor(" + TSF::toString(i) + ", " + TSF::toString(j) + ")";
              sensor.setLabel(label);
              CellSet sensorRegion = pointCells.labeledSubset(label);
              objectiveFunction += Integral(sensorRegion, (c-forwardSoln)*(c-forwardSoln));
            }
        }

			
      /* ------------------ set up and solve the inverse problem --------------- */

      /* adjoint variable */
      Expr lambda = new UnknownFunction(new Lagrange(2));
      /* inversion variable */
      Expr p = new UnknownFunction(new Lagrange(2));

      /* add a regularization term to the objective function */
      Expr R = 1.0;
      objectiveFunction += Integral(left, 0.5*R*p*p);

      /* form the Lagrangian */
      Expr lagrangian = objectiveFunction + Integral(left, k*p*lambda)
        - Integral(k*(grad*c)*(grad*lambda)) - Integral(lambda*(u*grad)*c);

      EssentialBC lagrangianBC = EssentialBC(top+bottom, lambda*c);
			
      Expr unkFuncs = List(c, lambda, p);
      Expr testFuncs = List(c.variation(), lambda.variation(), p.variation());

      Expr kktExpr = lagrangian.variation(unkFuncs);
      EssentialBC kktBC = lagrangianBC.variation(unkFuncs);

      /* Set up the inversion problem. 
       * Note that the equations are reordered unsymmetrically; if you 
       * don't reorder, the solver won't converge. */
      StaticLinearProblem prob(mesh, kktExpr, kktBC, 
                               List(lambda.variation(), c.variation(), p.variation()), 
                               List(c, lambda, p));
															 
      Expr soln = prob.solve(solver);

      FieldWriter writer2 = new MatlabWriter("uInv.dat");
      writer2.writeField(soln[0]);
    }
  catch(exception& e)
    {
      Sundance::handleError(e, __FILE__);
    }
  Sundance::finalize();
}