#pragma once

#include <utility>

#include <amdis/Assembler.hpp>
#include <amdis/LocalOperator.hpp>
#include <amdis/operations/Assigner.hpp>
#include <amdis/typetree/Traversal.hpp>
#include <amdis/utility/AssembleOperators.hpp>

namespace AMDiS {

template <class GB, class T, class Traits>
  template <class ContextTag, class Expr, class TreePath>
void LinearForm<GB,T,Traits>::
addOperator(ContextTag contextTag, Expr const& expr, TreePath path)
{
  static_assert( Concepts::PreTreePath<TreePath>,
      "path must be a valid treepath, or an integer/index-constant");

  auto i = makeTreePath(path);
  auto node = child(this->basis()->localView().tree(), i);

  using LocalContext = typename ContextTag::type;
  using Tr = DefaultAssemblerTraits<LocalContext, ElementVector>;
  auto op = makeLocalOperator<LocalContext>(expr, this->basis()->gridView());
  auto localAssembler = makeUniquePtr(makeAssembler<Tr>(std::move(op), node));

  operators_[i].push(contextTag, std::move(localAssembler));
}


template <class GB, class T, class Traits>
  template <class LocalView, REQUIRES_(Concepts::LocalView<LocalView>)>
void LinearForm<GB,T,Traits>::
assemble(LocalView const& localView)
{
  elementVector_.resize(localView.size());
  elementVector_ = 0;

  auto const& gv = this->basis()->gridView();
  auto const& element = localView.element();
  auto geometry = element.geometry();

  for_each_node(localView.tree(), [&](auto const& node, auto tp) {
    auto& rhsOp = operators_[tp];
    if (rhsOp) {
      rhsOp.bind(element, geometry);
      assembleOperators(gv, element, rhsOp, makeVectorAssembler(node, elementVector_));
      rhsOp.unbind();
    }
  });

  this->scatter(localView, elementVector_, Assigner::plus_assign{});
}


template <class GB, class T, class Traits>
void LinearForm<GB,T,Traits>::
assemble()
{
  auto localView = this->basis()->localView();

  this->init(sizeInfo(*this->basis()), true);
  for (auto const& element : elements(this->basis()->gridView(), typename Traits::PartitionSet{})) {
    localView.bind(element);
    this->assemble(localView);
    localView.unbind();
  }
  this->finish();
}

} // end namespace AMDiS