#pragma once

#include <cstddef>
#include <vector>

#include <dune/common/fmatrix.hh>
#include <dune/common/fvector.hh>
#include <dune/common/typetraits.hh>
#include <dune/geometry/type.hh>

#include <amdis/common/DerivativeTraits.hpp>
#include <amdis/common/FieldMatVec.hpp>
#include <amdis/utility/LocalBasisCache.hpp>

namespace AMDiS
{
  /// Traits class for local-to-global basis adaptors
  /**
   * \tparam LocalBasisTraits Traits class of the LocalBasis to be adapted.
   * \tparam dimGlobal        Dimension of the global coordinates,
   *                          i.e. Geometry::coorddimension, if the global
   *                          coordinates are determined by a Geometry.
   *
   * \implements BasisInterface::Traits
   */
  template <class LocalBasisTraits, std::size_t dimGlobal>
  struct LocalToGlobalBasisAdapterTraits
  {
    static const std::size_t dimDomainLocal = LocalBasisTraits::dimDomain;
    static const std::size_t dimDomainGlobal = dimGlobal;
    static const std::size_t dimRange = LocalBasisTraits::dimRange;

    using DomainField = typename LocalBasisTraits::DomainFieldType;
    using DomainLocal = typename LocalBasisTraits::DomainType;
    using DomainGlobal = FieldVector<DomainField, dimDomainGlobal>;

    using RangeField = typename LocalBasisTraits::RangeFieldType;
    using Range = typename LocalBasisTraits::RangeType;

    using GradientRange = typename DerivativeTraits<Range(DomainGlobal), tag::gradient>::Range;
    using PartialRange = typename DerivativeTraits<Range(DomainGlobal), tag::partial>::Range;
  };

  /// \brief Convert a simple (scalar) local basis into a global basis
  /**
   * The local basis must be scalar, i.e. LocalBasis::Traits::dimRange must be 1
   * It's values are not transformed.
   *
   * For scalar function \f$f\f$, the gradient is equivalent to the transposed
   * Jacobian \f$\nabla f|_x = J_f^T(x)\f$.  The Jacobian is thus transformed
   * using
   * \f[
   *   \nabla f|_{\mu(\hat x)} =
   *       \hat J_\mu^{-T}(\hat x) \cdot \hat\nabla\hat f|_{\hat x}
   * \f]
   * Here the hat \f$\hat{\phantom x}\f$ denotes local quantities and
   * \f$\mu\f$ denotes the local-to-global map of the geometry.
   *
   * \tparam BasisNode  Type of the typetree node containing a local basis to adopt.
   * \tparam Geometry   Type of the local-to-global transformation.
   *
   * NOTE: The adapter implements a caching of local basis evaluations at coordinates.
   */
  template <class BasisNode, class Geometry>
  class LocalToGlobalBasisAdapter
  {
    using LocalBasis = typename BasisNode::FiniteElement::Traits::LocalBasisType;

    static_assert(std::is_same<typename LocalBasis::Traits::DomainFieldType, typename Geometry::ctype>::value,
      "LocalToGlobalBasisAdapter: LocalBasis must use the same ctype as Geometry");

    static_assert(std::size_t(LocalBasis::Traits::dimDomain) == std::size_t(Geometry::mydimension),
      "LocalToGlobalBasisAdapter: LocalBasis domain dimension must match local dimension of Geometry");

  public:
    using Cache = NodeCache<BasisNode>;
    using Traits = LocalToGlobalBasisAdapterTraits<typename LocalBasis::Traits, Geometry::coorddimension>;

    /// \brief Construct a LocalToGlobalBasisAdapter
    /**
     * \param node     The basis node in the typetree containing the local basis to adopt.
     * \param geometry The geometry object to use for adaption.
     *
     * \note This class stores the references passed here.  Any use of this
     *       class after these references have become invalid results in
     *       undefined behaviour.  The exception is that the destructor of
     *       this class may still be called.
     */
    LocalToGlobalBasisAdapter(BasisNode const& node, Geometry const& geometry)
      : localBasis_(node.finiteElement().localBasis())
      , localBasisCache_(localBasis_)
      , geometry_(geometry)
      , size_(localBasis_.size())
    {}

    /// Return the number of local basis functions
    std::size_t size() const { return size_; }

    /// \brief Return maximum polynomial order of the base function
    /**
     * This is to determine the required quadrature order.  For an affine
     * geometry this is the same order as for the local basis.  For other
     * geometries this returns the order of the local basis plus the global
     * dimension minus 1.  The assumption for non-affine geometries is that
     * they are still multi-linear.
     */
    std::size_t order() const
    {
      if (geometry_.affine())
        // affine linear
        return localBasis_.order();
      else
        // assume at most order dim
        return localBasis_.order() + Traits::dimDomainGlobal - 1;
    }

    /// Evaluate the local basis functions in the local coordinate `x`
    void evaluateFunction(typename Traits::DomainLocal const& x,
                          std::vector<typename Traits::Range>& out) const
    {
      out = localBasisCache_.evaluateFunction(geometry_.type(), x);
    }

    /// Evaluate the local basis functions in the local coordinate `x` and
    /// return the result using a reference to a thread_local (or static) vector.
    auto const& valuesAt(typename Traits::DomainLocal const& x) const
    {
      return localBasisCache_.evaluateFunction(geometry_.type(), x);
    }

    /// Return the full (global) gradient of the local basis functions in
    /// the local coordinate `x`
    void evaluateGradient(typename Traits::DomainLocal const& x,
                          std::vector<typename Traits::GradientRange>& out) const
    {
      auto const& localJacobian = localBasisCache_.evaluateJacobian(geometry_.type(), x);
      auto const& geoJacobian = geometry_.jacobianInverseTransposed(x);

      out.resize(size_);
      for (std::size_t i = 0; i < size_; ++i)
        geoJacobian.mv(Dune::MatVec::as_vector(localJacobian[i]),
                       Dune::MatVec::as_vector(out[i]));
    }

    /// Return the full (global) gradient of the local basis functions in
    /// the local coordinate `x` and return the result using a reference
    /// to a thread_local (or static) vector.
    auto const& gradientsAt(typename Traits::DomainLocal const& x) const
    {
      thread_local std::vector<typename Traits::GradientRange> grad;
      evaluateGradient(x, grad);
      return grad;
    }

    /// Return the (global) partial derivative in direction `comp` of the
    /// local basis functions in the local coordinate `x`
    void evaluatePartial(typename Traits::DomainLocal const& x,
                         std::size_t comp,
                         std::vector<typename Traits::PartialRange>& out) const
    {
      auto const& localJacobian = localBasisCache_.evaluateJacobian(geometry_.type(), x);
      // NOTE: geoJacobian might be a Dune::DiagonalMatrix with limited interface!
      auto const& geoJacobian = geometry_.jacobianInverseTransposed(x);

      out.resize(size_);
      typename Traits::GradientRange grad;
      auto&& grad_ = Dune::MatVec::as_vector(grad);
      for (std::size_t i = 0; i < size_; ++i) {
        geoJacobian.mv(Dune::MatVec::as_vector(localJacobian[i]), grad_);
        out[i] = grad_[comp];
      }
    }

    /// Return the (global) partial derivative in direction `comp` of the
    /// local basis functions in the local coordinate `x` and return the
    /// result using a reference to a thread_local (or static) vector.
    auto const& partialsAt(typename Traits::DomainLocal const& x, std::size_t comp) const
    {
      thread_local std::vector<typename Traits::PartialRange> d_comp;
      evaluatePartial(x, comp, d_comp);
      return d_comp;
    }

  private:
    /// Reference to the local basis bound to the adapter
    LocalBasis const& localBasis_;

    /// A basis cache for the evaluation at local coordinates
    Cache localBasisCache_;

    /// Reference to the geometry this adapter is bound to
    Geometry const& geometry_;

    /// The number of basis functions
    std::size_t size_;
  };


  /// Generator function for the \ref LocalToGlobalBasisAdapter. \relates LocalToGlobalBasisAdapter.
  template <class BasisNode, class Geometry>
  LocalToGlobalBasisAdapter<BasisNode, Geometry>
  makeLocalToGlobalBasisAdapter(BasisNode const& node, Geometry const& geometry)
  {
    return LocalToGlobalBasisAdapter<BasisNode, Geometry>{node, geometry};
  }

} // namespace AMDiS