#pragma once

#include <dune/amdis/common/FieldMatVec.hpp>

namespace AMDiS {

template <class Traits, class TreePath>
typename DOFVectorConstView<Traits, TreePath>::Range DOFVectorConstView<Traits, TreePath>::
LocalFunction::operator()(LocalDomain const& x) const
{
  assert( bound_ );
  auto y = Range(0);

  auto&& coefficients = *globalFunction_->dofVector_;
  auto&& nodeToRangeEntry = globalFunction_->nodeToRangeEntry_;
  forEachLeafNode(*subTree_, [&,this](auto const& node, auto)
  {
    using Node = std::decay_t<decltype(node)>;
    using LocalBasisRange = typename LocalFunction::template LocalBasisRange<Node>;
    using MultiIndex = typename LocalIndexSet::MultiIndex;
    using CoefficientBlock = typename std::decay<decltype(std::declval<Vector>()[std::declval<MultiIndex>()])>::type;
    using RangeBlock = std::decay_t<decltype(std::declval<NodeToRangeEntry>()(node, y))>;

    auto&& fe = node.finiteElement();
    auto&& localBasis = fe.localBasis();

    auto&& shapeFunctionValues = shapeFunctionValueContainer_[node];
    localBasis.evaluateFunction(x, shapeFunctionValues);

    // Get range entry associated to this node
    auto&& re = nodeToRangeEntry(node, y);

    for (std::size_t i = 0; i < localBasis.size(); ++i) {
      auto&& multiIndex = localIndexSet_.index(node.localIndex(i));

      // Get coefficient associated to i-th shape function
      auto&& c = coefficients[multiIndex];

      // Get value of i-th shape function
      auto&& v = shapeFunctionValues[i];

      // Notice that the range entry re, the coefficient c, and the shape functions
      // value v may all be scalar, vector, matrix, or general container valued.
      std::size_t dimC = Flat<CoefficientBlock>::size(c);
      std::size_t dimV = Flat<LocalBasisRange>::size(v);
      assert(dimC*dimV == Flat<RangeBlock>::size(re));
      for(std::size_t j = 0; j < dimC; ++j) {
        auto&& c_j = Flat<CoefficientBlock>::getEntry(c, j);
        for(std::size_t k = 0; k < dimV; ++k) {
          auto&& v_k = Flat<LocalBasisRange>::getEntry(v, k);
          Flat<RangeBlock>::getEntry(re, j*dimV + k) += c_j*v_k;
        }
      }
    }
  });

  return y;
}


template <class Traits, class TreePath>
typename DOFVectorConstView<Traits, TreePath>::DerivativeRange DOFVectorConstView<Traits, TreePath>::
GradientLocalFunction::operator()(LocalDomain const& x) const
{
  assert( bound_ );
  Range dy;
  for (std::size_t j = 0; j < dy.size(); ++j)
    dy[j] = 0;

  auto&& coefficients = *globalFunction_->dofVector_;
  auto&& nodeToRangeEntry = globalFunction_->nodeToRangeEntry_;
  forEachLeafNode(*subTree_, [&,this](auto const& node, auto)
  {
    // TODO: may DOFVectorView::Range to FieldVector type if necessary
    using LocalDerivativeTraits = Dune::Functions::DefaultDerivativeTraits<Dune::FieldVector<double,1>(Domain)>;
    using GradientBlock = typename LocalDerivativeTraits::Range;
    using MultiIndex = typename LocalIndexSet::MultiIndex;
    using CoefficientBlock = typename std::decay<decltype(std::declval<Vector>()[std::declval<MultiIndex>()])>::type;
    using RangeBlock = std::decay_t<decltype(std::declval<NodeToRangeEntry>()(node, dy))>;

    auto&& fe = node.finiteElement();
    auto&& localBasis = fe.localBasis();

    // The transposed inverse Jacobian of the map from the reference element to the element
    auto&& jacobian = geometry_.value().jacobianInverseTransposed(x);

    auto&& referenceGradients = referenceGradientContainer_[node];
    localBasis.evaluateJacobian(x, referenceGradients);

    // Compute the shape function gradients on the real element
    std::vector<GradientBlock> gradients(referenceGradients.size());
    for (std::size_t i = 0; i < gradients.size(); ++i)
      multiplies_ABt(referenceGradients[i], jacobian, gradients[i]); // D[phi] * J^(-1) -> grad

    // Get range entry associated to this node
    auto&& re = nodeToRangeEntry(node, dy);

    for (std::size_t i = 0; i < localBasis.size(); ++i) {
      auto&& multiIndex = localIndexSet_.index(node.localIndex(i));

      // Get coefficient associated to i-th shape function
      auto&& c = coefficients[multiIndex];

      // Get value of i-th transformed reference gradient
      auto&& grad = gradients[i];

      // Notice that the range entry re, the coefficient c, and the transformed
      // reference gradient grad may all be scalar, vector, matrix, or general container valued.
      std::size_t dimC = Flat<CoefficientBlock>::size(c);
      std::size_t dimV = Flat<GradientBlock>::size(grad);
      assert(dimC*dimV == std::size_t(Flat<RangeBlock>::size(re)));
      for(std::size_t j = 0; j < dimC; ++j) {
        auto&& c_j = Flat<CoefficientBlock>::getEntry(c, j);
        for(std::size_t k = 0; k < dimV; ++k) {
          auto&& v_k = Flat<GradientBlock>::getEntry(grad, k);
          Flat<RangeBlock>::getEntry(re, j*dimV + k) += c_j*v_k;
        }
      }
    }
  });

  return dy;
}

} // end namespace AMDiS