#pragma once

#include <cassert>
#include <type_traits>

#include <amdis/GridFunctions.hpp>
#include <amdis/common/Utility.hpp>
#include <amdis/utility/FiniteElementType.hpp>

namespace AMDiS
{
  /**
   * \addtogroup operators
   * @{
   **/

  /// \brief The main implementation of an operator to be used in a \ref LocalAssembler.
  /**
   * An Operator that takes a GridFunction as coefficient.
   * Provides quadrature rules and handles the evaluation of the GridFunction at
   * local coordinates.
   *
   * The class is specialized, by deriving from it, in \ref GridFunctionOperator.
   *
   * \tparam GridFunction The GridFunction, a LocalFunction is created from, and
   *                      that is evaluated at quadrature points.
   * \tparam QuadratureCreator A functor that provides a \ref Dune::QuadratureRule.
   *
   * **Requirements:**
   * - `GridFunction` models the \ref Concepts::GridFunction
   * - `QuadratureCreator` models \ref Concepts::Callable<Dune::GeometryType, LocalFunction, F>
   *   where `F` is a functor of the signature `int(int)` that calculates the
   *   degree of the (bi)linear-form. The argument passed to `F` is the polynomial
   *   order of the GridFunction.
   **/
  template <class GridFunction, class QuadratureCreator>
  class GridFunctionOperatorBase
  {
    using LocalFunction = decltype(localFunction(std::declval<GridFunction>()));
    using Element = typename GridFunction::EntitySet::Element;
    using Geometry = typename Element::Geometry;

    using LocalCoordinate = typename GridFunction::EntitySet::LocalCoordinate;

  public:
    /// \brief Constructor. Stores a copy of `expr`.
    /**
     * An expression operator takes an expression, following the interface of
     * \ref ExpressionBase, and stores a copy. Additionally, it gets the
     * differentiation order, to calculate the quadrature degree in \ref getDegree.
     **/
    GridFunctionOperatorBase(GridFunction const& gridFct, QuadratureCreator creator, int order)
      : gridFct_(gridFct)
      , localFct_(localFunction(gridFct_))
      , quadCreator_(creator)
      , order_(order)
    {}

    // Copy constructor
    GridFunctionOperatorBase(GridFunctionOperatorBase const& that)
      : gridFct_(that.gridFct_)
      , localFct_(localFunction(gridFct_)) // recreate localFct_ on new gridFct_
      , quadCreator_(that.quadCreator_)
      , order_(that.order_)
    {}

    // Move constructor
    GridFunctionOperatorBase(GridFunctionOperatorBase&& that)
      : gridFct_(std::move(that.gridFct_))
      , localFct_(localFunction(gridFct_)) // recreate localFct_ on new gridFct_
      , quadCreator_(std::move(that.quadCreator_))
      , order_(std::move(that.order_))
    {}

    /// \brief Binds operator to `element` and `geometry`.
    /**
     * Binding an operator to the currently visited element in grid traversal.
     * Since all operators need geometry information, the `Element::Geometry`
     * object `geometry` is created once, during grid traversal, and passed in.
     *
     * By default, it binds the \ref localFct_ to the `element`.
     **/
    void bind(Element const& element, Geometry const& geometry)
    {
      if (!bound_) {
        localFct_.bind(element);
        isSimplex_ = geometry.type().isSimplex();
        isAffine_ = geometry.affine();
        bound_ = true;
      }
    }

    /// Unbinds operator from element.
    void unbind()
    {
      bound_ = false;
      localFct_.unbind();
    }

    /// Return expression value at LocalCoordinates
    auto coefficient(LocalCoordinate const& local) const
    {
      assert( bound_ );
      return localFct_(local);
    }


    /// Create a quadrature rule using the \ref QuadratureCreator by calculating the
    /// quadrature order necessary to integrate the (bi)linear-form accurately.
    template <class... Nodes>
    decltype(auto) getQuadratureRule(Dune::GeometryType type, Nodes const&... nodes) const
    {
      auto degreeFunctor = [&,this](int order) -> int { return this->getDegree(order, nodes...); };
      return quadCreator_(type, localFct_, degreeFunctor);
    }

  protected:
    /// \brief Return the quadrature degree for a vector operator.
    /**
     * The quadrature degree that is necessary, to integrate the expression
     * multiplied with (possibly derivatives of) shape functions. Thus it needs
     * the order of derivatives, this operator implements.
     **/
    template <class Node>
    int getDegree(int coeffDegree, Node const& node) const
    {
      assert( bound_ );

      int psiDegree = polynomialDegree(node);

      int degree = psiDegree + coeffDegree;
      if (isSimplex_)
        degree -= order_;
      if (isAffine_)
        degree += 1; // oder += (order+1)

      return degree;
    }


    /// \brief Return the quadrature degree for a matrix operator.
    /**
     * The quadrature degree that is necessary, to integrate the expression
     * multiplied with (possibly derivatives of) shape functions. Thus it needs
     * the order of derivatives, this operator implements.
     **/
    template <class RowNode, class ColNode>
    int getDegree(int coeffDegree, RowNode const& rowNode, ColNode const& colNode) const
    {
      assert( bound_ );

      int psiDegree = polynomialDegree(rowNode);
      int phiDegree = polynomialDegree(colNode);

      int degree = psiDegree + phiDegree + coeffDegree;
      if (isSimplex_)
        degree -= order_;
      if (isAffine_)
        degree += 1; // oder += (order+1)

      return degree;
    }

  private:
    GridFunction gridFct_;
    LocalFunction localFct_;

    QuadratureCreator quadCreator_; //< a creator to provide a quadrature rule
    int order_;                     //< the derivative order of this operator

    bool isSimplex_ = false;    //< the bound element is a simplex
    bool isAffine_ = false;     //< the bound geometry is affine
    bool bound_ = false;        //< an element is bound to the operator
  };


  /// A default implementation of an GridFunctionOperator if no specialization is available.
  /**
   * An operator must implement either \ref calculateElementVector, or
   * \ref calculateElementMatrix, if it is a vector or matrix operator,
   * respectively.
   **/
  template <class Tag, class GridFct, class QuadratureCreator>
  class GridFunctionOperator
      : public GridFunctionOperatorBase<GridFct, QuadratureCreator>
  {
  public:
    GridFunctionOperator(Tag, GridFct const& gridFct, QuadratureCreator const& quadCreator)
      : GridFunctionOperatorBase<GridFct, QuadratureCreator>(gridFct, 0, quadCreator)
    {}

    /// \brief Assemble a local element vector on the element that is bound.
    /**
     * \param contextGeometry Container for context related data
     * \param quad A quadrature formula
     * \param elementVector The output element vector
     * \param node The node of the test-basis
     **/
    template <class Context, class QuadratureRule,
              class ElementVector, class Node>
    void calculateElementVector(Context const& contextGeometry,
                                QuadratureRule const& quad,
                                ElementVector& elementVector,
                                Node const& node)
    {
      error_exit("Needs to be implemented!");
    }

    /// \brief Assemble a local element matrix on the element that is bound.
    /**
     * \param contextGeometry Container for context related data
     * \param quad A quadrature formula
     * \param elementMatrix The output element matrix
     * \param rowNode The node of the test-basis
     * \param colNode The node of the trial-basis
     * \param flag1 finiteelement is the same for test and trial
     * \param flag2 nodes are the same in the basis-tree
     **/
    template <class Context, class QuadratureRule,
              class ElementMatrix, class RowNode, class ColNode, bool sameFE, bool sameNode>
    void calculateElementMatrix(Context const& contextGeometry,
                                QuadratureRule const& quad,
                                ElementMatrix& elementMatrix,
                                RowNode const& rowNode,
                                ColNode const& colNode,
                                std::integral_constant<bool, sameFE> flag1,
                                std::integral_constant<bool, sameNode> flag2)
    {
      error_exit("Needs to be implemented!");
    }
  };


#ifndef DOXYGEN
  namespace Concepts
  {
    namespace Definition
    {
      struct HasGridFunctionOrder
      {
        template <class F, class GridView>
        auto requires_(F&& f, GridView const& gridView)
          -> decltype( order(localFunction(makeGridFunction(f, gridView))) );
      };
    }

    template <class F, class GridView = Dune::YaspGrid<2>::LeafGridView>
    constexpr bool HasGridFunctionOrder = models<Definition::HasGridFunctionOrder(F, GridView)>;

  } // end namespace Concepts


  namespace tag
  {
    struct deduce {};
    struct order {};
    struct rule {};

  } // end namespace tag


  template <class Tag, class Expr, class QuadTag, class Rule = void>
  struct ExpressionPreOperator;

  template <class Tag, class Expr>
  struct ExpressionPreOperator<Tag, Expr, tag::deduce>
  {
    Tag tag;
    Expr expr;
  };

  template <class Tag, class Expr>
  struct ExpressionPreOperator<Tag, Expr, tag::order>
  {
    Tag tag;
    Expr expr;
    int order;
    Dune::QuadratureType::Enum qt;
  };

  template <class Tag, class Expr, class Rule>
  struct ExpressionPreOperator<Tag, Expr, tag::rule, Rule>
  {
    Tag tag;
    Expr expr;
    Rule const& rule;
  };
#endif


  /// Store tag and expression to create a \ref GridFunctionOperator
  template <class Tag, class Expr>
  auto makeOperator(Tag tag, Expr const& expr)
  {
    using RawExpr = Underlying_t<Expr>;
    static_assert(Concepts::HasGridFunctionOrder<RawExpr>,
      "Polynomial degree of expression can not be deduced. You need to provide an explicit value for polynomial order or a quadrature rule in 'makeOperator()'.");

    return ExpressionPreOperator<Tag, Expr, tag::deduce>{tag, expr};
  }

  /// Store tag and expression and polynomial order of expression to create a \ref GridFunctionOperator
  template <class Tag, class Expr>
  auto makeOperator(Tag tag, Expr const& expr, int order,
                    Dune::QuadratureType::Enum qt = Dune::QuadratureType::GaussLegendre)
  {
    return ExpressionPreOperator<Tag, Expr, tag::order>{tag, expr, order, qt};
  }

  /// Store tag and expression and a quadrature rule to create a \ref GridFunctionOperator
  template <class Tag, class Expr, class ctype, int dim>
  auto makeOperator(Tag tag, Expr const& expr, Dune::QuadratureRule<ctype,dim> const& rule)
  {
    return ExpressionPreOperator<Tag, Expr, tag::rule, Dune::QuadratureRule<ctype,dim>>{tag, expr, rule};
  }

  /** @} **/

#ifndef DOXYGEN

  namespace Impl
  {
    // Deduce polynomial order of expression automatically. Create standard quadrature rule with degree
    // build up of operator derivative order, and polynomial degrees of trial/test functions.
    template <class Tag, class Expr, class GridView>
    auto makeQuadCreator(ExpressionPreOperator<Tag, Expr, tag::deduce> const& /*op*/, GridView const& /*gridView*/)
    {
      using QuadratureRules = Dune::QuadratureRules<typename GridView::ctype, GridView::dimension>;
      return [](auto type, auto&& localFct, auto getDegree) -> auto const&
      {
        return QuadratureRules::rule(type, getDegree(order(localFct)));
      };
    }

    // Provide polynomial order of expression explicitly
    template <class Tag, class Expr, class GridView>
    auto makeQuadCreator(ExpressionPreOperator<Tag, Expr, tag::order> const& op, GridView const& /*gridView*/)
    {
      using QuadratureRules = Dune::QuadratureRules<typename GridView::ctype, GridView::dimension>;
      return [order=op.order,qt=op.qt](auto type, auto&& /*localFct*/, auto getDegree) -> auto const&
      {
        return QuadratureRules::rule(type, getDegree(order), qt);
      };
    }

    // Provide quadrature rule explicitly
    template <class Tag, class Expr, class Rule, class GridView>
    auto makeQuadCreator(ExpressionPreOperator<Tag, Expr, tag::rule, Rule> const& op, GridView const& /*gridView*/)
    {
      return [&rule=op.rule](auto&&...) -> auto const&
      {
        return rule;
      };
    }

  } // end namespace Impl


  /// Generate an \ref GridFunctionOperator from a PreOperator (tag, expr).
  /// @{
  template <class Tag, class... Args, class GridView>
  auto makeGridOperator(ExpressionPreOperator<Tag, Args...> const& op, GridView const& gridView)
  {
    auto gridFct = makeGridFunction(op.expr, gridView);
    auto quadCreator = Impl::makeQuadCreator(op, gridView);
    using GridFctOp = GridFunctionOperator<Tag, decltype(gridFct), decltype(quadCreator)>;
    return GridFctOp{op.tag, gridFct, quadCreator};
  }

  template <class Tag, class... Args, class GridView>
  auto makeGridOperator(std::reference_wrapper<ExpressionPreOperator<Tag, Args...>> op, GridView const& gridView)
  {
    ExpressionPreOperator<Tag, Args...> const& op_ref = op;
    auto gridFct = makeGridFunction(std::ref(op_ref.expr), gridView);
    auto quadCreator = Impl::makeQuadCreator(op_ref, gridView);
    using GridFctOp = GridFunctionOperator<Tag, decltype(gridFct), decltype(quadCreator)>;
    return GridFctOp{op_ref.tag, gridFct, quadCreator};
  }
  /// @}


  /// Generate a shared_ptr to \ref GridFunctionOperator from a PreOperator (tag, expr).
  /// @{
  template <class Tag, class... Args, class GridView>
  auto makeGridOperatorPtr(ExpressionPreOperator<Tag, Args...> const& op, GridView const& gridView)
  {
    auto gridFct = makeGridFunction(op.expr, gridView);
    auto quadCreator = Impl::makeQuadCreator(op, gridView);
    using GridFctOp = GridFunctionOperator<Tag, decltype(gridFct), decltype(quadCreator)>;
    return std::make_shared<GridFctOp>(op.tag, gridFct, quadCreator);
  }

  template <class Tag, class... Args, class GridView>
  auto makeGridOperatorPtr(std::reference_wrapper<ExpressionPreOperator<Tag, Args...>> op, GridView const& gridView)
  {
    ExpressionPreOperator<Tag, Args...> const& op_ref = op;
    auto gridFct = makeGridFunction(std::ref(op_ref.expr), gridView);
    auto quadCreator = Impl::makeQuadCreator(op_ref, gridView);
    using GridFctOp = GridFunctionOperator<Tag, decltype(gridFct), decltype(quadCreator)>;
    return std::make_shared<GridFctOp>(op_ref.tag, gridFct, quadCreator);
  }
  /// @}

#endif // DOXYGEN

} // end namespace AMDiS