#ifndef DUNE_CURVEDGRID_ANALYTIC_DISCRETEFUNCTION_HH
#define DUNE_CURVEDGRID_ANALYTIC_DISCRETEFUNCTION_HH

#include <optional>
#include <type_traits>
#include <utility>

#if !HAVE_DUNE_LOCALFUNCTIONS
#error "Need dune-localfunctions for the definition of the AnalyticDiscreteFunction"
#endif

#include <dune/common/typeutilities.hh>
#include <dune/curvedgrid/gridfunctions/discretegridviewfunction.hh>
#include <dune/curvedgrid/gridfunctions/gridentityset.hh>
#include <dune/localfunctions/lagrange/lfecache.hh>

namespace Dune {

/// \brief LocalFunction associated to the \ref AnalyticDiscreteFunction
/**
 * \tparam LocalContext  Context this localFunction can be bound to, e.g. a grid element
 * \tparam F             Type of a function that can be evaluated in global coordinates.
 * \tparam derivativeOrder  Order of the derivative of the functor to return on evaluation
 **/
template <class LocalContext, class F, int derivativeOrder = 0>
class LocalAnalyticDiscreteFunction;

/// \brief Generator for \ref LocalAnalyticDiscreteFunction
/**
 * \param  ff            Function that can be evaluated at global coordinates of the \ref LocalContext
 * \param  order         Polynomial order of the approximation of ff
 * \tparam LocalContext  Context this localFunction can be bound to, e.g. a grid element
 **/
template <class LocalContext, class FF>
auto localAnalyticDiscreteFunction (FF&& ff, int order)
{
  using F = std::decay_t<FF>;
  return LocalAnalyticDiscreteFunction<LocalContext, F>{std::forward<FF>(ff), order};
}


template <class LC, class Functor, int derivativeOrder>
class LocalAnalyticDiscreteFunction
{
public:
  using LocalContext = LC;
  using Geometry = typename LocalContext::Geometry;
  using ctype = typename Geometry::ctype;

  using Domain = typename Geometry::GlobalCoordinate;
  using RangeType = std::invoke_result_t<Functor,Domain>;

  template <int degree>
  using DerivativeRange = typename Impl::DerivativeRangeType<RangeType(Domain), degree, Functions::DefaultDerivativeTraits>::type;

  using LocalDomain = typename Geometry::LocalCoordinate;
  using Range = DerivativeRange<derivativeOrder>;
  using Signature = Range(LocalDomain);

public:
  /// \brief Constructor. Stores the functor f by value
  template <class FF>
  LocalAnalyticDiscreteFunction (FF&& ff, int order)
    : f_(std::forward<FF>(ff))
    , order_(order)
    , cache_(order)
  {}

  /// \brief bind the LocalFunction to the local context
  /**
   * Stores the localContext and its geometry in a cache variable
   **/
  void bind (const LocalContext& localContext)
  {
    localContext_.emplace(localContext);
    geometry_.emplace(localContext_->geometry());

    lfe_ = &cache_.get(localContext_->type());
    lfe_->localInterpolation().interpolate([&](const auto& local)
    {
      return f_(geometry_->global(local));
    }, coefficients_);
  }

  /// \brief unbind the localContext from the localFunction
  /**
   * Reset the geometry
   **/
  void unbind ()
  {
    lfe_ = nullptr;
    geometry_.reset();
    localContext_.reset();
  }

  /// \brief evaluate the stored function in local coordinates
  /**
   * Transform the local coordinates to global coordinates first,
   * then evalute the stored functor.
   **/
  Range operator() (const LocalDomain& local) const
  {
    static_assert(derivativeOrder < 2, "Higher-order derivatives not implemented");

    if constexpr (derivativeOrder == 0)
      return evaluateFunction(local);
    else if constexpr (derivativeOrder == 1)
      return evaluateJacobian(local);
  }

  friend LocalAnalyticDiscreteFunction<LC,Functor,derivativeOrder+1> derivative (const LocalAnalyticDiscreteFunction& lf)
  {
    return {lf.f_, lf.order_};
  }

  /// \brief return the bound localContext.
  const LocalContext& localContext () const
  {
    assert(!!localContext_);
    return *localContext_;
  }

  /// \brief obtain the functor
  const Functor& f () const
  {
    return f_;
  }

private:
  DerivativeRange<0> evaluateFunction (const LocalDomain& local) const
  {
    assert(!!lfe_);

    const auto& localBasis = lfe_->localBasis();

    localBasis.evaluateFunction(local, shapeValues_);
    assert(coefficients_.size() == shapeValues_.size());

    DerivativeRange<0> nh(0);
    for (std::size_t i = 0; i < shapeValues_.size(); ++i)
      nh.axpy(shapeValues_[i], coefficients_[i]);

    return nh;
  }

  DerivativeRange<1> evaluateJacobian (const LocalDomain& local) const
  {
    assert(!!geometry_);
    assert(!!lfe_);

    const auto& localBasis = lfe_->localBasis();

    // evaluate basis functions in local coordinate
    localBasis.evaluateJacobian(local, shapeGradients_);
    assert(coefficients_.size() == shapeGradients_.size());

    // transform gradients to global coordinates
    auto jit = geometry_->jacobianInverseTransposed(local);
    gradients_.resize(shapeGradients_.size());
    for (std::size_t i = 0; i < shapeGradients_.size(); ++i)
      jit.mv(shapeGradients_[i][0], gradients_[i]);

    DerivativeRange<1> J(0);
    for (std::size_t i = 0; i < coefficients_.size(); ++i)
      for (std::size_t j = 0; j < J.N(); ++j)
        J[j].axpy(coefficients_[i][j], gradients_[i]);

    return J;
  }

private:
  Functor f_;
  int order_;

  using LFECache = LagrangeLFECache<ctype,ctype,Geometry::mydimension>;
  LFECache cache_;

  using LFE = typename LFECache::FiniteElementType;
  LFE const* lfe_ = nullptr;

  std::vector<RangeType> coefficients_;

  using LB = typename LFE::Traits::LocalBasisType;
  mutable std::vector<typename LB::Traits::RangeType> shapeValues_;
  mutable std::vector<typename LB::Traits::JacobianType> shapeGradients_;
  mutable std::vector<FieldVector<ctype,Geometry::coorddimension>> gradients_;

  // some caches
  std::optional<LocalContext> localContext_;
  std::optional<Geometry> geometry_;
};


/// \brief GridFunction associated to the mapping F
/**
 * \tparam Grid       The grid type with elements the corresponding LocalFunction can be bound to
 * \tparam F          Type of a function that can be evaluated in global coordinates.
 * \tparam ORDER      Polynomial order of the local lagrange basis functions (optional)
 **/
template <class Grid, class F, int ORDER = -1>
class AnalyticDiscreteFunction;

/// \brief Generator for \ref AnalyticDiscreteFunction
/**
 * \param  ff    Function that can be evaluated at global coordinates of the \ref Grid
 * \tparam Grid  The grid type with elements the corresponding LocalFunction can be bound to
 **/
template <class FF, class Grid>
auto analyticDiscreteFunction (FF&& ff, const Grid& /*grid*/, int order)
{
  using F = std::decay_t<FF>;
  return AnalyticDiscreteFunction<Grid, F>{std::forward<FF>(ff), order};
}

template <int ORDER, class FF, class Grid>
auto analyticDiscreteFunction (FF&& ff, const Grid& /*grid*/)
{
  using F = std::decay_t<FF>;
  return AnalyticDiscreteFunction<Grid, F, ORDER>{std::forward<FF>(ff)};
}


template <class GridType, class Functor, int ORDER>
class AnalyticDiscreteFunction
{
public:
  using Grid = GridType;
  using EntitySet = GridEntitySet<Grid,0>;

  using Domain = typename EntitySet::GlobalCoordinate;
  using Range = std::invoke_result_t<Functor,Domain>;
  using Signature = Range(Domain);

public:
  /// \brief Constructor. Stores the functor f by value
  template <class FF,
    disableCopyMove<AnalyticDiscreteFunction, FF> = 0>
  AnalyticDiscreteFunction (FF&& ff, int order = ORDER)
    : f_(std::forward<FF>(ff))
    , order_(order)
  {}

  /// \brief evaluate the stored function in global coordinates
  Range operator() (const Domain& x) const
  {
    return f_(x);
  }

  /// \brief construct the \ref LocalAnalyticDiscreteFunction
  using LocalFunction = LocalAnalyticDiscreteFunction<typename EntitySet::Element, Functor>;
  friend LocalFunction localFunction (const AnalyticDiscreteFunction& t)
  {
    return LocalFunction(t.f_, t.order_);
  }

  /// \brief obtain the stored \ref GridEntitySet
  const EntitySet& entitySet () const
  {
    return entitySet_;
  }

  /// \brief obtain the functor
  const Functor& f () const
  {
    return f_;
  }

  /// \brief obtain the order
  const int order () const
  {
    return order_;
  }

private:
  Functor f_;
  int order_;

  EntitySet entitySet_;
};

} // end namespace Dune

#endif // DUNE_CURVEDGRID_ANALYTIC_DISCRETEFUNCTION_HH