#include "AMDiS.h"

using namespace std;
using namespace AMDiS;

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

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

  /** \brief
   * Constructor
   */
  Project(int i) : comp(i) {};

  /** \brief
   * Implementation of AbstractFunction::operator().
   */
  double operator()(const WorldVector<double>& x) const {
    return x[comp];
  }

protected:
  /** \brief
   * coordinate plane to project at
   */
  int comp;
};


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

  ConstantFct(double c) 
    : AbstractFunction<double, WorldVector<double> >(0) ,
      constant(c)
  {};

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

protected:
  double constant;
};

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

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

  /** \brief
   * Implementation of AbstractFunction::operator().
   */
  double operator()(const WorldVector<double>& x) const {
    return (x[0] == 0.0) || (x[0] == 1.0) ? 0.0 : 1.0;
  }

protected:
  double constant;
};



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

/** \brief
 * Instationary problem
 */
class Navierstokes : public ProblemInstatVec
{
public:
  MEMORY_MANAGED(Navierstokes);

  /** \brief
   *  Constructor
   */
  Navierstokes(ProblemVec *navierstokesSpace) 
    : ProblemInstatVec("navierstokes", navierstokesSpace)
  {
    int i, dow = Global::getGeo(WORLD);
    TEST_EXIT(dow == navierstokesSpace->getNumComponents() - 1)
      ("number of components != dow + 1\n");
    boundaryFct.resize(dow+1, NULL);

    epsilon = -1.0;
    GET_PARAMETER(0, "epsilon", "%f", &epsilon);
    TEST_EXIT(epsilon != -1.0)("invalid epsilon\n");
  };

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

  /** \brief
   * Used by \ref problemInitial to solve the system of the initial problem
   */
  void solve(AdaptInfo *adaptInfo) 
  {
    int i, size = problemStat->getNumComponents()-1;
    for(i = 0; i < size; i++) {
      //problemStat->getMesh(i)->dofCompress();
      problemStat->getSolution()->getDOFVector(i)->interpol(boundaryFct[i]);
    }
  };

  /** \brief
   * Used by \ref problemInitial to do error estimation for the initial
   * problem.
   */
  void estimate(AdaptInfo *adaptInfo) 
  {
    int i, size = problemStat->getNumComponents()-1;
    double errMax, errSum;

    for(i = 0; i < size; i++) {
      errSum = 
	Error<double>::L2Err(*boundaryFct[i],
			     *(problemStat->getSolution()->getDOFVector(i)), 
			     0, &errMax, false);

      adaptInfo->setEstSum(errSum, i);
    }
  };

  /** \brief
   * Sets \ref boundaryFct;
   */
  void setBoundaryFct(AbstractFunction<double, WorldVector<double> > *fct, int i) 
  {
    boundaryFct[i] = fct;
  };

  /** \brief
   * Set the time in all needed functions! Called by adaption loop.
   */
  void setTime(AdaptInfo *adaptInfo) 
  {
    time = adaptInfo->getTime();
    tau1 = 1.0 / adaptInfo->getTimestep();
  };
  
  /** \brief
   * Returns address of \ref time.
   */
  double *getTimePtr() {
    return &time;
  };

  /** \brief
   * Returns address of \ref tau1.
   */
  double *getTau1Ptr() {
    return &tau1;
  };

  /** \brief
   * Implementation of ProblemInstatBase::closeTimestep()
   * pressure correction step
   */
  void closeTimestep(AdaptInfo *adaptInfo) {
    problemStat->writeFiles(adaptInfo, true);

    WAIT;

    int i, numComponents = problemStat->getNumComponents();
    int dow = Global::getGeo(WORLD);

    double timestep = adaptInfo->getTimestep();

    static DOFVector<double> pInterpol(problemStat->getFESpace(0), "p interpolated");
    static DOFVector<WorldVector<double> > gradPInterpol(problemStat->getFESpace(0), "gradient p interpolated");
    DOFVector<double> *p = problemStat->getSolution()->getDOFVector(numComponents-1);

    pInterpol.interpol(p, timestep * epsilon);
    pInterpol.getRecoveryGradient(&gradPInterpol);

    for(i = 0; i < dow; i++) {
      DOFVector<double>::Iterator 
	vIt(problemStat->getSolution()->getDOFVector(i), USED_DOFS);
      DOFVector<WorldVector<double> >::Iterator 
	gradIt(&gradPInterpol, USED_DOFS);

      for(vIt.reset(), gradIt.reset(); !vIt.end(); ++vIt, ++gradIt) {
	*vIt -= (*gradIt)[i];
      }
    }

    problemStat->writeFiles(adaptInfo, true);
  };

private:
  /** \brief
   * boundary function
   */
  std::vector<AbstractFunction<double, WorldVector<double> >*> boundaryFct;

  /** \brief
   * time
   */
  double time;

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

  /** \brief
   */
  double epsilon;
};

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

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

  // ===== check for init file =====
  TEST_EXIT(argc > 1)("usage: navierstokes initfile\n");

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

  // read nu (1/Re)
  double nu = -1.0;
  GET_PARAMETER(0, "nu", "%f", &nu);
  TEST_EXIT(nu != -1.0)("invalid nu\n");

  // ===== create and init stationary problem =====
  ProblemVec *navierstokesSpace = NEW ProblemVec("navierstokes->space");
  navierstokesSpace->initialize(INIT_ALL);

  int numComponents = navierstokesSpace->getNumComponents();

  // ===== create instationary problem =====
  Navierstokes *navierstokes = NEW Navierstokes(navierstokesSpace);
  navierstokes->initialize(INIT_ALL);

  // ===== create adapt Info =====
  AdaptInfo *adaptInfo = NEW AdaptInfo("navierstokes->adapt",
				       numComponents);

  // create initial adapt
  AdaptInfo *adaptInfoInitial = NEW AdaptInfo("navierstokes->initial->adapt",
					      numComponents);

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

  // ===== create boundary functions =====
  
  double *tau1 = navierstokes->getTau1Ptr();
  int i;

  // === boundary functions ===
  ConstantFct *boundaryFct = NEW ConstantFct(0.0);

  for(i = 0; i < numComponents-1; i++) {
    navierstokes->setBoundaryFct(boundaryFct, i);
    navierstokesSpace->addDirichletBC(1, i, boundaryFct);
  }
  navierstokesSpace->addDirichletBC(2, 0, NEW F());
  navierstokes->setBoundaryFct(boundaryFct, numComponents-1);
  

  // ===== create operators =====

  // ===== velocity =====
  for(i = 0; i < numComponents-1; i++) {
    // create laplace
    Operator *laplaceV = NEW Operator(Operator::MATRIX_OPERATOR,
				      navierstokesSpace->getFESpace(i),
				      navierstokesSpace->getFESpace(i));

    laplaceV->addSecondOrderTerm(NEW FactorLaplace_SOT(nu));
    navierstokesSpace->addMatrixOperator(laplaceV, i, i);

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

    dtv->setUhOld(navierstokes->getOldSolution()->getDOFVector(i));
    dtv->addZeroOrderTerm(NEW Simple_ZOT);

    navierstokesSpace->addMatrixOperator(dtv, i, i, tau1, tau1);
    navierstokesSpace->addVectorOperator(dtv, i, tau1, tau1);

    
    Operator *gradP = NEW Operator(Operator::VECTOR_OPERATOR,
				   navierstokesSpace->getFESpace(i));
    gradP->addZeroOrderTerm(NEW FctGradient_ZOT(navierstokes->getOldSolution()->
						getDOFVector(numComponents-1),
						NEW Project(i)));
    double minusOne = 1.0;
    navierstokesSpace->addVectorOperator(gradP, i, &minusOne, &minusOne);
  }

  // ===== pressure =====
  Operator *laplaceP = NEW Operator(Operator::MATRIX_OPERATOR,
				    navierstokesSpace->getFESpace(numComponents-1),
				    navierstokesSpace->getFESpace(numComponents-1));

  laplaceP->addSecondOrderTerm(NEW Laplace_SOT());

  navierstokesSpace->addMatrixOperator(laplaceP, 
				       numComponents-1, 
				       numComponents-1);
  
  WorldVector<double> b;

  for(i = 0; i < numComponents-1; i++) {
    Operator *divV = NEW Operator(Operator::MATRIX_OPERATOR,
				  navierstokesSpace->getFESpace(numComponents-1),
				  navierstokesSpace->getFESpace(i));

    b.set(0.0);
    b[i] = 1.0;

    divV->addFirstOrderTerm(NEW Vector_FOT(b), GRD_PHI);

    navierstokesSpace->addMatrixOperator(divV, numComponents-1, i, tau1, tau1);    
  }

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