#include "AMDiS.h"

using namespace std;
using namespace AMDiS;

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

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

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

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

  const double& operator()(const WorldVector<double>& x) const {
    return 0.0;
  };
};

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

  /** \brief
   * Constructor
   */
  F(int degree)
    : AbstractFunction<double, WorldVector<double> >(degree)
  {};

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

/** \brief
 * Needed for zero order term.
 */
class X3 : public AbstractFunction<double, double>
{
public:
  MEMORY_MANAGED(X3);

  X3() : AbstractFunction<double, double>(3) {};

  /** \brief
   * Implementation of AbstractFunction::operator().
   */
  double operator()(const double& x) const {
    return x * x * x;
  }
};

// ===========================================================================
// ===== class NonLin ========================================================
// ===========================================================================

/** \brief
 * Non linear problem.
 */
class NonLin : public ProblemScal
{
public:
  NonLin(const char *name)
    : ProblemScal(name)
  {
    newtonTolerance = 1e-8;
    GET_PARAMETER(0, std::string(name) + "->newton->tolerance", "%f", 
		  &newtonTolerance);    
    newtonMaxIter = 50;
    GET_PARAMETER(0, std::string(name) + "->newton->max iteration", "%d", 
		  &newtonMaxIter);    
  };

  void initialize(Flag initFlag,
		  ProblemScal *adoptProblem = NULL,
		  Flag adoptFlag = INIT_NOTHING)
  {
    ProblemScal::initialize(initFlag, adoptProblem, adoptFlag);
    correction = NEW DOFVector<double>(this->getFESpace(), "old solution");
    correction->set(0.0);

    dirichletZero = NEW DirichletBC(1, &zero, feSpace_);
    dirichletG    = NEW DirichletBC(1, &g, feSpace_);

    systemMatrix_->getBoundaryManager()->addBoundaryCondition(dirichletZero);
    solution_->getBoundaryManager()->addBoundaryCondition(dirichletG);
    rhs_->getBoundaryManager()->addBoundaryCondition(dirichletZero);
    correction->getBoundaryManager()->addBoundaryCondition(dirichletZero);
  };

  ~NonLin()
  { 
    DELETE correction; 
    DELETE dirichletZero;
    DELETE dirichletG;
  };

  void buildAfterCoarsen(AdaptInfo *adaptInfo, Flag flag) {};

  void solve(AdaptInfo *adaptInfo) 
  {
    FUNCNAME("NonLin::solve()");
    DOFVector<double> *sol = solution_;
    solution_ = correction;
    double res = 0;
    int newtonIteration = 0;
    Flag flag;

    // fill boundary conditions
    sol->getBoundaryManager()->initVector(sol);

    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh_, -1, 
					 Mesh::CALL_LEAF_EL | 
					 Mesh::FILL_COORDS |
					 Mesh::FILL_BOUND);

    while(elInfo) {
      sol->getBoundaryManager()->fillBoundaryConditions(elInfo, sol);
      elInfo = stack.traverseNext(elInfo);
    }

    sol->getBoundaryManager()->exitVector(sol);

    do {    
      newtonIteration++;
      ProblemScal::buildAfterCoarsen(adaptInfo, flag);    
      ProblemScal::solve(adaptInfo);
      *sol -= *correction;
      res = correction->L2Norm();
      MSG("newton iteration %d: residual %f (tol: %f)\n",
	  newtonIteration, res, newtonTolerance);
    } while((res > newtonTolerance) && (newtonIteration < newtonMaxIter));
    solution_ = sol;
  };

private:
  double newtonTolerance;
  int newtonMaxIter;
  DOFVector<double> *correction;
  Zero zero;
  G g;
  DirichletBC *dirichletZero;
  DirichletBC *dirichletG;
};

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

int main(int argc, char* argv[])
{
  FUNCNAME("main");

  // ===== check for init file =====
  TEST_EXIT(argc == 2)("usage: nonlin initfile\n");

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

  NonLin nonlin("nonlin");
  nonlin.initialize(INIT_ALL);

  // === create adapt info ===
  AdaptInfo *adaptInfo = NEW AdaptInfo("nonlin->adapt", 1);

  // === create adapt ===
  AdaptStationary *adapt = NEW AdaptStationary("nonlin->adapt",
					       &nonlin,
					       adaptInfo);

  // ===== create operators =====
  double four = 4.0;
  double one = 1.0;
  double zero = 0.0;
  double minusOne = -1.0;

  Operator *nonlinOperator0 = NEW Operator(Operator::MATRIX_OPERATOR | 
					   Operator::VECTOR_OPERATOR, 
					   nonlin.getFESpace());

  nonlinOperator0->setUhOld(nonlin.getSolution());
  nonlinOperator0->addZeroOrderTerm(NEW VecAtQP_ZOT(nonlin.getSolution(), 
						    NEW X3));

  nonlin.addMatrixOperator(nonlinOperator0, &four, &one);
  nonlin.addVectorOperator(nonlinOperator0, &one, &zero);

  Operator *nonlinOperator2 = NEW Operator(Operator::MATRIX_OPERATOR | 
					   Operator::VECTOR_OPERATOR,
					   nonlin.getFESpace());

  nonlinOperator2->setUhOld(nonlin.getSolution());
  nonlinOperator2->addSecondOrderTerm(NEW Laplace_SOT);

  nonlin.addMatrixOperator(nonlinOperator2, &one, &one);
  nonlin.addVectorOperator(nonlinOperator2, &one, &zero);

  int degree = nonlin.getFESpace()->getBasisFcts()->getDegree();

  Operator* rhsFunctionOperator = NEW Operator(Operator::VECTOR_OPERATOR, 
					       nonlin.getFESpace());
  rhsFunctionOperator->addZeroOrderTerm(NEW CoordsAtQP_ZOT(NEW F(degree)));
  
  nonlin.addVectorOperator(rhsFunctionOperator, &minusOne, &one);

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

  nonlin.writeFiles(adaptInfo, true);
}