#include "ProblemScal.h"
#include "AbstractFunction.h"
#include "DirichletBC.h"
#include "RobinBC.h"
#include "FixVec.h"
#include "Flag.h"
#include "Serializer.h"
#include "RecoveryEstimator.h"
#include "Operator.h"
#include "DOFMatrix.h"
#include "FiniteElemSpace.h"
#include "Estimator.h"
#include "OEMSolver.h"
#include "Preconditioner.h"
#include "MatVecMultiplier.h"
#include "DOFVector.h"
#include "Marker.h"
#include "AdaptInfo.h"
#include "ElInfo.h"
#include "FileWriter.h"
#include "RefinementManager.h"
#include "CoarseningManager.h"
#include "Lagrange.h"
#include "PeriodicBC.h"
#include "ValueReader.h"
#include "ElementFileWriter.h"

namespace AMDiS {

  ProblemScal *ProblemScal::traversePtr_ = NULL;

  void ProblemScal::writeFiles(AdaptInfo *adaptInfo, bool force) {
    for (int i = 0; i < static_cast<int>(fileWriters_.size()); i++) {
      fileWriters_[i]->writeFiles(adaptInfo, force);
    }
  }

  void ProblemScal::writeDelayedFiles()
  {
    for (int i = 0; i < static_cast<int>(fileWriters_.size()); i++) {
      fileWriters_[i]->writeDelayedFiles();
    }
  }

  bool ProblemScal::existsDelayedCalculation()
  {
    for (int i = 0; i < static_cast<int>(fileWriters_.size()); i++) {
      if (fileWriters_[i]->isWritingDelayed())
	return true;   
    }

    return false;
  }

  void ProblemScal::interpolInitialSolution(AbstractFunction<double, WorldVector<double> > *fct) 
  {
    solution_->interpol(fct);
  }

  void ProblemScal::addMatrixOperator(Operator *op, 
				      double *factor,
				      double *estFactor) 
  {
    systemMatrix_->addOperator(op, factor, estFactor);
  }

  void ProblemScal::addVectorOperator(Operator *op, 
				      double *factor,
				      double *estFactor) 
  {
    rhs_->addOperator(op, factor, estFactor);
  }
  
  void ProblemScal::addDirichletBC(BoundaryType type, 
				   AbstractFunction<double, WorldVector<double> >* b) 
  {
    DirichletBC *dirichlet = new DirichletBC(type, b, feSpace_);
    if (systemMatrix_) 
      systemMatrix_->getBoundaryManager()->addBoundaryCondition(dirichlet);
    if (rhs_) 
      rhs_->getBoundaryManager()->addBoundaryCondition(dirichlet);
    if (solution_)
      solution_->getBoundaryManager()->addBoundaryCondition(dirichlet);
  }

  void ProblemScal::addDirichletBC(BoundaryType type, 
				   DOFVector<double> *vec) 
  {
    DirichletBC *dirichlet = new DirichletBC(type, vec);
    if (systemMatrix_) 
      systemMatrix_->getBoundaryManager()->addBoundaryCondition(dirichlet);
    if (rhs_) 
      rhs_->getBoundaryManager()->addBoundaryCondition(dirichlet);
    if (solution_)
      solution_->getBoundaryManager()->addBoundaryCondition(dirichlet);
  }

  void ProblemScal::addRobinBC(BoundaryType type, 
			       AbstractFunction<double, WorldVector<double> > *n,
			       AbstractFunction<double, WorldVector<double> > *r)
  {
    RobinBC *robin = new RobinBC(type, n, r, feSpace_);
    if (rhs_) 
      rhs_->getBoundaryManager()->addBoundaryCondition(robin);
    if (systemMatrix_)
      systemMatrix_->getBoundaryManager()->addBoundaryCondition(robin);
  }

  void ProblemScal::addNeumannBC(BoundaryType type, 
				 AbstractFunction<double, WorldVector<double> > *n)
  {
    NeumannBC *neumann = new NeumannBC(type, n, feSpace_);
    if(rhs_) 
      rhs_->getBoundaryManager()->addBoundaryCondition(neumann);
  }

  void ProblemScal::addRobinBC(BoundaryType type, 
			       DOFVector<double> *n,
			       DOFVector<double> *r)
  {
    RobinBC *robin = new RobinBC(type, n, r, feSpace_);
    if (rhs_) 
      rhs_->getBoundaryManager()->addBoundaryCondition(robin);
    if (systemMatrix_)
      systemMatrix_->getBoundaryManager()->addBoundaryCondition(robin);
  }

  void ProblemScal::addPeriodicBC(BoundaryType type) 
  {
    PeriodicBC *periodic = new PeriodicBC(type, feSpace_);

    if (systemMatrix_) 
      systemMatrix_->getBoundaryManager()->addBoundaryCondition(periodic);
    if (rhs_) 
      rhs_->getBoundaryManager()->addBoundaryCondition(periodic);
  }

  void ProblemScal::createMesh()
  {
    TEST_EXIT(Parameters::initialized())("parameters not initialized\n");

    // === create problems mesh ===
    ::std::string meshName("");

    GET_PARAMETER(0, name_ + "->info", "%d", &info_);
    GET_PARAMETER(0, name_ + "->mesh", &meshName);

    TEST_EXIT(meshName != "")
      ("no mesh name spezified\n");

    // get problem dimension
    int dim = 0;
    GET_PARAMETER(0, name_ + "->dim", "%d", &dim);
    TEST_EXIT(dim)("no problem dimension spezified!\n");

    // create the mesh
    mesh_ = NEW Mesh(meshName, dim);

    switch (dim) {
    case 1:
      coarseningManager_ = NEW CoarseningManager1d();
      refinementManager_ = NEW RefinementManager1d();
      break;
    case 2:
      coarseningManager_ = NEW CoarseningManager2d();
      refinementManager_ = NEW RefinementManager2d();
      break;
    case 3:
      coarseningManager_ = NEW CoarseningManager3d();
      refinementManager_ = NEW RefinementManager3d();
      break;
    default:
      ERROR_EXIT("invalid dim!\n");
    }
  }


  void ProblemScal::setMeshFromProblemVec(ProblemVec* pv, int i) 
  { 
    mesh_ = pv->getMesh(i);
    coarseningManager_ = pv->getCoarseningManager(i);
    refinementManager_ = pv->getRefinementManager(i);
  }


  Flag ProblemScal::markElements(AdaptInfo *adaptInfo) 
  { 
    if (marker_)
      return marker_->markMesh(adaptInfo, mesh_);
    else
      WARNING("no marker\n");

    return 0;
  }


  Flag ProblemScal::refineMesh(AdaptInfo *adaptInfo) 
  { 
    return refinementManager_->refineMesh(mesh_); 
  }


  Flag ProblemScal::coarsenMesh(AdaptInfo *adaptInfo) 
  {
    if (adaptInfo->isCoarseningAllowed(0)) {
      return coarseningManager_->coarsenMesh(mesh_);
    } else {
      WARNING("coarsening not allowed\n");
      return 0;
    }
  }

  void ProblemScal::solve(AdaptInfo *adaptInfo) 
  {
    FUNCNAME("Problem::solve()");
    if (!solver_) {
      WARNING("no solver\n");
      return;
    }

#ifdef _OPENMP
    double wtime = omp_get_wtime();
#endif

    clock_t first = clock();
    int iter = solver_->solve(matVec_, solution_, rhs_, leftPrecon_, rightPrecon_); 

#ifdef _OPENMP
    INFO(info_, 8)("solution of discrete system needed %.5f seconds system time / %.5f seconds wallclock time\n",
		   TIME_USED(first, clock()),
		   omp_get_wtime() - wtime);
#else
    INFO(info_, 8)("solution of discrete system needed %.5f seconds\n",
		   TIME_USED(first, clock()));
#endif

    adaptInfo->setSolverIterations(iter);
    adaptInfo->setMaxSolverIterations(solver_->getMaxIterations());
    adaptInfo->setSolverTolerance(solver_->getTolerance());
    adaptInfo->setSolverResidual(solver_->getResidual());
  }

  void ProblemScal::initialize(Flag initFlag,
			       ProblemScal *adoptProblem,
			       Flag adoptFlag) 
  {
    FUNCNAME("Problem::initialize()");

    // === create mesh ===
    if (mesh_) { 
      WARNING("mesh already created\n");
    } else {
      if (initFlag.isSet(CREATE_MESH) || 
	 ((!adoptFlag.isSet(INIT_MESH))&&
	  (initFlag.isSet(INIT_SYSTEM)||initFlag.isSet(INIT_FE_SPACE)))) {
	createMesh();
      } 
      
      if (adoptProblem && 
	 (adoptFlag.isSet(INIT_MESH) || 
	  adoptFlag.isSet(INIT_SYSTEM) ||
	  adoptFlag.isSet(INIT_FE_SPACE))) {
	TEST_EXIT(!mesh_)("mesh already created\n");
	mesh_ = adoptProblem->getMesh();
	refinementManager_ = adoptProblem->refinementManager_;
	coarseningManager_ = adoptProblem->coarseningManager_;
      }
    }

    if (!mesh_) 
      WARNING("no mesh created\n");

    // === create fespace ===
    if (feSpace_) {
      WARNING("feSpace already created\n");
    } else {
      if (initFlag.isSet(INIT_FE_SPACE) || (initFlag.isSet(INIT_SYSTEM) && !adoptFlag.isSet(INIT_FE_SPACE))) {
	createFESpace();
      } 
      if (adoptProblem && (adoptFlag.isSet(INIT_FE_SPACE) || adoptFlag.isSet(INIT_SYSTEM))) {
	TEST_EXIT(!feSpace_)("feSpace already created");
	feSpace_ = dynamic_cast<ProblemScal*>(adoptProblem)->getFESpace();
      }
    }

    if (!feSpace_) 
      WARNING("no feSpace created\n");

    // === create system ===
    if (initFlag.isSet(INIT_SYSTEM)) {
      createMatricesAndVectors();
    } 
    if (adoptProblem && adoptFlag.isSet(INIT_SYSTEM)) {
      TEST_EXIT(!solution_)("solution already created\n");
      TEST_EXIT(!rhs_)("rhs already created\n");
      TEST_EXIT(!systemMatrix_)("systemMatrix already created\n");
      solution_ = adoptProblem->getSolution();
      rhs_ = adoptProblem->getRHS();
      systemMatrix_  = adoptProblem->getSystemMatrix();
    }

    // === create solver ===
    if (solver_) {
      WARNING("solver already created\n");
    } else {
      if (initFlag.isSet(INIT_SOLVER)) {
	createSolver();
      } 
      if (adoptProblem && adoptFlag.isSet(INIT_SOLVER)) {
	TEST_EXIT(!solver_)("solver already created\n");
	solver_ = adoptProblem->getSolver();
      }
    }

    if (!solver_) 
      WARNING("no solver created\n");

    // === create estimator ===
    if (estimator_) {
      WARNING("estimator already created\n");
    } else {
      if (initFlag.isSet(INIT_ESTIMATOR)) {
	createEstimator();
      } 
      if (adoptProblem && adoptFlag.isSet(INIT_ESTIMATOR)) {
	TEST_EXIT(!estimator_)("estimator already created\n");
	estimator_ = adoptProblem->getEstimator();
      } 
    }

    if (!estimator_) 
      WARNING("no estimator created\n");

    // === create marker ===
    if (marker_) {
      WARNING("marker already created\n");
    } else {
      if (initFlag.isSet(INIT_MARKER)) {
	createMarker();
      } 
      if (adoptProblem && adoptFlag.isSet(INIT_MARKER)) {
	TEST_EXIT(!marker_)("marker already created\n");
	marker_ = adoptProblem->getMarker();
      } 
    }

    if (!marker_) 
      WARNING("no marker created\n");

    // === create file writer ===
    if (initFlag.isSet(INIT_FILEWRITER)) {
      createFileWriter();
    }

  

    // === read serialization and init mesh ===

    // There are two possiblities where the user can define a serialization
    // to be read from disk. Either by providing the parameter -rs when executing
    // the program or in the init file. The -rs parameter is always checked first,
    // because it can be added automatically when  rescheduling the program
    // before timeout of the runqueue.

    int readSerialization = 0;
    ::std::string serializationFilename = "";
    GET_PARAMETER(0, "argv->rs", &serializationFilename);

    // If the parameter -rs is set, we do nothing here, because the problem will be
    // deserialized in the constructor of a following AdaptInstationary initialization.
    if (!serializationFilename.compare("")) {
      GET_PARAMETER(0, name_ + "->input->read serialization", "%d", 
		    &readSerialization);
      if (readSerialization) {
	GET_PARAMETER(0, name_ + "->input->serialization filename", 
		      &serializationFilename);
	TEST_EXIT(serializationFilename != "")("no serialization file\n");

	MSG("Deserialization from file: %s\n", serializationFilename.c_str());
	::std::ifstream in(serializationFilename.c_str());
	deserialize(in);
	in.close();
      } else {
	if (initFlag.isSet(INIT_MESH) && mesh_ && !mesh_->isInitialized()) {
	  mesh_->initialize();
	}
    
	// === read value file and use it for the mesh values ===
	::std::string valueFilename("");
	GET_PARAMETER(0, mesh_->getName() + "->value file name", &valueFilename);     
	if (valueFilename.length()) {     
	  ValueReader::readValue(valueFilename.c_str(),
				 mesh_,
				 solution_,
				 mesh_->getMacroFileInfo());
	  mesh_->clearMacroFileInfo();
	}
	
		
	// === do global refinements ===
	if (initFlag.isSet(INIT_GLOBAL_REFINES)) {
	  int globalRefinements = 0;
	  GET_PARAMETER(0, mesh_->getName() + "->global refinements", "%d", &globalRefinements);
	  refinementManager_->globalRefine(mesh_, globalRefinements);	
	}
      }
    }
  }

  void ProblemScal::createFESpace()
  {
    // create finite element space
    int degree = 1;
    GET_PARAMETER(0, name_ + "->polynomial degree" ,"%d", &degree);
    feSpace_ = FiniteElemSpace::provideFESpace(NULL,
					       Lagrange::getLagrange(mesh_->getDim(), degree),
					       mesh_,
					       name_ + "->feSpace");  

    // create dof admin for vertex dofs if neccessary
    if (mesh_->getNumberOfDOFs(VERTEX) == 0) {
      DimVec<int> ln_dof(mesh_->getDim(), DEFAULT_VALUE, 0);
      ln_dof[VERTEX]= 1;
      mesh_->createDOFAdmin("vertex dofs", ln_dof);
    }
  }

  void ProblemScal::createMatricesAndVectors()
  {
    // === create vectors and system matrix ===
    systemMatrix_ = NEW DOFMatrix(feSpace_, feSpace_, "system matrix");
    rhs_          = NEW DOFVector<double>(feSpace_, "rhs");
    solution_     = NEW DOFVector<double>(feSpace_, "solution");

    solution_->refineInterpol(true);
    solution_->setCoarsenOperation(COARSE_INTERPOL);
    solution_->set(0.0);      /*  initialize u_h  !                      */

    // === create matVec ===
    matVec_ = NEW StandardMatVec<DOFMatrix, DOFVector<double> >(systemMatrix_);
  }

  void ProblemScal::createSolver()
  {
    // === create solver ===
    ::std::string solverType("no");
    GET_PARAMETER(0, name_ + "->solver", &solverType);
    OEMSolverCreator<DOFVector<double> > *solverCreator = 
      dynamic_cast<OEMSolverCreator<DOFVector<double> >*>(
							  CreatorMap<OEMSolver<DOFVector<double> > >::getCreator(solverType));
    TEST_EXIT(solverCreator)("no solver type\n");
    solverCreator->setName(name_ + "->solver");
    solver_ = solverCreator->create();
    solver_->initParameters();

    // === create preconditioners ===
    ::std::string preconType("no");
    Preconditioner<DOFVector<double> > *precon;
    GET_PARAMETER(0, name_ + "->solver->left precon", &preconType);
    CreatorInterface<PreconditionerScal> *preconCreator = 
      CreatorMap<PreconditionerScal>::getCreator(preconType);

    if (!preconCreator->isNullCreator()) {
      dynamic_cast<PreconditionerScalCreator*>(preconCreator)->setSizeAndRow(1, 0);
      dynamic_cast<PreconditionerScalCreator*>(preconCreator)->setName(name_ + "->solver->left precon");
    }
  
    precon = preconCreator->create();

    if (precon) {
      dynamic_cast<PreconditionerScal*>(precon)->setMatrix(&systemMatrix_);
      leftPrecon_ = precon;
    }

    preconType.assign("no");
    GET_PARAMETER(0, name_ + "->solver->right precon", &preconType);
    preconCreator = CreatorMap<PreconditionerScal>::getCreator(preconType);

    if (!preconCreator->isNullCreator()) {
      dynamic_cast<PreconditionerScalCreator*>(preconCreator)->setSizeAndRow(1, 0);
      dynamic_cast<PreconditionerScalCreator*>(preconCreator)->setName(name_ + "->solver->left precon");
    }

    precon = preconCreator->create();
    if (precon) {
      dynamic_cast<PreconditionerScal*>(precon)->setMatrix(&systemMatrix_);
      rightPrecon_ = precon;
    }

    // === create vector creator ===
    solver_->setVectorCreator(new DOFVector<double>::Creator(feSpace_));

  }

  void ProblemScal::createEstimator()
  {
    // create and set leaf data prototype
    mesh_->
      setElementDataPrototype(NEW LeafDataEstimatable(NEW LeafDataCoarsenable));

    // create estimator
    //estimator = NEW ResidualEstimator(name + "->estimator");

    // === create estimator ===
    ::std::string estimatorType("no");
    GET_PARAMETER(0, name_ + "->estimator", &estimatorType);
    EstimatorCreator *estimatorCreator = 
      dynamic_cast<EstimatorCreator*>(
				      CreatorMap<Estimator>::getCreator(estimatorType));
    if(estimatorCreator) {
      estimatorCreator->setName(name_ + "->estimator");
      if(estimatorType == "recovery") {
	dynamic_cast<RecoveryEstimator::Creator*>(estimatorCreator)->setSolution(solution_);
      }
      estimator_ = estimatorCreator->create();

      // init estimator
      estimator_->addSystem(systemMatrix_, solution_, rhs_);
    }
  }

  void ProblemScal::createMarker()
  {
    marker_ = dynamic_cast<Marker*>(Marker::createMarker(name_ + "->marker", -1));
  }

  void ProblemScal::createFileWriter()
  {
    fileWriters_.push_back(NEW FileWriter(name_ + "->output", mesh_, solution_));
    int writeSerialization = 0;
    GET_PARAMETER(0, name_ + "->output->write serialization", "%d", 
		  &writeSerialization);
    if(writeSerialization) {
      fileWriters_.push_back(NEW Serializer<ProblemScal>(this));
    }
  }

  void ProblemScal::estimate(AdaptInfo *adaptInfo) 
  {
    FUNCNAME("Problem::estimate()");

    if(estimator_) {
      clock_t first = clock();
      estimator_->estimate(adaptInfo->getTimestep());
      INFO(info_,8)("estimation of the error needed %.5f seconds\n",
		    TIME_USED(first,clock()));

      adaptInfo->setEstSum(estimator_->getErrorSum(), 0);

      adaptInfo->
	setTimeEstSum(estimator_->getTimeEst(), 0);

      adaptInfo->
	setEstMax(estimator_->getErrorMax(), 0);

      adaptInfo->
	setTimeEstMax(estimator_->getTimeEstMax(), 0);

    } else {
      WARNING("no estimator\n");
    }
  }

  void ProblemScal::buildAfterCoarsen(AdaptInfo *adaptInfo, Flag flag)
  {
    FUNCNAME("ProblemScal::buildAfterCoarsen()");

    clock_t first = clock();

    mesh_->dofCompress();

    MSG("%d DOFs for %s\n", 
	feSpace_->getAdmin()->getUsedSize(), 
	feSpace_->getName().c_str());

    Flag assembleFlag = 
      flag | 
      systemMatrix_->getAssembleFlag() | 
      rhs_->getAssembleFlag();

    if (useGetBound_)
      assembleFlag |= Mesh::FILL_BOUND;

    systemMatrix_->clear();
    rhs_->set(0.0);

    traversePtr_ = this;

    mesh_->traverse(-1, assembleFlag, &buildAfterCoarsenFct);

    // fill boundary conditions
    if (systemMatrix_->getBoundaryManager())
      systemMatrix_->getBoundaryManager()->initMatrix(systemMatrix_);
    if (rhs_->getBoundaryManager())
      rhs_->getBoundaryManager()->initVector(rhs_);
    if (solution_->getBoundaryManager())
      solution_->getBoundaryManager()->initVector(solution_);

    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh_, -1, 
					 Mesh::CALL_LEAF_EL | 
					 Mesh::FILL_BOUND |
					 Mesh::FILL_COORDS |
					 Mesh::FILL_DET |
					 Mesh::FILL_GRD_LAMBDA |
					 Mesh::FILL_NEIGH);
    // for all elements ...
    while (elInfo) {
      if (systemMatrix_->getBoundaryManager())
	systemMatrix_->getBoundaryManager()->fillBoundaryConditions(elInfo, systemMatrix_);
      if (rhs_->getBoundaryManager())
	rhs_->getBoundaryManager()->fillBoundaryConditions(elInfo, rhs_);
      if (solution_->getBoundaryManager())
	solution_->getBoundaryManager()->fillBoundaryConditions(elInfo, solution_);

      elInfo = stack.traverseNext(elInfo);
    }

    if (systemMatrix_->getBoundaryManager())
      systemMatrix_->getBoundaryManager()->exitMatrix(systemMatrix_);
    if (rhs_->getBoundaryManager())
      rhs_->getBoundaryManager()->exitVector(rhs_);
    if (solution_->getBoundaryManager())
      solution_->getBoundaryManager()->exitVector(solution_);

    INFO(info_, 8)("buildAfterCoarsen needed %.5f seconds\n",
		   TIME_USED(first,clock()));

    return;
  }

  int ProblemScal::buildAfterCoarsenFct(ElInfo *elInfo) 
  {
    const BoundaryType *bound;

    if (traversePtr_->getBoundUsed())
      bound = traversePtr_->getFESpace()->getBasisFcts()->getBound(elInfo, NULL);
    else
      bound = NULL;
  
    traversePtr_->getSystemMatrix()->assemble(1.0, elInfo, bound);
    traversePtr_->getRHS()->assemble(1.0, elInfo, bound);

    return 0;
  }

  void ProblemScal::writeResidualMesh(AdaptInfo *adaptInfo, const ::std::string name)
  {
    FUNCNAME("ProblemVec::writeResidualMesh()");

    Mesh *mesh = this->getMesh();
    FiniteElemSpace *fe = this->getFESpace();
    
    std::map<int, double> vec;
    
    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh,
					 -1, 
					 Mesh::CALL_LEAF_EL | 
					 Mesh::FILL_COORDS);
    
    while (elInfo) {		  
      Element *el = elInfo->getElement();
      double lError = el->getEstimation(0);
      
      vec[elInfo->getElement()->getIndex()] = lError;
      elInfo = stack.traverseNext(elInfo);
    }
    
    ElementFileWriter fw(name, mesh, fe, vec);
    fw.writeFiles(adaptInfo, true);    
  }

  void ProblemScal::serialize(::std::ostream &out) 
  {
    FUNCNAME("ProblemScal::serialize()");

    mesh_->serialize(out);
    solution_->serialize(out);
  }

  void ProblemScal::deserialize(::std::istream &in) 
  {
    FUNCNAME("ProblemScal::deserialize()");

    mesh_->deserialize(in);    
    solution_->deserialize(in);
  }

}