// -*- tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 2 -*-
// vi: set et ts=4 sw=2 sts=2:
#ifndef DUNE_MULTIMESH_ITERATOR_HH
#define DUNE_MULTIMESH_ITERATOR_HH

#include <numeric>
#include <stack>
#include <variant>

#include <dune/common/iteratorfacades.hh>
#include <dune/common/std/type_traits.hh>
#include <dune/grid/common/exceptions.hh>
#include <dune/grid/common/gridenums.hh>

#include "mmentity.hh"
#include "mmhierarchiciterator.hh"

namespace Dune
{
  /** \brief Iterator over all entities of a given codimension and level of a grid.
   *  \ingroup MultiMesh
   */
  template <int codim, PartitionIteratorType pitype, class HostGrid>
  class MultiMeshIterator
  {
  public:
    template <class... Args>
    MultiMeshIterator(Args&&... args)
    {
      DUNE_THROW(NotImplemented, "Not implemented for entity of codim != 0!");
    }

    int operator*() const { return 0; }
    MultiMeshIterator& operator++() { return *this; }
    MultiMeshIterator operator++(int) { return *this; }
    bool operator!=(MultiMeshIterator const&) const { return false; }
  };


  // Implemented for codim 0 entities only
  template <PartitionIteratorType pitype, class HostGrid>
  class MultiMeshIterator<0, pitype, HostGrid>
      : public ForwardIteratorFacade<MultiMeshIterator<0,pitype,HostGrid>,
            MultiEntity<HostGrid>,
            MultiEntity<HostGrid> const&>
  {
  private:
    // LevelIterator to the equivalent entity in the host grid
    using HostGridLevelIterator =
      typename HostGrid::template Codim<0>::template Partition<pitype>::LevelIterator;

    using HostEntity = typename HostGrid::template Codim<0>::Entity;

    using EntityTest  = typename MultiMeshHierarchicIterator<HostGrid>::EntityTest;
    using EntityStack = typename MultiMeshHierarchicIterator<HostGrid>::EntityStack;

  public:
    /// Constructor. Stores a pointer to the grid
    template <class GridImp>
    MultiMeshIterator (const GridImp* multiMesh, int level)
      : incrementAllowed_(multiMesh->size(), true)
      , contains_(multiMesh->size(), EntityTest{[level](const HostEntity& entity) { return entity.level() == level; }})
      , maxLevel_(multiMesh->size(), level)
      , multiEntity_(multiMesh->size())
    {
      for (auto const& grid : multiMesh->grids_) {
        macroIterators_.push_back(grid->levelGridView(0).template begin<0,pitype>());
        macroEndIterators_.push_back(grid->levelGridView(0).template end<0,pitype>());
      }

      // go to first leaf entity on all grids
      entityStacks_.reserve(multiMesh->size());
      for (std::size_t i = 0; i < multiMesh->size(); ++i) {
        maxLevel_[i] = (*multiMesh)[i].maxLevel();
        entityStacks_.emplace_back((*multiMesh)[i].maxLevel());
        initialIncrement(i);
      }
    }

    /// Constructor which create the end iterator
    /**
     * \param  multiMesh  Pointer to grid instance
     * \param  endDummy   Here only to distinguish it from the other constructor
     */
    template <class GridImp>
    MultiMeshIterator (const GridImp* multiMesh, int level, bool endDummy)
    {
      for (auto const& grid : multiMesh->grids_)
        macroIterators_.push_back(grid->levelGridView(0).template end<0,pitype>());
    }

    /// Constructor which create the leaf-iterator
    template <class GridImp>
    MultiMeshIterator (const GridImp* multiMesh)
      : MultiMeshIterator{multiMesh, -1}
    {}

    /// Constructor which create the end leaf-iterator
    template <class GridImp>
    MultiMeshIterator (const GridImp* multiMesh, bool endDummy)
      : MultiMeshIterator{multiMesh, -1, endDummy}
    {}

    /// Construct an iterator from n gridViews of possibly different type
    template <class... GridViews,
      std::enable_if_t<not Std::disjunction<std::is_same<GridViews,bool>...>::value, int> = 0>
    MultiMeshIterator (GridViews const&... gridViews)
      : incrementAllowed_(sizeof...(GridViews), true)
      , maxLevel_{gridViews.grid().maxLevel()...}
      , multiEntity_(sizeof...(GridViews))
      , macroIterators_{gridViews.grid().levelGridView(0).template begin<0,pitype>()...}
      , macroEndIterators_{gridViews.grid().levelGridView(0).template end<0,pitype>()...}
      , entityStacks_{EntityStack{gridViews.grid().maxLevel()}...}
    {
      contains_.reserve(sizeof...(GridViews));
      Hybrid::forEach(std::forward_as_tuple(gridViews...), [this](auto const& gv) {
        contains_.emplace_back([gv](const HostEntity& entity) { return gv.contains(entity); });
      });

      for (std::size_t i = 0; i < sizeof...(GridViews); ++i)
        initialIncrement(i);
    }

    template <class... GridViews>
    MultiMeshIterator (bool endDummy, GridViews const&... gridViews)
      : macroIterators_{gridViews.grid().levelGridView(0).template end<0,pitype>()...}
    {}

#if DUNE_HAVE_CXX_VARIANT
    /// Construct an iterator from a vector of GridViews, where each GridView can be
    /// either a leaf- or a level-gridView, thus we need to use a variant.
    template <class... GV>
    MultiMeshIterator (const std::vector<std::variant<GV...>>& gridViews)
      : incrementAllowed_(gridViews.size(), true)
      , multiEntity_(gridViews.size())
    {
      contains_.reserve(gridViews.size());
      entityStacks_.reserve(gridViews.size());
      maxLevel_.reserve(gridViews.size());

      for (auto const& gvs : gridViews) {
        macroIterators_.push_back(std::visit([](auto const& gv) {
          return gv.grid().levelGridView(0).template begin<0,pitype>(); }, gvs));
        macroEndIterators_.push_back(std::visit([](auto const& gv) {
          return gv.grid().levelGridView(0).template end<0,pitype>(); }, gvs));
        contains_.emplace_back(std::visit([](auto const& gv) {
          return EntityTest{[gv](const HostEntity& entity) { return gv.contains(entity); }}; }, gvs));
        maxLevel_.push_back(std::visit([](auto const& gv) { return gv.grid().maxLevel(); }, gvs));
        entityStacks_.emplace_back(std::visit([](auto const& gv) { return gv.grid().maxLevel(); }, gvs));
      }

      // go to first leaf entity on all grids
      for (std::size_t i = 0; i < entityStacks_.size(); ++i)
        initialIncrement(i);
    }

    template <class... GV>
    MultiMeshIterator (const std::vector<std::variant<GV...>>& gridViews, bool endDummy)
    {
      for (auto const& gvs : gridViews) {
        macroIterators_.push_back(std::visit([](auto const& gv) {
          return gv.grid().levelGridView(0).template end<0,pitype>(); }, gvs));
      }
    }
#endif

    /// prefix increment
    void increment ()
    {
      for (std::size_t i = 0; i < entityStacks_.size(); ++i)
        incrementAllowed_[i] = incrementAllowed(i);

      for (std::size_t i = 0; i < entityStacks_.size(); ++i) {
        if (incrementAllowed_[i])
          increment(i);
      }
    }

    /// dereferencing
    MultiEntity<HostGrid> const& dereference () const
    {
      // update entries in multiEntity that have changed
      for (std::size_t i = 0; i < entityStacks_.size(); ++i) {
        if (incrementAllowed_[i])
          multiEntity_[i] = dereference(i);
      }
      return multiEntity_;
    }

    /// equality
    bool equals (const MultiMeshIterator& that) const
    {
      return macroIterators_ == that.macroIterators_;
    }

  protected:
    // got to next entity in grid i
    void increment (std::size_t i)
    {
      auto& entityStack = entityStacks_[i];
      auto& macroIt = macroIterators_[i];
      auto const& macroEnd = macroEndIterators_[i];

      // 1. go up in tree or to next entity on current level until we can go down again
      while (!entityStack.empty()) {
        auto& top = entityStack.top();
        ++top.it;
        if (top.it == top.end) {
          entityStack.pop();
        } else {
          break;
        }
      }

      // 2. if entityStack is empty, go to next macroElement
      if (entityStack.empty()) {
        ++macroIt;
        if (macroIt == macroEnd)
          return;
      }

      // 3. go down in tree until leaf entity
      auto child = dereference(i);
      for (; !(contains_[i](child) || child.isLeaf()); child = dereference(i)) {
        assert(child.isRegular() && "No irregular elements allowed in multi-mesh traversal");
        entityStack.emplace(child);
        assert(entityStack.size() <= maxLevel_[i]);
      }

      // assert(contains_[i](child) && "No valid child element found in gridView");
    }

    /// Return true, if all stacks with size > stack[i].size are finished
    bool incrementAllowed (std::size_t i) const
    {
      std::size_t size = entityStacks_[i].size();
      return std::accumulate(entityStacks_.begin(), entityStacks_.end(), true,
        [i,size](bool allowed, auto const& entityStack) {
          return allowed && (entityStack.size() <= size || entityStack.finished(size));
        });
    }

    // got to first leaf entity on grid i
    void initialIncrement (std::size_t i)
    {
      auto& entityStack = entityStacks_[i];
      auto& macroIt = macroIterators_[i];
      auto const& macroEnd = macroEndIterators_[i];

      assert(entityStack.empty());
      if (macroIt == macroEnd)
        return;

      // 1. go down in tree until leaf entity
      auto child = dereference(i);
      for (; !(contains_[i](child) || child.isLeaf()); child = dereference(i)) {
        assert(child.isRegular() && "No irregular elements allowed in multi-mesh traversal");
        entityStack.emplace(child);
        assert(entityStack.size() <= maxLevel_[i]);
      }

      // assert(contains_[i](child) && "No valid child element found in gridView");
    }

    HostEntity dereference (std::size_t i) const
    {
      if (entityStacks_[i].empty()) {
        assert(macroIterators_[i] != macroEndIterators_[i]);
        return *macroIterators_[i];
      } else {
        assert(entityStacks_[i].top().it != entityStacks_[i].top().end);
        return *entityStacks_[i].top().it;
      }
    }

  private:
    std::vector<bool> incrementAllowed_;
    std::vector<EntityTest> contains_;
    std::vector<int> maxLevel_;

    mutable MultiEntity<HostGrid> multiEntity_;

    std::vector<HostGridLevelIterator> macroIterators_;
    std::vector<HostGridLevelIterator> macroEndIterators_;
    std::vector<EntityStack> entityStacks_;
  };

  template <class... GridViews>
  inline auto multi_elements(GridViews const&... gridViews)
  {
    using GridView0 = std::tuple_element_t<0,std::tuple<GridViews...>>;
    using Iterator = MultiMeshIterator<0,All_Partition,typename GridView0::Grid>;
    using Range = IteratorRange<Iterator>;
    return Range{ Iterator{gridViews...},
                  Iterator{true, gridViews...} };
  }

}  // end namespace Dune

#endif // DUNE_MULTIMESH_ITERATOR_HH