/******************************************************************************
 *
 * Mass-conserving ProblemStat for AMDiS
 *
 * Copyright (C) 2013 Dresden University of Technology. All Rights Reserved.
 * Web: https://fusionforge.zih.tu-dresden.de/projects/amdis/trunk/extensions
 *
 * Author: Simon Praetorius
 *
 * This file is provided AS IS with NO WARRANTY OF ANY KIND, INCLUDING THE
 * WARRANTY OF DESIGN, MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 *
 *
 * See also license.opensource.txt in the distribution.
 * 
 ******************************************************************************/

/** \file ProblemStatMassConserve.h */


#ifndef PROBLEM_STAT_MASS_CONSERVE
#define PROBLEM_STAT_MASS_CONSERVE

#include "ExtendedProblemStat.h"

/// Implementation of ProblemStat to allow for the conservation of mass of one
/// solution component, i.e. project the old-solution to the new mesh so that
/// the scalar-product is conserved: \f$ (u^*, v) = (u^\text{old}, v) \f$, where \f$ u^*,v \f$ live on the
/// new mesh and \f$ u^\text{old} \f$ lives on the old mesh.
struct ProblemStatMassConserve : public ExtendedProblemStat
{
  /// constructor: create a temporary problem that is initialized with 
  /// similar parameters as the regular one, except one additional component
  /// with an extra mesh for a temporary oldSolution
  ProblemStatMassConserve(std::string name_)
  : ExtendedProblemStat(name_),
    initialMeshAdoption(false),
    prob2(nullptr),
    comp(0)
  {
#ifdef HAVE_PARALLEL_DOMAIN_AMDIS
    TEST_EXIT(false)("Multi-Mesh approach does not work in parallel mode for the moment. Need to be implemented!\n");
#endif
    
    Parameters::get(name_ + "->mass conservation component", comp);
    
    // set parameters for the temporary problem
    std::string name2 = name_ + "_tmp";
    int zero = 0;
    int one = 1;
    int two = 2;
    double tol = 1.e-8;
    int degree = 1;
    int dim = 1;
    std::string meshName = "";
    
    Parameters::get(name_ + "->polynomial degree[" + boost::lexical_cast<std::string>(comp) + "]", degree);
    Parameters::get(name_ + "->dim", dim);
    Parameters::get(name_ + "->mesh", meshName);
    
    Parameters::set(name2 + "->components", two);
    Parameters::set(name2 + "->polynomial degree[0]", degree);
    Parameters::set(name2 + "->polynomial degree[1]", degree);
    Parameters::set(name2 + "->dim", dim);
    Parameters::set(name2 + "->mesh", meshName);
    Parameters::set(name2 + "->refinement set[0]", one);
    Parameters::set(name2 + "->refinement set[1]", two);
#if defined HAVE_UMFPACK || defined HAVE_PARALLEL_PETSC || defined HAVE_PETSC
    std::string solverName = "direct";
#else
    std::string solverName = "cg";
#endif
    Parameters::set(name2 + "->solver", solverName);
    Parameters::set(name2 + "->tolerance", tol);
    Parameters::set(name2 + "->info", zero);
    
    // create a new temporary problem
    prob2 = new ProblemStat(name2);
  }
  
  
  ~ProblemStatMassConserve()
  {
    if (prob2 != nullptr) {
      delete prob2;
      prob2 = nullptr;
    }
  }
  
  
  /// make shure that temporary problem has an equal mesh as regular problem
  /// and initialize this mesh only once
  void initialize(Flag initFlag,
		  ProblemStatSeq *adoptProblem = NULL,
		  Flag adoptFlag = INIT_NOTHING) override
  {
    ExtendedProblemStat::initialize(initFlag - INIT_MESH, adoptProblem, adoptFlag); // create meshes for prob
    
    std::vector<Mesh*> meshes2;
    meshes2.push_back(getMesh(comp));
    prob2->setComponentMesh(0, getMesh(comp));
      
    int refSet = -1;
    Parameters::get(name + "_tmp->refinement set[0]", refSet);
    if (refSet < 0)
      refSet = 0;
    
    std::map<int, Mesh*> meshForRefinementSet;
    meshForRefinementSet[refSet] = meshes2[0];
        
    for (int i = 1; i < prob2->getNumComponents(); i++) {
      refSet = -1;
      Parameters::get(name + "_tmp->refinement set[" + lexical_cast<string>(i) + "]", 
		      refSet);
      if (refSet < 0)
	refSet = 0;
      
      if (meshForRefinementSet[refSet] == nullptr) {
	Mesh *newMesh = new Mesh(getMesh(comp)->getName(), getMesh(comp)->getDim());
	meshForRefinementSet[refSet] = newMesh;
	meshes2.push_back(newMesh);
      }
      prob2->setComponentMesh(i, meshForRefinementSet[refSet]);
    }
    prob2->setMeshes(meshes2);    
    prob2->initialize(INIT_ALL - INIT_MESH);
    
    for (size_t i = 0; i < meshes2.size(); i++) {
      int globalRefinements = 0;

      // If AMDiS is compiled for parallel computations, the global refinements are
      // ignored here. Later, each rank will add the global refinements to its
      // private mesh.
#ifndef HAVE_PARALLEL_DOMAIN_AMDIS
      Parameters::get(meshes2[i]->getName() + "->global refinements",
		      globalRefinements);
#endif

      bool initMesh = true;

      // Initialize the meshes if there is no serialization file.
      if (initMesh && meshes2[i] && !(meshes2[i]->isInitialized())) {
	meshes2[i]->initialize();
	prob2->getRefinementManager()->globalRefine(meshes2[i], globalRefinements);
      }
    }
    
    
    for (int i = 0; i < prob2->getNumComponents(); i++)
      prob2->getSolution(i)->setCoarsenOperation(NO_OPERATION);
    
    fillOperators();
  }
  
  
  /// in the first iteration copy initial-solution to temporary problem
  /// and prepare the temporary mesh
  void buildBeforeRefine(AdaptInfo* adaptInfo, 
			 Flag markFlag) override
  {
    if (!initialMeshAdoption) {
      MeshStructure meshStructure2;
      meshStructure2.init(getMesh(comp));
      prob2->getCoarseningManager()->globalCoarsen(prob2->getMesh(1), -2);
      meshStructure2.fitMeshToStructure(prob2->getMesh(1), prob2->getRefinementManager());
      
      prob2->getSolution(1)->interpol(getSolution(comp));
      initialMeshAdoption = true;
    }
    
    ExtendedProblemStat::buildBeforeRefine(adaptInfo, markFlag);
  }
  
  
  /// solve temporary problem to project solution to new mesh and then
  /// copy temporary solution back to regular solution
  void buildAfterCoarsen(AdaptInfo* adaptInfo, 
			 Flag markFlag, 
			 bool asmM, 
			 bool asmV) override
  {  
    prob2->buildAfterCoarsen(adaptInfo, markFlag, true, true);
    solve2(adaptInfo);
    getSolution(comp)->copy(*prob2->getSolution(0));
    
    // call original buildAfterCoarsen
    ExtendedProblemStat::buildAfterCoarsen(adaptInfo, markFlag, asmM, asmV); 
  }
  
  
  /// copy solution to temporary problem after regular solve
  void solve(AdaptInfo *adaptInfo, 
	     bool createMatrixData = true, 
	     bool storeMatrixData = false) override
  {
    ExtendedProblemStat::solve(adaptInfo, createMatrixData, storeMatrixData);
    
    MeshStructure meshStructure2;
    meshStructure2.init(getMesh(comp));
    prob2->getCoarseningManager()->globalCoarsen(prob2->getMesh(1), -2);
    meshStructure2.fitMeshToStructure(prob2->getMesh(1), prob2->getRefinementManager());
    
    prob2->getSolution(1)->interpol(getSolution(comp));
  }
  
  /// solve only one block of temporary problem
  void solve2(AdaptInfo *adaptInfo, 
	      bool createMatrixData = true, 
	      bool storeMatrixData = false)
  {
    SolverMatrix<Matrix<DOFMatrix*> > solverMatrix_;
    Matrix<DOFMatrix*> mat(1,1);
    mat[0][0] = prob2->getSystemMatrix(0,0);
    solverMatrix_.setMatrix(mat);
    
    std::vector<const FiniteElemSpace*> feSpaces0; 
    feSpaces0.push_back(prob2->getFeSpace(0));
    
    SystemVector vec_sol("solution_0", feSpaces, 1, false);
    vec_sol.setDOFVector(0, prob2->getSolution(0));
    SystemVector vec_rhs("rhs_0", feSpaces, 1, false);
    vec_rhs.setDOFVector(0, prob2->getRhsVector(0));
    
    // solve only one block of the system
    prob2->getSolver()->solveSystem(solverMatrix_, vec_sol, vec_rhs, 
			createMatrixData, storeMatrixData);
  }

  
  /// add mass-matrix and transfer-operator to temporary problem, to
  /// project solution to new mesh
  void fillOperators()
  {
    // conversion propblem 2
    const FiniteElemSpace* feSpace2_0 = prob2->getFeSpace(0);
    const FiniteElemSpace* feSpace2_1 = prob2->getFeSpace(1);

    Operator *opMnew0 = new Operator(feSpace2_0, feSpace2_0);
    opMnew0->addTerm(new Simple_ZOT);
    Operator *opMold01 = new Operator(feSpace2_0, feSpace2_1);
    opMold01->addTerm(new VecAtQP_ZOT(prob2->getSolution(1)));
    
    // assemble only one block of the system
    prob2->addMatrixOperator(*opMnew0,0,0);
    prob2->addVectorOperator(*opMold01, 0); 
  }
  
private:
  /// true if initial-solution/-mesh is copied to temporary problem
  bool initialMeshAdoption; 	
  
  /// temporary problem for the projection of the old-solution to the new mesh
  ProblemStat* prob2;
  
  /// component the should be copnserved
  int comp;
};

#endif // PROBLEM_STAT_MASS_CONSERVE