#include "AMDiS.h"

using namespace std;
using namespace AMDiS;

class G : public AbstractFunction<double, WorldVector<double> >
{
public:
  MEMORY_MANAGED(G);

  const double& operator()(const WorldVector<double>& x) const
  {
    static double result;
    result = exp(-10.0*(x*x));
    return result;  
  };
};

class F : public AbstractFunction<double, WorldVector<double> >
{
public:
  MEMORY_MANAGED(F);

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

  const double& operator()(const WorldVector<double>& x) const {
    static double result = 0.0;
    int dow = x.getSize();
    double r2 = (x*x);
    double ux = exp(-10.0*r2);
    result = -(400.0*r2 - 20.0*dow)*ux;
    return result;
  };
};

class MyCoupledIteration : public ProblemIterationInterface
{
public:
  MyCoupledIteration(ProblemStatBase *prob1,
		     ProblemStatBase *prob2)
    : problem1(prob1),
      problem2(prob2)
  {};

  void beginIteration(AdaptInfo *adaptInfo) 
  {
    FUNCNAME("StandardProblemIteration::beginIteration()");
    MSG("\n");
    MSG("begin of iteration number: %d\n", adaptInfo->getSpaceIteration()+1);
    MSG("=============================\n");
  };

  void endIteration(AdaptInfo *adaptInfo) {
    FUNCNAME("StandardProblemIteration::endIteration()");
    MSG("\n");
    MSG("end of iteration number: %d\n", adaptInfo->getSpaceIteration()+1);
    MSG("=============================\n");
  };

  Flag oneIteration(AdaptInfo *adaptInfo, Flag toDo = FULL_ITERATION) 
  {
    Flag flag, markFlag;
    if(toDo.isSet(MARK)) markFlag = problem1->markElements(adaptInfo);
    if(toDo.isSet(ADAPT) && markFlag.isSet(MESH_REFINED)) flag = problem1->refineMesh(adaptInfo);

    if(toDo.isSet(BUILD)) problem1->buildAfterCoarsen(adaptInfo, markFlag);
    if(toDo.isSet(SOLVE)) problem1->solve(adaptInfo);
 
    if(toDo.isSet(BUILD)) problem2->buildAfterCoarsen(adaptInfo, markFlag);
    if(toDo.isSet(SOLVE)) problem2->solve(adaptInfo);

    if(toDo.isSet(ESTIMATE)) problem1->estimate(adaptInfo);    
    return flag;
  };

  int getNumProblems() 
  {
    return 2;
  };

  ProblemStatBase *getProblem(int number = 0) 
  {
    FUNCNAME("CoupledIteration::getProblem()");
    if(number == 0) return problem1;
    if(number == 1) return problem2;
    ERROR_EXIT("invalid problem number\n");
    return NULL;
  };

private:
  ProblemStatBase *problem1;
  ProblemStatBase *problem2;
};

class Identity : public AbstractFunction<double, double>
{
public:
  MEMORY_MANAGED(Identity);

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

  const double& operator()(const double& x) const {
    static double result;
    result = x;
    return result;
  };
};

int main(int argc, char* argv[])
{
  FUNCNAME("main");
  TEST_EXIT(argc == 2)("usage: couple initfile\n");
  Parameters::init(true, argv[1]);

  // ===== create and init the first problem ===== 
  ProblemScal problem1("problem1");
  problem1.initialize(INIT_ALL);

  // ===== create and init the second problem ===== 
  Flag initFlag = 
    INIT_FE_SPACE |
    INIT_SYSTEM |
    INIT_SOLVER |
    INIT_FILEWRITER;

  Flag adoptFlag =
    CREATE_MESH |
    INIT_MESH;

  ProblemScal problem2("problem2");
  problem2.initialize(initFlag,
		      &problem1,
		      adoptFlag);

  // ===== add boundary conditions for problem1 =====
  problem1.addDirichletBC(1, NEW G);

  // ===== add boundary conditions for problem1 =====
  //problem2.addDirichletBC(1, NEW G);

  // ===== create operators for problem1 =====
  Operator matrixOperator1(Operator::MATRIX_OPERATOR, problem1.getFESpace());
  matrixOperator1.addSecondOrderTerm(NEW Laplace_SOT);
  problem1.addMatrixOperator(&matrixOperator1);

  int degree = problem1.getFESpace()->getBasisFcts()->getDegree();
  Operator rhsOperator1(Operator::VECTOR_OPERATOR, problem1.getFESpace());
  rhsOperator1.addZeroOrderTerm(NEW CoordsAtQP_ZOT(NEW F(degree)));
  problem1.addVectorOperator(&rhsOperator1);

  // ===== create operators for problem2 =====
  Operator matrixOperator2(Operator::MATRIX_OPERATOR, problem2.getFESpace());
  matrixOperator2.addZeroOrderTerm(NEW Simple_ZOT);
  problem2.addMatrixOperator(&matrixOperator2);
  
  Operator rhsOperator2(Operator::VECTOR_OPERATOR, problem2.getFESpace());
  rhsOperator2.addZeroOrderTerm(NEW VecAtQP_ZOT(problem1.getSolution(), 
						NEW Identity(degree)));
  problem2.addVectorOperator(&rhsOperator2);

  // ===== create adaptation loop and iteration interface =====
  AdaptInfo *adaptInfo = NEW AdaptInfo("couple->adapt", 1);
  MyCoupledIteration coupledIteration(&problem1, &problem2);
  AdaptStationary *adapt = NEW AdaptStationary("couple->adapt",
					       &coupledIteration,
					       adaptInfo);

  // ===== start adaptation loop =====
  adapt->adapt();

  // ===== write solution ===== 
  problem1.writeFiles(adaptInfo, true);
  problem2.writeFiles(adaptInfo, true);
}