//
// Software License for AMDiS
//
// Copyright (c) 2010 Dresden University of Technology 
// All rights reserved.
// Authors: Simon Vey, Thomas Witkowski et al.
//
// This file is part of AMDiS
//
// See also license.opensource.txt in the distribution.


#include "time/ExtendedRosenbrockStationary.h"
#include "io/VtkWriter.h"
#include "SystemVector.h"
#include "OEMSolver.h"
#include "Debug.h"
#include "POperators_ZOT.h"
#include "VectorOperations.h"

#ifdef HAVE_PARALLEL_DOMAIN_AMDIS
#include "parallel/MeshDistributor.h"
#endif

using namespace AMDiS;

  void ExtendedRosenbrockStationary::initialize(Flag initFlag,
						ProblemStat *adoptProblem,
						Flag adoptFlag)
  { FUNCNAME_DBG("ExtendedRosenbrockStationary::initialize()");
    
    super::initialize(initFlag, adoptProblem, adoptFlag);
    
    // create rosenbrock method
    std::string str("");
    std::string initFileStr(name + "->rosenbrock->method");
    Parameters::get(initFileStr, str);
    RosenbrockMethodCreator *creator = 
      dynamic_cast<RosenbrockMethodCreator*>(CreatorMap<RosenbrockMethod>::getCreator(str, initFileStr));
    rm = creator->create();
    TEST_EXIT_DBG(rm)("No Rosenbrock method created!\n");
    
    init();
    
    // read error weights
    Parameters::get(name + "->rosenbrock->error weights", errorEstWeights);
    if (errorEstWeights.size() == 0) {
      WARNING("No weights for time-error estimation given.\n");
    }

    errorEstWeights.resize(super::getNumComponents(), 0.0);
    double SumErrorEstWeights = vector_operations::sum(errorEstWeights);
    if (SumErrorEstWeights<DBL_TOL) {
      WARNING("Using equal weights for all components!\n");
      for (size_t i = 0; i < errorEstWeights.size(); i++)
	errorEstWeights[i] = 1.0/static_cast<double>(super::getNumComponents());
    } else {
      for (size_t i = 0; i < errorEstWeights.size(); i++)
	errorEstWeights[i] *= 1.0/SumErrorEstWeights;
    }
    
    std::cout << "errorEstWeights = " << errorEstWeights << "\n";
  }

  void ExtendedRosenbrockStationary::init()
  {
    stageSolution = new SystemVector(*solution);
    unVec = new SystemVector(*solution);
    timeRhsVec = new SystemVector(*solution);
    newUn = new SystemVector(*solution);
    tmp = new SystemVector(*solution);
    if (!fixedTimestep)
      lowSol = new SystemVector(*solution);    

    stageSolution->set(0.0);
    unVec->set(0.0);
    
    stageSolutions.resize(rm->getStages());
    for (int i = 0; i < rm->getStages(); i++) {
      stageSolutions[i] = new SystemVector(*solution);
      stageSolutions[i]->set(0.0);
    }
    order = rm->getOrder();
  }

  Flag ExtendedRosenbrockStationary::oneIteration(AdaptInfo *adaptInfo, Flag toDo)
  {
    order = rm->getOrder();
    Flag flag = 0, markFlag = 0;
    

    if (toDo.isSet(MARK))
      markFlag = problem->markElements(adaptInfo);
    else
      markFlag = 3;

    // refine
    if (toDo.isSet(ADAPT) && markFlag.isSet(MESH_REFINED))
      flag = problem->refineMesh(adaptInfo);

    // coarsen
    if (toDo.isSet(ADAPT) && markFlag.isSet(MESH_COARSENED))
      flag |= problem->coarsenMesh(adaptInfo);
    
    if (toDo.isSet(BUILD) || toDo.isSet(SOLVE)) {
      flag = stageIteration(adaptInfo, toDo, true, true);
      estimateTimeError(adaptInfo);
    }
    
    if (toDo.isSet(ESTIMATE))
      problem->estimate(adaptInfo);

    return flag;
  }
  
  Flag ExtendedRosenbrockStationary::stageIteration(AdaptInfo *adaptInfo, Flag flag,
					       bool asmMatrix, bool asmVector)
  {
    FUNCNAME("ExtendedRosenbrockStationary::stageIteration()");

    TEST_EXIT(tauPtr)("No tau pointer defined in stationary problem!\n");

    if (first) {
      first = false;
      *unVec = *solution;
    }
    
    *newUn = *unVec;    
    if (!fixedTimestep)
      *lowSol = *unVec;
    
    invTauGamma = 1.0 / ((*tauPtr) * rm->getGamma());

    for (int i = 0; i < rm->getStages(); i++) {
      
      stageTime = oldTime + rm->getAlphaI(i) * (*tauPtr);
      tauGammaI = rm->getGammaI(i) * (*tauPtr);
      
      // stage-solution: u_s(i) = u_old + sum_{j=0}^{i-1} a_ij*U_j
      *stageSolution = *unVec;
      for (int j = 0; j < i; j++) {
	*tmp = *(stageSolutions[j]);
	*tmp *= rm->getA(i, j);
	*stageSolution += *tmp;
      }

      // Dirichlet-BC implemented as additional algebraic equation u = g(x,t) on boundary
      // => U_i = -u_s(i) + g(x,t_s(i)) + tau*gamma_i* d_t(g)(t_old) on boundary
      // where u_s(i) = ith stage-solution, t_s(i) = ith stage-time
      for (unsigned int j = 0; j < boundaries.size(); j++) {
	boundaries[j].vec->interpol(boundaries[j].fct);
	*(boundaries[j].vec) -= *(stageSolution->getDOFVector(boundaries[j].col));
	
	if (boundaries[j].fctDt != NULL) {
	// time derivative of dirichlet bc is given
	  DOFVector<double>* tmpDt = new DOFVector<double>(getFeSpace(boundaries[j].col), "tmp");
	  tmpDt->interpol(boundaries[j].fctDt);
	  *tmpDt *= tauGammaI;
	  *(boundaries[j].vec) += *tmpDt; // alternativ transformDOF(...) benutzen, für x = x + a*y 
	  delete tmpDt;
	}
      }

      timeRhsVec->set(0.0);
      for (int j = 0; j < i; j++) {
	*tmp = *(stageSolutions[j]);
	*tmp *= (rm->getC(i, j) / *tauPtr);
	*timeRhsVec += *tmp;
      }

      super::buildAfterCoarsen(adaptInfo, flag | UPDATE_DIRICHLET_BC, (i == 0), true);
      
      super::solve(adaptInfo, i == 0, true);

      *(stageSolutions[i]) = *solution;
      
      *tmp = *solution;
      *tmp *= rm->getM1(i);
      *newUn += *tmp;

      if (!fixedTimestep) {
	// lower order approximation, with coefficients of embedded method
	// used for local error estimator
	*tmp = *solution;
	*tmp *= rm->getM2(i);
	*lowSol += *tmp;
      }
    }
    
    stageTime = oldTime + (*tauPtr);
    
    Flag flag_;
    return flag_;
  }

  
  double ExtendedRosenbrockStationary::estimateTimeError(AdaptInfo* adaptInfo)
  {
    oldErrorEst = errorEst;
    errorEst = 0.0;
    if (!fixedTimestep) {
      for (int i = 0; i < lowSol->getSize(); i++) {
	(*(lowSol->getDOFVector(i))) -= (*(newUn->getDOFVector(i)));
	adaptInfo->setTimeEstSum(lowSol->getDOFVector(i)->l2norm(), i+componentShift);
	errorEst += errorEstWeights[i] * adaptInfo->getTimeEstSum(i+componentShift);
      }
    }
    
    adaptInfo->setTimeEst(errorEst);
    return errorEst;
  }

  void ExtendedRosenbrockStationary::acceptTimestep(AdaptInfo* adaptInfo)
  {
    *solution = *newUn;
    *unVec = *newUn;
    oldTime = adaptInfo->getTime();
  }

      
  double ExtendedRosenbrockStationary::getNewTimestep(double tol, double tau, bool restrict)
  {
    double fac = pow((tol / errorEst), 1.0 / order);
    if (restrict)
      fac = std::min(3.0, fac);
    return 0.95 * (fac * tau);
  }
  
  double ExtendedRosenbrockStationary::getNewTimestep(double tol, double oldTau, double tau, bool restrict)
  {
    double fac = tau / oldTau * pow((tol * oldErrorEst / sqr(errorEst)), 1.0 / order);
    if (restrict)
      fac = std::min(3.0, fac);
    return 0.95 * (fac * tau);
  }
  
  void ExtendedRosenbrockStationary::addOperator(Operator &op, int row, int col,
					 double *factor, double *estFactor)
  {
    FUNCNAME("ExtendedRosenbrockStationary::addOperator()");

    TEST_EXIT(op.getUhOld() == NULL)("UhOld is not allowed to be set!\n");

    op.setUhOld(stageSolution->getDOFVector(col));
    super::addVectorOperator(op, row, factor, estFactor);
  }
  

  void ExtendedRosenbrockStationary::addJacobianOperator(Operator &op, int row, int col,
						 double *factor, double *estFactor)
  {
    FUNCNAME("ExtendedRosenbrockStationary::addJacobianOperator()");
    
    TEST_EXIT(factor == NULL)("Not yet implemented!\n");
    TEST_EXIT(estFactor == NULL)("Not yet implemented!\n");

    super::addMatrixOperator(op, row, col, &minusOne, &minusOne);
  }


  void ExtendedRosenbrockStationary::addTimeOperator(int row, int col, double factor)
  {
    FUNCNAME("ExtendedRosenbrockStationary::addTimeOperator()");

    Operator *op = new Operator(componentSpaces[row], componentSpaces[col]);
    op->addTerm(new Simple_ZOT(factor));
    super::addMatrixOperator(op, row, col, &invTauGamma, &invTauGamma);

    Operator *opRhs = new Operator(componentSpaces[row]);
    opRhs->addTerm(new Phase_ZOT(timeRhsVec->getDOFVector(col), factor));
    super::addVectorOperator(opRhs, row);
  }

  
  void ExtendedRosenbrockStationary::addDirichletBC(BoundaryType type, int row, int col,
					    AbstractFunction<double, WorldVector<double> > *fct,
					    AbstractFunction<double, WorldVector<double> > *fctDt)
  {
    FUNCNAME("ExtendedRosenbrockStationary::addDirichletBC()");

    DOFVector<double>* vec = new DOFVector<double>(componentSpaces[col], "vec");
    MyRosenbrockBoundary bound(fct, fctDt, vec, row, col);
    boundaries.push_back(bound);

    super::addDirichletBC(type, row, col, vec);
  }
  

  void ExtendedRosenbrockStationary::addSingularDirichletBC(WorldVector<double> &pos, int row, int col,
							    AbstractFunction<double, WorldVector<double> > *fct)
  {
    DOFVector<double>* vec = new DOFVector<double>(componentSpaces[col], "vec");
    MyRosenbrockBoundary bound(fct, NULL, vec, row, col);
    boundaries.push_back(bound);

    super::addSingularDirichletBC(pos, row, col, *vec);
  }


  void ExtendedRosenbrockStationary::addSingularDirichletBC(DegreeOfFreedom idx, int row, int col,
							    AbstractFunction<double, WorldVector<double> > *fct)
  {
    DOFVector<double>* vec = new DOFVector<double>(componentSpaces[col], "vec");
    MyRosenbrockBoundary bound(fct, NULL, vec, row, col);
    boundaries.push_back(bound);

    super::addSingularDirichletBC(idx, row, col, *vec);
  }
  

  void ExtendedRosenbrockStationary::addImplicitDirichletBC(AbstractFunction<double, WorldVector<double> > &signedDist, int row, int col,
							    AbstractFunction<double, WorldVector<double> > &fct)
  {
    DOFVector<double>* vec = new DOFVector<double>(componentSpaces[col], "vec");
    MyRosenbrockBoundary bound(&fct, NULL, vec, row, col);
    boundaries.push_back(bound);

    super::addImplicitDirichletBC(signedDist, row, col, *vec);
  }


  void ExtendedRosenbrockStationary::addImplicitDirichletBC(DOFVector<double> &signedDist, int row, int col,
							    AbstractFunction<double, WorldVector<double> > &fct)
  {
    DOFVector<double>* vec = new DOFVector<double>(componentSpaces[col], "vec");
    MyRosenbrockBoundary bound(&fct, NULL, vec, row, col);
    boundaries.push_back(bound);

    super::addImplicitDirichletBC(signedDist, row, col, *vec);
  }


  void ExtendedRosenbrockStationary::addManualDirichletBC(BoundaryType nr, AbstractFunction<bool, WorldVector<double> >* meshIndicator,
							  int row, int col,
							  AbstractFunction<double, WorldVector<double> > &fct)
  {
    DOFVector<double>* vec = new DOFVector<double>(componentSpaces[col], "vec");
    MyRosenbrockBoundary bound(&fct, NULL, vec, row, col);
    boundaries.push_back(bound);

    super::addManualDirichletBC(nr, meshIndicator, row, col, *vec);
  }


  void ExtendedRosenbrockStationary::addManualDirichletBC(AbstractFunction<bool, WorldVector<double> >* meshIndicator,
							  int row, int col,
							  AbstractFunction<double, WorldVector<double> > &fct)
  {
    DOFVector<double>* vec = new DOFVector<double>(componentSpaces[col], "vec");
    MyRosenbrockBoundary bound(&fct, NULL, vec, row, col);
    boundaries.push_back(bound);

    super::addManualDirichletBC(meshIndicator, row, col, *vec);
  }