/******************************************************************************
 *
 * Extension of AMDiS - Adaptive multidimensional simulations
 *
 * Copyright (C) 2013 Dresden University of Technology. All Rights Reserved.
 * Web: https://fusionforge.zih.tu-dresden.de/projects/amdis
 *
 * Authors: Simon Praetorius et al.
 *
 * 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.
 * 
 ******************************************************************************/


#ifndef EXTENSIONS_REFINEMENT_H
#define EXTENSIONS_REFINEMENT_H

#include "ElementFunction.h"

using namespace AMDiS;


/** \brief
 * Abstract class that can be passed to RefinementLevel* as indicator where
 * to refine the mesh up to which level. It is an AbstractFunction that 
 * overloads the operator() method to return a level or a meshsize depending
 * on the coords/data passed to the operator.
 * You can switch between meshsize and level with the methods hToLevel(double) and 
 * levelToH(int)
 **/
template<typename T, typename T2>
class MeshRefinementFunction : public AbstractFunction<T2, T>
{
public:

  MeshRefinementFunction(Mesh* mesh_) :
    AbstractFunction<T2, T>(0),
    mesh(mesh_), 
    globalSize(0)
  {
      h0 = getMacroMeshSize(mesh);
      reduction = 1.0 / sqrt(2.0); // if dim==2
  }

  int getGlobalSize() { return globalSize; }

  double meshSize() { return h0; }
  
  virtual void init(ElInfo* elInfo) {}

  virtual T2 operator()(const T &value) const { return globalSize; }
  virtual T2 getLevel() const { return globalSize; }
  
  virtual double indicator(const T &value) const { return 1.0; }

protected:

  int hToLevel(double h) {
      int level = static_cast<int>(floor(log(h / h0) / log(reduction)));
      return level;
  }

  double levelToH(int level) {
      double h = pow(reduction,level)*h0;
      return h;
  }

  double getMacroMeshSize(Mesh* mesh) {
      FixVec<WorldVector<double>, VERTEX> coords = mesh->getMacroElement(0)->getCoord();
      double h = 0.0;
      for (int i = 0; i < coords.size(); ++i)
          for (int j = i + 1; j < coords.size(); ++j)
              h = std::max(h, norm(coords[i]-coords[j]));
      return h;
  }

protected:

  Mesh* mesh;

  int globalSize;

  double h0;
  double reduction;
};

/// Operations performed before and after refinement/ coarsening
struct StandardRefineOperation
{
  virtual ~StandardRefineOperation() {};
  virtual void beforeRefine(AdaptInfo* adaptInfo, Flag markFlag) {}
  virtual void beforeCoarsen(AdaptInfo* adaptInfo, Flag markFlag) {}
  virtual void afterCoarsen(AdaptInfo* adaptInfo, Flag markFlag) {}  
};

/// vector of operations performed before and after refinement/ coarsening
struct CoupledRefineOperation : public StandardRefineOperation
{
  CoupledRefineOperation(std::vector<StandardRefineOperation*> operations_) 
    : operations(operations_) {}
    
  void beforeRefine(AdaptInfo* adaptInfo, Flag markFlag) override 
  {
    for (size_t i = 0; i < operations.size(); i++)
	operations[i]->beforeRefine(adaptInfo, markFlag);
  }
  
  void beforeCoarsen(AdaptInfo* adaptInfo, Flag markFlag) override
  {
    for (size_t i = 0; i < operations.size(); i++)
      operations[i]->beforeCoarsen(adaptInfo, markFlag);
  }
  void afterCoarsen(AdaptInfo* adaptInfo, Flag markFlag) override
  {
    for (size_t i = 0; i < operations.size(); i++)
      operations[i]->afterCoarsen(adaptInfo, markFlag);
  }

protected:
   std::vector<StandardRefineOperation*> operations;
}; 


/** \brief
 * Base class for Refinement structure to perform local anisotropic refinement
 */
template<typename T, typename T2>
class RefinementLevel
{
public:

  RefinementLevel(const FiniteElemSpace *feSpace_, 
		  MeshRefinementFunction<T,T2>* refineFct_ = NULL) :
    feSpace(feSpace_), 
    refineFct(refineFct_),
    adaptInfo(NULL),
    refineOperation(NULL),
    numRefinements0(15),
    globalRefined(false)
  {	
    FUNCNAME("RefinementLevel::RefinementLevel()");
    
    mesh = feSpace->getMesh();
    switch (mesh->getDim()) {
    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");
    }

    numRefinements = numRefinements0;
    
    refineOperation = new StandardRefineOperation;
  }
  
  virtual ~RefinementLevel() {
    delete coarseningManager;
    delete refinementManager;
    if (refineOperation)
      delete refineOperation;
  }
  
  void setRefineOperation(AdaptInfo* adaptInfo_,
			  StandardRefineOperation* refineOperation_)
  {
    if (refineOperation)
      delete refineOperation;
    
    adaptInfo = adaptInfo_;
    refineOperation = refineOperation_;
  }

  void refine(bool onlyRefine= false) 
  {
    FUNCNAME("RefinementLevel::refine()");

    if (!globalRefined && refineFct) {
      MSG("nr of global refinements: %d\n", refineFct->getGlobalSize());
      refinementManager->globalRefine(mesh, refineFct->getGlobalSize());
      globalRefined = true;
    }
    double minH = 0.0, maxH = 1.0;
    int minLevel = 100, maxLevel = 0;
    
    // build mesh for phasefield-function
    bool meshChanged = true;
    Flag markFlag;
    int oldNr = 0, oldOldNr = 0;
    int i = 0;
    while (meshChanged && i < numRefinements) {
      markElements(markFlag);
      meshChanged = refineMesh(markFlag, onlyRefine);
      
      calcMeshSizes(minH, maxH, minLevel, maxLevel, false); 
      int nr = mesh->getNumberOfVertices();
      meshChanged = meshChanged && oldOldNr!=nr && oldNr!=nr;
      if (meshChanged) {
	MSG("(local) mesh sizes: [%f, %f], Vs: %d, ELs: %d\n", 
          minH, maxH, nr, mesh->getNumberOfElements());
      }
      i++;
			
      oldOldNr = oldNr;
      oldNr = nr; 
    }
    calcMeshSizes(minH, maxH, minLevel, maxLevel, true); 
    MSG("Final (global) mesh: [%f, %f], Vs: %d, ELs: %d, Level: [%d, %d]\n",
      minH, maxH, mesh->getNumberOfVertices(), mesh->getNumberOfElements(), minLevel, maxLevel);
  }

  void refine(int numRefinements_, bool onlyRefine= false) { 
    numRefinements = numRefinements_; 
    refine(onlyRefine);  
    numRefinements = numRefinements0; 
  }
  
  
  void initial_refine()
  {
    
  }

  int getNumRefinements(){
      return numRefinements;
  }

  void calcMeshSizes(double& minH, double& maxH, int& minLevel, int& maxLevel, bool allReduce = false) 
  {
    FixVec<WorldVector<double>, VERTEX> coords(mesh->getDim(), NO_INIT);

    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh, -1, Mesh::CALL_LEAF_EL | Mesh::FILL_COORDS);
    minH = 1e15; maxH = 0.0;
    int k = 0;
    minLevel = 100;
    maxLevel = 0;
    while (elInfo) {
      maxLevel = std::max(maxLevel,elInfo->getLevel());
      minLevel = std::min(minLevel,elInfo->getLevel());
      coords = elInfo->getCoords();
      double h = 0.0;
      for (int i = 0; i < coords.size(); i++) {
        for (int j = 0; j < coords.size(); j++) {
	  if (i != j)
            h = std::max(h,norm(coords[i]-coords[j]));
        }
      }
      minH = std::min(h, minH);
      maxH = std::max(h, maxH);
      elInfo = stack.traverseNext(elInfo);
      k++;
    }
    minLevel += mesh->getMacroElementLevel();
    maxLevel += mesh->getMacroElementLevel();
    
#ifdef HAVE_PARALLEL_DOMAIN_AMDIS
    if (allReduce) {
      Parallel::mpi::globalMin(minH);
      Parallel::mpi::globalMax(maxH);
      Parallel::mpi::globalMin(minLevel);
      Parallel::mpi::globalMax(maxLevel);
    }
#endif
  }


  double calcMeshSize(ElInfo *elInfo, bool allReduce = false) 
  {
    FixVec<WorldVector<double>, VERTEX> coords(mesh->getDim(), NO_INIT);
    coords = elInfo->getCoords();
    double h = 0.0;
    for (int i = 0; i < coords.size(); i++) {
      for (int j = 0; j < coords.size(); j++) {
        if (i != j)
          h = std::max(h,norm(coords[i]-coords[j]));
      }
    }
#ifdef HAVE_PARALLEL_DOMAIN_AMDIS
    if (allReduce)
      Parallel::mpi::globalMax(h);
#endif

    return h;
  }


  int calcMark(double refineH, double currentH)
  {
    return (refineH < currentH ? 
        1 : (refineH > currentH * (mesh->getDim() == 1 ? 
          2.0 : (mesh->getDim() == 2 ? 
          sqrt(2.0) : 
          sqrt(2.0)/2.0 + 0.5)) ? 
        -1 : 
        0));
  }


  virtual int calcMark(int refineLevel, int currentLevel)
  {
    int levelDiff = refineLevel - currentLevel;
    return (levelDiff > 0 ? 1 : (levelDiff < 0 ? -1 : 0));
  }

	
  bool refineMesh(Flag markFlag, bool onlyRefine) 
  {
    int oldSize = mesh->getNumberOfVertices();
    refineOperation->beforeRefine(adaptInfo, markFlag);
    if (markFlag.isSet(1))
      refinementManager->refineMesh(mesh);
    refineOperation->beforeCoarsen(adaptInfo, markFlag);
    if (markFlag.isSet(2) && !onlyRefine)
      coarseningManager->coarsenMesh(mesh);
    refineOperation->afterCoarsen(adaptInfo, markFlag);
    if (markFlag.isSet(1) || markFlag.isSet(2)) {
      int newSize = mesh->getNumberOfVertices();
      if (oldSize != newSize) 
        return true;
    }
    return false;
  }

  virtual void markElements(Flag &markFlag) = 0;

  void setGlobalRefined(bool refined) { globalRefined = refined; }
  RefinementManager* getRefinementManager() { return refinementManager; }
  CoarseningManager* getCoarseningManager() { return coarseningManager; }

protected:

  const FiniteElemSpace *feSpace;
  Mesh* mesh;
  MeshRefinementFunction<T,T2>* refineFct;
  RefinementManager* refinementManager;
  CoarseningManager* coarseningManager;
  
  AdaptInfo* adaptInfo;
  StandardRefineOperation* refineOperation;

  int numRefinements;
  int numRefinements0;
  
  bool globalRefined;
};



/** \brief
 * Base class for Refinement structure to perform local anisotropic refinement
 */
class RefinementLocal
{
public:

  RefinementLocal(Mesh *mesh_) :
    mesh(mesh_), 
    adaptInfo(nullptr),
    refineOperation(nullptr),
    numRefinements0(15),
    globalRefined(false)
  {	
    FUNCNAME("RefinementLevel::RefinementLevel()");
    
    switch (mesh->getDim()) {
    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");
    }

    numRefinements = numRefinements0;
    
    refineOperation = new StandardRefineOperation;
  }
  
  virtual ~RefinementLocal() {
    delete coarseningManager;
    delete refinementManager;
    if (refineOperation)
      delete refineOperation;
  }
  
  void setRefineOperation(AdaptInfo* adaptInfo_,
			  StandardRefineOperation* refineOperation_)
  {
    if (refineOperation)
      delete refineOperation;
    
    adaptInfo = adaptInfo_;
    refineOperation = refineOperation_;
  }

  void setOnlyRefine(bool onlyRefine_)
  {
    onlyRefine = onlyRefine_;
  }
  
  template<typename MeshRefFunc>
  void refine(MeshRefFunc* refineFct) 
  {
    FUNCNAME("RefinementLevel::refine()");

    refineFct->init(mesh);
    if (!globalRefined) {
      MSG("nr of global refinements: %d\n", refineFct->getGlobalSize());
      refinementManager->globalRefine(mesh, refineFct->getGlobalSize());
      globalRefined = true;
    }
    double minH = 0.0, maxH = 1.0;
    int minLevel = 100, maxLevel = 0;
    
    // build mesh for phasefield-function
    bool meshChanged = true;
    Flag markFlag;
    int oldNr = 0, oldOldNr = 0;
    int i = 0;
    while (meshChanged && i < numRefinements) {
      markElements(refineFct, markFlag);
      meshChanged = refineMesh(markFlag, onlyRefine);
      
      calcMeshSizes(minH, maxH, minLevel, maxLevel, false); 
      int nr = mesh->getNumberOfVertices();
      meshChanged = meshChanged && oldOldNr!=nr && oldNr!=nr;
      if (meshChanged) {
	MSG("(local) mesh sizes: [%f, %f], Vs: %d, ELs: %d\n", 
          minH, maxH, nr, mesh->getNumberOfElements());
      }
      i++;
			
      oldOldNr = oldNr;
      oldNr = nr; 
    }
    calcMeshSizes(minH, maxH, minLevel, maxLevel, true); 
    MSG("Final (global) mesh: [%f, %f], Vs: %d, ELs: %d, Level: [%d, %d]\n",
      minH, maxH, mesh->getNumberOfVertices(), mesh->getNumberOfElements(), minLevel, maxLevel);
  }

  template<typename MeshRefFunc>
  void refine(int numRefinements_, MeshRefFunc* refineFct, bool onlyRefine_= false) { 
    numRefinements = numRefinements_; 
    setOnlyRefine(onlyRefine_);
    refine(refineFct);  
    numRefinements = numRefinements0; 
  }
  
  template<typename MeshRefFunc>
  void markElements(MeshRefFunc* refineFct, Flag &markFlag) 
  {    
    bool elMarkRefine = false, elMarkCoarsen = false;

    Flag traverseFlag = Mesh::CALL_LEAF_EL | Mesh::FILL_COORDS;
    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh, -1, traverseFlag);
    while(elInfo) {
      refineFct->init(elInfo);
      int refineLevel = refineFct->getLevel();
      int oldLevel = elInfo->getLevel();
      elInfo->getElement()->setMark( calcMark(refineLevel, oldLevel) );

      elMarkRefine |= elInfo->getElement()->getMark() == 1;
      elMarkCoarsen |= elInfo->getElement()->getMark() == -1;
      elInfo = stack.traverseNext(elInfo);
    }
      
    markFlag= 0;
    if(elMarkRefine) 
	    markFlag= 1;
    if(elMarkCoarsen) 
	    markFlag|= 2;
  }

  int getNumRefinements() {
      return numRefinements;
  }

  void calcMeshSizes(double& minH, double& maxH, int& minLevel, int& maxLevel, bool allReduce = false) 
  {
    FixVec<WorldVector<double>, VERTEX> coords(mesh->getDim(), NO_INIT);

    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh, -1, Mesh::CALL_LEAF_EL | Mesh::FILL_COORDS);
    minH = 1e15; maxH = 0.0;
    int k = 0;
    minLevel = 100;
    maxLevel = 0;
    while (elInfo) {
      maxLevel = std::max(maxLevel,elInfo->getLevel());
      minLevel = std::min(minLevel,elInfo->getLevel());
      coords = elInfo->getCoords();
      double h = 0.0;
      for (int i = 0; i < coords.size(); i++) {
        for (int j = 0; j < coords.size(); j++) {
	  if (i != j)
            h = std::max(h,norm(coords[i]-coords[j]));
        }
      }
      minH = std::min(h, minH);
      maxH = std::max(h, maxH);
      elInfo = stack.traverseNext(elInfo);
      k++;
    }
    minLevel += mesh->getMacroElementLevel();
    maxLevel += mesh->getMacroElementLevel();
    
#ifdef HAVE_PARALLEL_DOMAIN_AMDIS
    if (allReduce) {
      Parallel::mpi::globalMin(minH);
      Parallel::mpi::globalMax(maxH);
      Parallel::mpi::globalMin(minLevel);
      Parallel::mpi::globalMax(maxLevel);
    }
#endif
  }


  double calcMeshSize(ElInfo *elInfo, bool allReduce = false) 
  {
    FixVec<WorldVector<double>, VERTEX> coords(mesh->getDim(), NO_INIT);
    coords = elInfo->getCoords();
    double h = 0.0;
    for (int i = 0; i < coords.size(); i++) {
      for (int j = 0; j < coords.size(); j++) {
        if (i != j)
          h = std::max(h,norm(coords[i]-coords[j]));
      }
    }
#ifdef HAVE_PARALLEL_DOMAIN_AMDIS
    if (allReduce)
      Parallel::mpi::globalMax(h);
#endif

    return h;
  }


  int calcMark(double refineH, double currentH)
  {
    return (refineH < currentH ? 
        1 : (refineH > currentH * (mesh->getDim() == 1 ? 
          2.0 : (mesh->getDim() == 2 ? 
          sqrt(2.0) : 
          sqrt(2.0)/2.0 + 0.5)) ? 
        -1 : 
        0));
  }


  virtual int calcMark(int refineLevel, int currentLevel)
  {
    int levelDiff = refineLevel - currentLevel;
    return (levelDiff > 0 ? 1 : (levelDiff < 0 ? -1 : 0));
  }

	
  bool refineMesh(Flag markFlag, bool onlyRefine) 
  {
    int oldSize = mesh->getNumberOfVertices();
    refineOperation->beforeRefine(adaptInfo, markFlag);
    if (markFlag.isSet(1))
      refinementManager->refineMesh(mesh);
    refineOperation->beforeCoarsen(adaptInfo, markFlag);
    if (markFlag.isSet(2) && !onlyRefine)
      coarseningManager->coarsenMesh(mesh);
    refineOperation->afterCoarsen(adaptInfo, markFlag);
    if (markFlag.isSet(1) || markFlag.isSet(2)) {
      int newSize = mesh->getNumberOfVertices();
      if (oldSize != newSize) 
        return true;
    }
    return false;
  }

  void setGlobalRefined(bool refined) { globalRefined = refined; }
  RefinementManager* getRefinementManager() { return refinementManager; }
  CoarseningManager* getCoarseningManager() { return coarseningManager; }

protected:

  Mesh* mesh;
  RefinementManager* refinementManager;
  CoarseningManager* coarseningManager;
  
  AdaptInfo* adaptInfo;
  StandardRefineOperation* refineOperation;

  int numRefinements;
  int numRefinements0;
  
  bool globalRefined;
  bool onlyRefine;
};

#include "Refinement_Level.h"
// #include "Refinement_MeshSize.h"

#endif // EXTENSIONS_REFINEMENT_H