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

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

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

namespace Dune
{
  /** \brief Iterator over all entities of a given codimension and level of a grid.
   *  \ingroup MultiMesh
   */
  template <class HostGrid>
  class MultiMeshHierarchicIterator
      : public ForwardIteratorFacade<MultiMeshHierarchicIterator<HostGrid>,
            typename HostGrid::template Codim<0>::Entity,
            typename HostGrid::template Codim<0>::Entity>
  {
  private:
    template <PartitionIteratorType pitype, class HG, class D>
    friend class MultiMeshIteratorBase;

    using HostEntity = typename HostGrid::template Codim<0>::Entity;
    using EntityTest = std::function<bool(HostEntity)>;

    struct EntityStackEntry
    {
      template <class Entity>
      explicit EntityStackEntry (Entity&& entity)
        : it(entity.hbegin(entity.level()+1))
        , end(entity.hend(entity.level()+1))
      {}

      typename HostEntity::HierarchicIterator it;
      typename HostEntity::HierarchicIterator end;

      friend bool operator== (const EntityStackEntry& lhs, const EntityStackEntry& rhs)
      {
        return lhs.it == rhs.it;
      }
    };

    class EntityStack
        : public std::stack<EntityStackEntry, std::vector<EntityStackEntry>>
    {
      using Super = std::stack<EntityStackEntry, std::vector<EntityStackEntry>>;

    public:
      explicit EntityStack (int maxLevel)
      {
        c.reserve(maxLevel);
      }

      // return true if all levels >= l are finished, i.e. hierarchic iterators it+1 == end
      bool finished (std::size_t l = 0) const
      {
        bool f = true;
        for (auto tmp = c[l].it; f && l < c.size(); ++l) {
          tmp = c[l].it;
          f = f && (++tmp) == c[l].end;
        }
        return f;
      }

      friend bool operator== (const EntityStack& lhs, const EntityStack& rhs)
      {
        return lhs.size() == rhs.size() && lhs.c == rhs.c;
      }

    protected:
      using Super::c;
    };

  public:
    /// Constructor. Stores a pointer to the entity
    template <class Entity, class GridView>
    MultiMeshHierarchicIterator (const Entity& entity, const GridView& gridView)
      : entity_{&entity}
      , contains_{[gv=gridView](const HostEntity& e) { return gv.contains(e); }}
      , maxLevel_{gridView.grid().maxLevel()}
      , entityStack_{EntityStack{gridView.grid().maxLevel()}}
    {
      initialIncrement();
    }

    /// Constructor which create the end iterator
    /**
     * \param  endDummy   Here only to distinguish it from the other constructor
     */
    MultiMeshHierarchicIterator (bool endDummy)
      : entityStack_{0}
    {}

    /// prefix increment
    void increment ()
    {
      // 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, iteration is finished
      if (entityStack_.empty())
        return;

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

      assert(contains_(child) && "No valid child element found in gridView");
    }

    /// dereferencing
    decltype(auto) dereference () const
    {
      if (entityStack_.empty()) {
        return *entity_;
      } else {
        assert(entityStack_.top().it != entityStack_.top().end);
        return *entityStack_.top().it;
      }
    }

    /// equality
    bool equals (const MultiMeshHierarchicIterator& that) const
    {
      return entityStack_ == that.entityStack_;
    }

  protected:

    // got to first leaf entity on gridView
    void initialIncrement ()
    {
      // 1. go down in tree until leaf entity
      auto child = dereference();
      for (; !(contains_(child) || child.isLeaf()); child = dereference()) {
        assert(child.isRegular() && "No irregular elements allowed in multi-mesh traversal");
        entityStack_.emplace(child);
        assert( entityStack_.size() <= maxLevel_ );
      }

      assert(contains_(child) && "No valid child element found in gridView");
    }

  private:
    const HostEntity* entity_ = nullptr;

    EntityTest contains_;
    int maxLevel_;
    EntityStack entityStack_;
  };

  template <class Entity, class GridView>
  inline auto childs(const Entity& entity, const GridView& gridView)
  {
    using Iterator = MultiMeshHierarchicIterator<typename GridView::Grid>;
    return IteratorRange<Iterator>{ Iterator{entity, gridView}, Iterator{true} };
  }

}  // end namespace Dune

#endif // DUNE_MULTIMESH_HIERARCHIC_ITERATOR_HH