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

#include <type_traits>
#include <variant>
#include <vector>

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

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

namespace Dune
{
  template <int codim, PartitionIteratorType pitype, class HostGrid>
  class MultiIterator : public DummyIterator
  {
    using DummyIterator::DummyIterator;
  };

  template <PartitionIteratorType pitype, class HostGrid>
  class MultiIterator<0,pitype,HostGrid>
      : public MultiMeshIteratorBase<pitype,HostGrid,MultiIterator<0,pitype,HostGrid>>
  {
    using Self = MultiIterator;
    using Super = MultiMeshIteratorBase<pitype,HostGrid,Self>;
    using HostEntity = typename HostGrid::template Codim<0>::Entity;
    using EntityTest = typename Super::EntityTest;

    template <class GV>
    using IsGridView = Std::disjunction<
      std::is_same<std::decay_t<GV>,typename HostGrid::LeafGridView>,
      std::is_same<std::decay_t<GV>,typename HostGrid::LevelGridView> >;

  public:
    template <class... GridViews,
      std::enable_if_t<Std::conjunction<IsGridView<GridViews>...>::value, int> = 0>
    MultiIterator (tag::begin_iterator, GridViews const&... gridViews)
      : Super{tag::begin_iterator{}, gridViews...}
      , incrementAllowed_(sizeof...(GridViews), true)
      , multiEntity_(sizeof...(GridViews))
    {
      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); });
      });

      this->initialIncrement();
    }

#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>
    MultiIterator (tag::begin_iterator, const std::vector<std::variant<GV...>>& gridViews)
      : Super{tag::begin_iterator{}, gridViews}
      , incrementAllowed_(gridViews.size(), true)
      , multiEntity_(gridViews.size())
    {
      contains_.reserve(gridViews.size());

      for (auto const& gvs : gridViews) {
        contains_.emplace_back(std::visit([](auto const& gv) {
          return EntityTest{[gv](const HostEntity& entity) { return gv.contains(entity); }}; }, gvs));
      }

      this->initialIncrement();
    }
#endif

    // End iterator
    template <class... Args>
    MultiIterator (tag::end_iterator, Args&&... args)
      : Super{tag::end_iterator{}, std::forward<Args>(args)...}
    {}

    void initialIncrement_impl ()
    {
      for (std::size_t i = 0; i < this->size(); ++i)
        this->initialIncrement(i);
    }

    void increment_impl ()
    {
      for (std::size_t i = 0; i < this->size(); ++i)
        incrementAllowed_[i] = this->incrementAllowed(i);

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

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

    bool levelReached_impl (std::size_t i, const HostEntity& entity) const
    {
      return contains_[i](entity) || entity.isLeaf();
    }

  private:
    std::vector<bool> incrementAllowed_;
    std::vector<EntityTest> contains_;
    mutable MultiEntity<HostGrid> multiEntity_;
  };


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

}  // end namespace Dune

#endif // DUNE_MULTI_ITERATOR_HH