#pragma once

#include <cmath>
#include <functional>
#include <limits>
#include <map>
#include <memory>
#include <numeric>
#include <type_traits>
#include <unordered_map>
#include <utility>
#include <vector>

#include <dune/common/fvector.hh>
#include <dune/common/hash.hh>

#include <dune/grid/common/geometry.hh>
#include <dune/grid/common/rangegenerators.hh>

#include <amdis/Output.hpp>
#include <amdis/common/ConcurrentCache.hpp>
#include <amdis/functions/FunctionFromCallable.hpp>
#include <amdis/operations/Assigner.hpp>
#include <amdis/typetree/Traversal.hpp>

namespace AMDiS {
namespace Impl {

// Hash function for cache container
struct CoordHasher
{
  template <class LocalCoord>
  std::size_t operator()(LocalCoord const& coord) const
  {
    std::size_t seed = 0;
    for (std::size_t i = 0; i < coord.size(); ++i)
      Dune::hash_combine(seed, coord[i]);
    return seed;
  }
};

} // end namespace Impl


template <class C, class B>
DataTransfer<C,B>::DataTransfer(std::shared_ptr<B const> basis)
  : basis_(std::move(basis))
  , mapper_(basis_->gridView().grid(), Dune::mcmgElementLayout())
  , nodeDataTransfer_()
{}


template <class C, class B>
void DataTransfer<C,B>::
preAdapt(C const& coeff, bool mightCoarsen)
{
  GridView gv = basis_->gridView();
  LocalView lv = basis_->localView();
  auto const& idSet = gv.grid().localIdSet();

  for_each_leaf_node(lv.tree(), [&](auto const& node, auto const& tp) {
    nodeDataTransfer_[tp].preAdaptInit(lv, coeff, node);
  });

  // Make persistent DoF container
  persistentContainer_.clear(); // Redundant if postAdapt was correctly called last cycle
  for (const auto& e : elements(gv))
  {
    auto it = persistentContainer_.emplace(idSet.id(e), TypeTree::treeContainer<NodeElementData,true>(lv.tree()));

    lv.bind(e);
    auto& treeContainer = it.first->second;
    for_each_leaf_node(lv.tree(), [&](auto const& node, auto const& tp) {
      nodeDataTransfer_[tp].cacheLocal(treeContainer[tp]);
    });
  }

  if (!mightCoarsen)
    return;

  // Interpolate from possibly vanishing elements
  auto maxLevel = gv.grid().maxLevel();
  using std::sqrt;
  typename Grid::ctype const checkInsideTolerance = sqrt(std::numeric_limits<typename Grid::ctype>::epsilon());
  for (auto const& e : elements(gv, typename C::Traits::PartitionSet{}))
  {
    auto father = e;
    while (father.mightVanish() && father.hasFather())
    {
      father = father.father();
      auto it = persistentContainer_.emplace(idSet.id(father), TypeTree::treeContainer<NodeElementData,true>(lv.tree()));
      if (!it.second)
        continue;

      auto& treeContainer = it.first->second;
      auto geo = father.geometry();
      bool init = true; // init flag for first call on new father element
      bool restrictLocalCompleted = false;
      auto hItEnd = father.hend(maxLevel);
      for (auto hIt = father.hbegin(maxLevel); hIt != hItEnd; ++hIt) {
        if (!hIt->isLeaf())
          continue;

        auto const& child = *hIt;
        auto search = persistentContainer_.find(idSet.id(child));
        assert(search != persistentContainer_.end());
        auto const& childContainer = search->second;
        lv.bind(child);

        auto const& childGeo = child.geometry();
        auto refTypeId = Dune::referenceElement(childGeo).type().id();

        using BoolCoordPair = std::pair<bool, LocalCoordinate>;
        using CacheImp = std::unordered_map<LocalCoordinate, BoolCoordPair, Impl::CoordHasher>;
        using ChildCache = ConcurrentCache<LocalCoordinate, BoolCoordPair, ConsecutivePolicy, CacheImp>;

        // Transfers input father-local point x into child-local point y
        // Returns false if x is not inside the child
        auto xInChild = [&](LocalCoordinate const& x) -> BoolCoordPair {
          LocalCoordinate local = childGeo.local(geo.global(x));
          // TODO(FM): Using an implementation detail as workaround for insufficient
          //   tolerance, see https://gitlab.dune-project.org/core/dune-grid/issues/84
          bool isInside = Dune::Geo::Impl::checkInside(refTypeId, Geometry::coorddimension, local, checkInsideTolerance);
          return BoolCoordPair(isInside, std::move(local));
        };
        // TODO(FM): Disable for single-node basis
        ChildCache childCache;
        auto xInChildCached = [&](LocalCoordinate const& x) -> BoolCoordPair {
          return childCache.get(x, [&](LocalCoordinate const& x) { return xInChild(x); });
        };

        restrictLocalCompleted = true;
        for_each_leaf_node(lv.tree(), [&](auto const& node, auto const& tp) {
          restrictLocalCompleted &=
            nodeDataTransfer_[tp].restrictLocal(father, treeContainer[tp], xInChildCached,
                                                childContainer[tp], init);
        });
        init = false;
      }
      // test if restrictLocal was completed on all nodes
      assert(restrictLocalCompleted);

    } // end while (father.mightVanish)
  } // end for (elements)
}


template <class C, class B>
void DataTransfer<C,B>::adapt(C& coeff)
{
  coeff.resize();

  // No data was saved before adapting the grid, make
  // sure to call DataTransfer::preAdapt before calling adapt() on the grid
  if (persistentContainer_.empty())
    return;

  GridView gv = basis_->gridView();
  LocalView lv = basis_->localView();
  auto const& idSet = gv.grid().localIdSet();
  for_each_leaf_node(lv.tree(), [&](auto const& node, auto const& tp) {
    nodeDataTransfer_[tp].adaptInit(lv, coeff, node);
  });

  mapper_.update();
  std::vector<bool> finished(mapper_.size(), false);
  for (const auto& e : elements(gv, typename C::Traits::PartitionSet{}))
  {
    auto index = mapper_.index(e);
    if (finished[index])
      continue;

    auto it = persistentContainer_.find(idSet.id(e));

    // Data already exists and no interpolation is required
    if (it != persistentContainer_.end()) {
      lv.bind(e);
      auto const& treeContainer = it->second;
      for_each_leaf_node(lv.tree(), [&](auto const& node, auto const& tp) {
        nodeDataTransfer_[tp].copyLocal(treeContainer[tp]);
      });
      finished[index] = true;
      continue;
    }

    // Data needs to be interpolated
    auto father = e;
    while (father.hasFather() && father.isNew())
      father = father.father();

    auto maxLevel = gv.grid().maxLevel();
    auto fatherGeo = father.geometry();
    bool init = true; // init flag for first call on new father element

    auto father_it = persistentContainer_.find(idSet.id(father));
    assert(father_it != persistentContainer_.end());
    auto const& treeContainer = father_it->second;

    auto hItEnd = father.hend(maxLevel);
    for (auto hIt = father.hbegin(maxLevel); hIt != hItEnd; ++hIt) {
      if (!hIt->isLeaf())
        continue;

      auto const& child = *hIt;
      lv.bind(child);

      // coordinate transform from child to father element
      auto xInFather = [&fatherGeo, childGeo = child.geometry()]
        (LocalCoordinate const& x) -> LocalCoordinate
      {
        return fatherGeo.local(childGeo.global(x));
      };

      for_each_leaf_node(lv.tree(), [&](auto const& node, auto const& tp) {
        nodeDataTransfer_[tp].prolongLocal(father, treeContainer[tp], xInFather, init);
      });

      finished[mapper_.index(child)] = true;
      init = false;
    }
  } // end for (elements)

  coeff.finish();
}


template <class C, class B>
void DataTransfer<C,B>::postAdapt(C&)
{
  persistentContainer_.clear();
}


/** Element-local data transfer on a single leaf node of the basis tree
 *  Handles computations related to the finite element basis node
 */
template <class Node, class Container, class Basis>
class NodeDataTransfer
{
  using T = typename Container::value_type;
  using LocalView = typename Basis::LocalView;
  using Element = typename Node::Element;

  using LocalBasis = typename Node::FiniteElement::Traits::LocalBasisType;
  using LBRangeType = typename LocalBasis::Traits::RangeType;
  using LocalInterpolation = typename Node::FiniteElement::Traits::LocalBasisType;
  using LIDomainType = typename LocalInterpolation::Traits::DomainType;
  using LIRangeType = typename LocalInterpolation::Traits::RangeType;

public:
  using NodeElementData = std::vector<T>;

public:
  NodeDataTransfer() = default;

  /// To be called once before cacheLocal/restrictLocal are called within the preAdapt step
  void preAdaptInit(LocalView const& lv, Container const& coeff, Node const& node)
  {
    lv_ = &lv;
    node_ = &node;
    fatherNode_ = std::make_unique<Node>(node);
    constCoeff_ = &coeff;
  }

  /// \brief Cache data on the element bound to node_
  /**
   * This functions is used whenever the element does not vanish and thus the
   * data can trivially be transferred to the new element
   **/
  // [[expects: preAdaptInit to be called before]]
  void cacheLocal(NodeElementData& dofs) const
  {
    constCoeff_->gather(*lv_, *node_, dofs);
  }

  /** \brief Evaluate data on the child element bound to node_ and interpolate onto
   *  father entity using the coordinate transformation trafo from father to child.
   *
   * Stores cached data in the NodeElementData argument. After grid adaption the
   * data is copied by \ref copyLocal or \ref prolongLocal to the target element
   * in the new grid.
   *
   * \param father      The father element to interpolate to
   * \param fatherDOFs  Container to store the interpolated DOFs
   * \param trafo       Coordinate transform from local coordinates in father to local
   *                    coordinates in child element
   * \param childDOFs   DOF values from the child element
   * \param init        The father element is visited for the first time
   **/
  // [[expects: preAdaptInit to be called before]]
  template <class Trafo>
  bool restrictLocal(Element const& father, NodeElementData& fatherDOFs, Trafo const& trafo,
                      NodeElementData const& childDOFs, bool init);


  /// To be called once before copyLocal/prolongLocal are called within the adapt step
  void adaptInit(LocalView const& lv, Container& coeff, Node const& node)
  {
    lv_ = &lv;
    node_ = &node;
    fatherNode_ = std::make_unique<Node>(node);
    coeff_ = &coeff;
  }

  /// \brief Copy already existing data to element bound to node_
  // [[expects: adaptInit to be called before]]
  void copyLocal(NodeElementData const& dofs) const
  {
    coeff_->scatter(*lv_, *node_, dofs, Assigner::assign{});
  }

  /** \brief Interpolate data from father onto the child element bound to node_ using
   *  the transformation trafo from child to father
   *
   *  Stores the interpolated data from father to child in the container \ref coeff_.
   *
   * \param father      The father element
   * \param fatherDOFs  DOF values cached on the father element before adapt
   * \param trafo       Coordinate transform from local coordinates in child to local
   *                    coordinates in father element
   * \param init        Father element is visited for the first time
   **/
  // [[expects: adaptInit to be called before]]
  template <class Trafo>
  void prolongLocal(Element const& father, NodeElementData const& fatherDOFs,
                    Trafo const& trafo, bool init);

private:
  LocalView const* lv_ = nullptr;
  Node const* node_ = nullptr;
  std::unique_ptr<Node> fatherNode_;
  Container const* constCoeff_ = nullptr;
  Container* coeff_ = nullptr;
  std::vector<bool> finishedDOFs_;
  NodeElementData fatherDOFsTemp_;
};


template <class N, class C, class B>
  template <class Trafo>
bool NodeDataTransfer<N,C,B>::
restrictLocal(Element const& father, NodeElementData& fatherDOFs, Trafo const& trafo,
              NodeElementData const& childDOFs, bool init)
{
  auto& fatherNode = *fatherNode_;
  std::size_t currentDOF = 0;
  if (init)
  {
    // TODO(FM): This is UB, replace with FE cache for father
    bindTree(fatherNode, father);
  }
  auto const& childNode = *node_;
  auto const& childFE = childNode.finiteElement();
  auto const& fatherFE = fatherNode.finiteElement();

  if (init) {
    finishedDOFs_.assign(fatherFE.size(), false);
    fatherDOFsTemp_.assign(fatherFE.size(), 0);
  }

  auto evalLeaf = [&](LIDomainType const& x) -> LIRangeType {
    if (!finishedDOFs_[currentDOF])
    {
      auto const& insideLocal = trafo(x);
      bool isInside = insideLocal.first;
      if (isInside)
      {
        auto const& local = insideLocal.second;
        thread_local std::vector<LBRangeType> shapeValues;
        childFE.localBasis().evaluateFunction(local, shapeValues);

        assert(childDOFs.size() == shapeValues.size());

        LIRangeType y(0);
        for (std::size_t i = 0; i < shapeValues.size(); ++i)
          y += shapeValues[i] * childDOFs[i];

        fatherDOFsTemp_[currentDOF] = T(y);
        finishedDOFs_[currentDOF++] = true;
        return y;
      }
    }
    return fatherDOFsTemp_[currentDOF++];
  };

  auto evalLeafFct = functionFromCallable<LIRangeType(LIDomainType)>(evalLeaf);
  fatherFE.localInterpolation().interpolate(evalLeafFct, fatherDOFs);

  // Return true if all father DOFs have been evaluated
  return std::accumulate(finishedDOFs_.begin(), finishedDOFs_.end(), true,
                          std::logical_and<bool>());
}


template <class N, class C, class B>
  template <class Trafo>
void NodeDataTransfer<N,C,B>::
prolongLocal(Element const& father, NodeElementData const& fatherDOFs,
             Trafo const& trafo, bool init)
{
  auto& fatherNode = *fatherNode_;
  if (init)
  {
    // TODO(FM): This is UB, replace with FE cache for father
    bindTree(fatherNode, father);
  }
  auto const& childNode = *node_;

  // evaluate father in child coordinate x
  auto evalFather = [&](LIDomainType const& x) -> LIRangeType
  {
    thread_local std::vector<LBRangeType> shapeValues;
    fatherNode.finiteElement().localBasis().evaluateFunction(trafo(x), shapeValues);
    assert(shapeValues.size() == fatherDOFs.size());

    LIRangeType y(0);
    for (std::size_t i = 0; i < shapeValues.size(); ++i)
      y += shapeValues[i] * fatherDOFs[i];

    return y;
  };

  auto const& childFE = childNode.finiteElement();
  thread_local std::vector<T> childDOFs;
  auto evalFatherFct = functionFromCallable<LIRangeType(LIDomainType)>(evalFather);
  childFE.localInterpolation().interpolate(evalFatherFct, childDOFs);

  coeff_->scatter(*lv_, childNode, childDOFs, Assigner::assign{});
}

} // end namespace AMDiS