#include "AMDiS.h"

using namespace std;
using namespace AMDiS;

// ===========================================================================
// ===== function definitions ================================================
// ===========================================================================

/** \brief
 * Dirichlet boundary function
 */
class G : public AbstractFunction<double, WorldVector<double> >,
	  public TimedObject
{
public:
  MEMORY_MANAGED(G);

  /** \brief
   * Implementation of AbstractFunction::operator().
   */
  double operator()(const WorldVector<double>& x) const {
    return sin(M_PI * (*timePtr)) * exp(-10.0 * (x * x));
  }
};

/** \brief
 * RHS function
 */
class F : public AbstractFunction<double, WorldVector<double> >,
	  public TimedObject
{
public:
  MEMORY_MANAGED(F);

  F(int degree) : AbstractFunction<double, WorldVector<double> >(degree) {};

  /** \brief
   * Implementation of AbstractFunction::operator().
   */
  double operator()(const WorldVector<double>& x) const {
    int dim = x.getSize();
    double r2 = x * x;
    double ux = sin(M_PI * (*timePtr)) * exp(-10.0 * r2);
    double ut = M_PI * cos(M_PI * (*timePtr)) * exp(-10.0 * r2);
    return ut -(400.0 * r2 - 20.0 * dim) * ux;
  }
};

// ===========================================================================
// ===== instationary problem ================================================
// ===========================================================================

/** \brief
 * Instationary problem
 */
class Heat : public ProblemInstatScal
{
public:
  MEMORY_MANAGED(Heat);

  /** \brief
   *  Constructor
   */
  Heat(ProblemScal *heatSpace)
    : ProblemInstatScal("heat", heatSpace)
  {
    // init theta scheme
    theta = -1.0;
    GET_PARAMETER(0, name + "->theta", "%f", &theta);
    TEST_EXIT(theta >= 0)("theta not set!\n");
    if (theta == 0.0) {
      WARNING("You are using the explicit Euler scheme\n");
      WARNING("Use a sufficiently small time step size!!!\n");
    }
    MSG("theta = %f\n", theta);
    theta1 = theta - 1;
  }


  // ===== ProblemInstatBase methods ===================================

  /** \brief
   * set the time in all needed functions!
   */
  void setTime(AdaptInfo *adaptInfo) {
    rhsTime = adaptInfo->getTime() - (1 - theta) * adaptInfo->getTimestep();
    boundaryTime = adaptInfo->getTime();
    tau1 = 1.0 / adaptInfo->getTimestep();    
  }

  void closeTimestep(AdaptInfo *adaptInfo) {
    ProblemInstatScal::closeTimestep(adaptInfo);
    WAIT;
  }

  // ===== initial problem methods =====================================

  /** \brief
   * Used by \ref problemInitial to solve the system of the initial problem
   */
  void solve(AdaptInfo *adaptInfo, bool) {
    problemStat->getSolution()->interpol(exactSolution);
  }

  /** \brief
   * Used by \ref problemInitial to do error estimation for the initial
   * problem.
   */
  void estimate(AdaptInfo *adaptInfo) {
    double errMax, errSum;

    errSum = Error<double>::L2Err(*exactSolution,
				  *(problemStat->getSolution()), 0, &errMax, false);

    adaptInfo->setEstSum(errSum, 0);
    adaptInfo->setEstMax(errMax, 0);
  }

  // ===== setting methods ===============================================

  /** \brief
   * Sets \ref exactSolution;
   */
  void setExactSolution(AbstractFunction<double, WorldVector<double> > *fct) {
    exactSolution = fct;
  } 

  // ===== getting methods ===============================================

  /** \brief
   * Returns pointer to \ref theta.
   */
  double *getThetaPtr() { 
    return &theta; 
  }

  /** \brief
   * Returns pointer to \ref theta1.
   */
  double *getTheta1Ptr() { 
    return &theta1; 
  }

  /** \brief
   * Returns pointer to \ref tau1
   */
  double *getTau1Ptr() { 
    return &tau1; 
  }

  /** \brief
   * Returns pointer to \ref rhsTime.
   */
  double *getRHSTimePtr() { 
    return &rhsTime; 
  }

  /** \brief
   * Returns pointer to \ref theta1.
   */
  double *getBoundaryTimePtr() {
    return &boundaryTime; 
  }

private:
  /** \brief
   * Used for theta scheme.
   */
  double theta;

  /** \brief
   * theta - 1
   */
  double theta1;

  /** \brief
   * 1.0 / timestep
   */
  double tau1;

  /** \brief
   * time for right hand side functions.
   */
  double rhsTime;

  /** \brief
   * time for boundary functions.
   */
  double boundaryTime;

  /** \brief
   * Pointer to boundary function. Needed for initial problem.
   */
  AbstractFunction<double, WorldVector<double> > *exactSolution;
};

// ===========================================================================
// ===== main program ========================================================
// ===========================================================================

int main(int argc, char** argv)
{
  // ===== check for init file =====
  TEST_EXIT(argc > 1)("usage: heat initfile\n");

  // ===== init parameters =====
  Parameters::init(false, argv[1]);
  Parameters::readArgv(argc, argv);

  // ===== create and init stationary problem =====
  ProblemScal *heatSpace = NEW ProblemScal("heat->space");
  heatSpace->initialize(INIT_ALL);

  // ===== create instationary problem =====
  Heat *heat = new Heat(heatSpace);
  heat->initialize(INIT_ALL);

  // create adapt info
  AdaptInfo *adaptInfo = NEW AdaptInfo("heat->adapt");

  // create initial adapt info
  AdaptInfo *adaptInfoInitial = NEW AdaptInfo("heat->initial->adapt");

  // create instationary adapt
  AdaptInstationary *adaptInstat =
    NEW AdaptInstationary("heat->adapt",
			  heatSpace,
			  adaptInfo,
			  heat,
			  adaptInfoInitial);

  // ===== create boundary functions =====
  G *boundaryFct = NEW G;
  boundaryFct->setTimePtr(heat->getBoundaryTimePtr());
  heat->setExactSolution(boundaryFct);

  heatSpace->addDirichletBC(DIRICHLET, boundaryFct);

  // ===== create rhs functions =====
  int degree = heatSpace->getFESpace()->getBasisFcts()->getDegree();
  F *rhsFct = NEW F(degree);
  rhsFct->setTimePtr(heat->getRHSTimePtr());

  // ===== create operators =====
  double one = 1.0;
  double zero = 0.0;

  // create laplace
  Operator *A = NEW Operator(Operator::MATRIX_OPERATOR | 
			     Operator::VECTOR_OPERATOR,
			     heatSpace->getFESpace());

  A->addSecondOrderTerm(new Laplace_SOT);

  A->setUhOld(heat->getOldSolution());

  if((*(heat->getThetaPtr())) != 0.0)
    heatSpace->addMatrixOperator(A, heat->getThetaPtr(), &one);

  if((*(heat->getTheta1Ptr())) != 0.0)
    heatSpace->addVectorOperator(A, heat->getTheta1Ptr(), &zero);

  // create zero order operator
  Operator *C = NEW Operator(Operator::MATRIX_OPERATOR | 
			     Operator::VECTOR_OPERATOR,
			     heatSpace->getFESpace());

  C->addZeroOrderTerm(NEW Simple_ZOT);

  C->setUhOld(heat->getOldSolution());

  heatSpace->addMatrixOperator(C, heat->getTau1Ptr(), heat->getTau1Ptr());
  heatSpace->addVectorOperator(C, heat->getTau1Ptr(), heat->getTau1Ptr());

  // create RHS operator
  Operator *F = NEW Operator(Operator::VECTOR_OPERATOR,
			     heatSpace->getFESpace());

  F->addZeroOrderTerm(NEW CoordsAtQP_ZOT(rhsFct));

  heatSpace->addVectorOperator(F);

  // ===== start adaption loop =====
  int errorCode = adaptInstat->adapt();

  return errorCode;
}