#pragma once

#include <array>
#include <functional>
#include <type_traits>
#include <utility>

#include <dune/common/indices.hh>
#include <dune/common/tuplevector.hh>

#include <dune/typetree/treepath.hh>

#include <amdis/common/Apply.hpp>
#include <amdis/common/TypeTraits.hpp>
#include <amdis/typetree/TreePath.hpp>

// NOTE: backport of dune/typetree/treecontainer.hh

namespace AMDiS
{
  /// \brief Vector data-structure with tree-path index access and hierarchic structure
  /// given by the `Container` template type
  /**
   * This Vector container is parametrized with the actual container type that is stored
   * internally. Access to the elements of the container is possible by using a tree-path
   * index.
   *
   * The internal container is constructed by the \ref ContainerFactory, storing for each
   * tree node a corresponding array or tuple plus a value. The data-structure to hold
   * both, the value and the container is defined in \ref ValueAndContainer.
   **/
  template <class Container>
  class TreeContainerStorage
  {
    using Self = TreeContainerStorage;

    template <class C>
    static constexpr decltype(auto)
    accessByTreePath(C&& container, Dune::TypeTree::HybridTreePath<> const&)
    {
      return container.value();
    }

    template <class C, class... T>
    static constexpr decltype(auto)
    accessByTreePath(C&& container, Dune::TypeTree::HybridTreePath<T...> const& path)
    {
      auto head = Dune::TypeTree::treePathEntry(path,Dune::Indices::_0);
      return accessByTreePath(container[head], pop_front(path));
    }

  public:
    /// \brief Default-construct the tree-container
    TreeContainerStorage()
      : container_()
    {}

    TreeContainerStorage(Container const& container)
      : container_(container)
    {}

    /// \brief Construct the tree-container from a given container storage
    TreeContainerStorage(Container&& container)
      : container_(std::move(container))
    {}

    /// \brief Access a (const) element of the container by treepath
    template <class... T>
    decltype(auto) operator[](Dune::TypeTree::HybridTreePath<T...> const& path) const
    {
      return accessByTreePath(container_, path);
    }

    /// \brief Access a (mutable) element of the container by treepath
    template <class... T>
    decltype(auto) operator[](Dune::TypeTree::HybridTreePath<T...> const& path)
    {
      return accessByTreePath(container_, path);
    }

    /// \brief Obtain the container (const)
    Container const& data() const
    {
      return container_;
    }

    /// \brief Obtain the container (mutable)
    Container& data()
    {
      return container_;
    }

    /// \brief Compare two containers for equality
    bool operator==(TreeContainerStorage const& other) const
    {
      return container_ == other.container_;
    }

    /// \brief Compare two containers for inequality
    bool operator!=(TreeContainerStorage const& other) const
    {
      return container_ != other.container_;
    }

  private:
    Container container_;
  };


  namespace Impl
  {
    template <class Value, class Container>
    class ValueAndContainer
    {
    public:
      template <class V, class C>
      ValueAndContainer(V&& value, C&& container)
        : value_(FWD(value))
        , container_(FWD(container))
      {}

      ValueAndContainer()
        : value_()
        , container_()
      {}

      template <class I>
      auto& operator[](I const& i) { return container_[i]; }

      template <class I>
      auto const& operator[](I const& i) const { return container_[i]; }

      Value& value() { return value_; }
      Value const& value() const { return value_; }

      Container& container() { return container_; }
      Container const& container() const { return container_; }

      bool operator==(ValueAndContainer const& other) const
      {
        return value_ == other.value_ && container_ == other.container_;
      }

    private:
      Value value_;
      Container container_;
    };

    template <class Value>
    class ValueAndContainer<Value,void>
    {
    public:
      template <class V, Dune::disableCopyMove<ValueAndContainer,V> = 0>
      ValueAndContainer(V&& value)
        : value_(FWD(value))
      {}

      ValueAndContainer()
        : value_()
      {}

      Value& value() { return value_; }
      Value const& value() const { return value_; }

      bool operator==(ValueAndContainer const& other) const
      {
        return value_ == other.value_;
      }

    private:
      Value value_;
    };

    struct Ignore{};

    template <class Container>
    class ValueAndContainer<Ignore,Container>
    {
    public:
      template <class C>
      ValueAndContainer(Ignore, C&& container)
        : container_(FWD(container))
      {}

      ValueAndContainer()
        : container_()
      {}

      template <class I>
      auto& operator[](I const& i) { return container_[i]; }

      template <class I>
      auto const& operator[](I const& i) const { return container_[i]; }

      Ignore value() { return {}; }
      Ignore value() const { return {}; }

      Container& container() { return container_; }
      Container const& container() const { return container_; }

      bool operator==(ValueAndContainer const& other) const
      {
        return container_ == other.container_;
      }

    private:
      Container container_;
    };


    template <class Value, class Container>
    ValueAndContainer(Value const&, Container const&)
      -> ValueAndContainer<Value,Container>;

    template <class Value>
    ValueAndContainer(Value const&)
      -> ValueAndContainer<Value,void>;


    /// \brief A factory class creating a hybrid container compatible with a type tree
    /**
     * This class allows to create a nested hybrid container having the same structure
     * as a given type tree. Power nodes are represented as std::array's while composite
     * nodes are represented as Dune::TupleVector's. The stored values for the leaf nodes
     * are creating using a given predicate. Once created, the factory provides an
     * operator() creating the container for the tree given as argument.
     *
     * \tparam NodeToValue Type of a predicate that determines the stored values at the
     *                     leafs
     **/
    template <class NodeToValue, bool leafOnly = false>
    class ContainerFactory
    {
    public:
      /// \brief Create ContainerFactory
      /**
        * The given predicate will be stored by value.
        *
        * \param A predicate used to generate the stored values for the leaves
        */
      ContainerFactory(NodeToValue nodeToValue)
        : nodeToValue_(std::move(nodeToValue))
      {}

      /// \brief Return a container for storing the node content
      template <class Node>
      auto operator()(Node const& node) const
      {
        if constexpr (Node::isLeaf)
          return ValueAndContainer{value(node)};
        else
        if constexpr (Node::isPower) {
          using TransformedChild = decltype((*this)(node.child(0)));
          return ValueAndContainer{
            value(node),
            std::array<TransformedChild, Node::degree()>()};
        }
        else
        if constexpr (Node::isComposite) {
          return ValueAndContainer{
            value(node),
            Ranges::applyIndices<Node::degree()>(
            [&](auto... ii) { return Dune::makeTupleVector((*this)(node.child(ii))...); })};
        }
        else {
          static_assert(Node::isLeaf || Node::isPower || Node::isComposite,
            "Node must be one of leaf,power,composite.");
          return;
        }
      }

      template <class Node>
      decltype(auto) value(Node const& node) const
      {
        if constexpr(!Node::isLeaf && leafOnly)
          return Ignore{};
        else
          return nodeToValue_(node);
      }

    private:
      NodeToValue nodeToValue_;
    };

  } // end namespace Impl


  /// \brief Create container having the same structure as the given tree
  /**
   * This class allows to create a nested hybrid container having the same structure
   * as a given type tree. Power nodes are represented as std::array's while composite
   * nodes are represented as Dune::TupleVector's. The stored values for the leaf nodes
   * are creating using a given predicate. For convenience the created container is
   * not returned directly. Instead, the returned object stores the container and
   * provides operator[] access using a HybridTreePath.
   *
   * \param tree        The tree which should be mapper to a container
   * \param nodeToValue A predicate used to generate the stored values for the nodes
   *
   * \returns A container matching the tree structure
   */
  template <bool leafOnly = false, class Tree, class NodeToValue>
  auto makeTreeContainer(Tree const& tree, NodeToValue&& nodeToValue)
  {
    auto f = std::ref(nodeToValue);
    auto factory = Impl::ContainerFactory<decltype(f),leafOnly>(f);
    return TreeContainerStorage{factory(tree)};
  }


  /// \brief Create container having the same structure as the given tree
  /**
   * This class allows to create a nested hybrid container having the same structure
   * as a given type tree. Power nodes are represented as std::array's while composite
   * nodes are represented as Dune::TupleVector's. The stored values for the leaf nodes
   * are of the given type Value. For convenience the created container is
   * not returned directly. Instead, the returned object stores the container and
   * provides operator[] access using a HybridTreePath.
   *
   * \tparam Value      Type of the values to be stored for the leafs. Should be default
   *                    constructible.
   * \param nodeToValue A predicate used to generate the stored values for the nodes
   *
   * \returns           A container matching the tree structure
   */
  template <class Value, bool leafOnly = false, class Tree>
  auto makeTreeContainer(Tree const& tree)
  {
    return makeTreeContainer<leafOnly>(tree, [](auto&&) { return Value{}; });
  }

  template <template<class> class NodeData, bool leafOnly = false, class Tree>
  auto makeTreeContainer(Tree const& tree)
  {
    return makeTreeContainer<leafOnly>(tree, [](auto&& node) { return NodeData<TYPEOF(node)>{}; });
  }

  /// Alias to container type generated by makeTreeContainer for given tree type and
  /// uniform value type
  template <class Value, class Tree, bool leafOnly = false>
  using UniformTreeContainer
    = TYPEOF(makeTreeContainer<Value,leafOnly>(std::declval<const Tree&>()));

  /// Alias to matrix-container type generated by makeTreeContainer for given tree types
  /// and uniform value type
  template <class Value, class RowTree, class ColTree = RowTree, bool leafOnly = false>
  using UniformTreeMatrix
    = UniformTreeContainer<UniformTreeContainer<Value,ColTree,leafOnly>,RowTree,leafOnly>;

  /// Alias to container type generated by makeTreeContainer for give tree type and when
  /// using NodeToValue to create values
  template <template <class Node> class NodeData, class Tree, bool leafOnly = false>
  using TreeContainer
    = TYPEOF(makeTreeContainer<NodeData,leafOnly>(std::declval<const Tree&>()));


  namespace Impl
  {
    template <template <class,class> class NodeData, class Tree, bool leafOnly>
    struct RowNodeData
    {
      template <class RowNode>
      struct ColNodeData
      {
        template <class ColNode>
        using type = NodeData<RowNode, ColNode>;
      };

      template <class RowNode>
      using type = TreeContainer<ColNodeData<RowNode>::template type, Tree, leafOnly>;
    };

  } // end namespace Impl

  /// Alias to matrix-container type generated by makeTreeContainer for give tree type
  /// and when using NodeToValue to create values
  template <template <class,class> class NodeData, class RowTree, class ColTree = RowTree, bool leafOnly = false>
  using TreeMatrix
    = TreeContainer<Impl::RowNodeData<NodeData,ColTree,leafOnly>::template type,RowTree,leafOnly>;

} //namespace AMDiS