#include "Sundance.h"


/*

twoZoneHeat2D solves the heat conduction equation in a two-material domain.
The domain is rectangular, with the left half having thermal resistivity alpha1
and the right half having thermal resistivity alpha2. Boundary conditions
are insulating on the top and bottom, and T=0 on the left and right edges.

The exact solution in subdomain n is
   u_n = 1/2 alpha_n x^2 + beta_n x + gamma_n

where beta and gamma are determined by the boundary and jump conditions:

   u1[-1] == 0                         (left bc)
   u2[1] == 0                          (right bc) 
   u1[0] == u2[0]                      (continuity at internal bdry)
   dx*u1[0]/alpha1 = dx*u2[0]/alpha2   (flux conservation at internal bdry)


 */

// Functions to identify cells in left and right zones.
bool inLeftZone(const Cell& cell);
bool inRightZone(const Cell& cell);

// Function that evaluates exact solution
ListBatch solnFunc(const PointBatch& pts);

// constants: thermal resistivities for zones 1 and 2.
static const double alpha1 = 1.0;
static const double alpha2 = 2.0;

// main solver code


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

      int n = 4;
      double h = 2.0/((double) n);
      Mesh mesh = rectMesh(-1.0, 1.0, n, -1.0, 1.0, n, false);

      mesh.labelCells(2, "leftZone", inLeftZone);
      mesh.labelCells(2, "rightZone", inRightZone);
      CellSet leftEdge(1, "left");
      CellSet rightEdge(1, "right");
      CellSet leftZone(2, "leftZone");
      CellSet rightZone(2, "rightZone");

			
      // define variational and unknown functions for the problem
      Expr delU = new TestFunction(Lagrange(2), "delU");
      Expr U = new UnknownFunction(Lagrange(2), "U");

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

      // source term
      double f = -1.0;

      // expr for the left subdomain
      Expr e1 = -1.0/alpha1*(grad*delU)*(grad*U)  + f*delU;
      // expr for the right subdomain
      Expr e2 = -1.0/alpha2*(grad*delU)*(grad*U)  + f*delU;

      // integrate the variational form over the whole domain,
      // with different integrands for the left and right zones.
      Expr eqn = Integral(leftZone, e1, GaussLegendre(4))
        + Integral(rightZone, e2, GaussLegendre(4));

      // impose essential bcs
      EssentialBC bc = EssentialBC(leftEdge, delU*U/h) 
        && EssentialBC(rightEdge, delU*U/h);

      // solver object
      LinearSolver solver = new BICGSTABSolver(1.e-15, 500);

      // problem object, combining symbolic expression, mesh, bcs, and solver
      StaticLinearProblem prob(mesh, eqn, bc, delU, U, solver);

      // solve it.
      Expr soln;
      prob.solve(soln);
			
      // compute the error and represent as a discrete function
      Expr exactSoln(solnFunc, Scalar(), "soln");
      Expr errorExpr = new DiscreteFunction(mesh, exactSoln-soln, 
                                            Lagrange(1), "errorExpr");
	
			
      // compute the norm of the error
      double errorNorm = errorExpr.norm();
      double tolerance = 1.0e-13;

      Testing::passFailCheck(__FILE__, errorNorm, tolerance);
      Testing::timeStamp(__FILE__, __DATE__, __TIME__);
    }
  catch(exception& e)
    {
      e.print();
      Testing::crash(__FILE__);
      Testing::timeStamp(__FILE__, __DATE__, __TIME__);
    }

}



bool inLeftZone(const Cell& cell)
{
  // x < 0.0
  const Point& a = cell.facet(0,0).point(0);
  const Point& b = cell.facet(0,1).point(0);
  const Point& c = cell.facet(0,2).point(0);
		
  if (a[0] <= 0.0 && b[0] <= 0.0 && c[0] <= 0.0)
    {
      return true;
    }
  return false;
}

bool inRightZone(const Cell& cell)
{
  return !inLeftZone(cell);
}

/*

 Function to evaluate exact soln at a point batch. 

 The problem is multimaterial, so the solution will take a different form
 in each subdomain. Therefore, for each cell on which we evaluate, we
 first identify the subdomain in which it lives and branch accordingly.

*/

ListBatch solnFunc(const PointBatch& pts)
{
  // calculate constants in exact soln. In a real problem, we'd do this once
  // and for all somewhere and store the results, but for this toy problem
  // we just recompute them at every call.
  double gamma1 = -alpha1*alpha2/(alpha1+alpha2);
  double gamma2 = gamma1;
  double beta1 = 0.5*alpha1*(alpha1-alpha2)/(alpha1+alpha2);
  double beta2 = 0.5*alpha2*(alpha1-alpha2)/(alpha1+alpha2);

  const TSFArray<Point>& p = pts.physPts();
  const Cell& cell = pts.cell();
  Vector soln(p.length());

  // loop over points in batch, evaluating function for every point.
  for (int i=0; i<p.length(); i++)
    {
      const Point& x = p[i];
      // identify current subdomain, and use the solution for that subdomain.
      if (inLeftZone(cell))
        {
          soln[i] = gamma1 + x[0]*(beta1 + 0.5*alpha1*x[0]);
        }
      else
        {
          soln[i] = gamma2 + x[0]*(beta2 + 0.5*alpha2*x[0]);
        }
    }
  return ListBatch(soln);
}