#pragma once

#include <cmath>

#include <amdis/LinearAlgebra.hpp>
#include <amdis/Observer.hpp>
#include <amdis/OperatorList.hpp>
#include <amdis/common/FlatMatrix.hpp>
#include <amdis/common/TypeTraits.hpp>
#include <amdis/linearalgebra/SymmetryStructure.hpp>
#include <amdis/typetree/TreePath.hpp>

namespace AMDiS
{
  /**
   * Basis implementation of DOFMatrix, i.e. a sparse matrix storing all the
   * assembled Operators indexed with DOF indices. The matrix data is associated
   * to a row and column global basis.
   *
   * \tparam RB  Basis of the matrix rows
   * \tparam CB  Basis of matrix columns
   * \tparam T   Coefficient type to store in the matrix
   **/
  template <class RB, class CB, class T = double, class Traits = BackendTraits<RB>>
  class BiLinearForm
      : public MatrixFacade<T, Traits::template MatrixImpl>
      , private ObserverSequence<event::adapt,2>
  {
    using Self = BiLinearForm;
    using Super = MatrixFacade<T, Traits::template MatrixImpl>;

  public:
    /// The type of the finite element space / basis of the row
    using RowBasis = RB;
    using RowLocalView = typename RowBasis::LocalView;

    /// The type of the finite element space / basis of the column
    using ColBasis = CB;
    using ColLocalView = typename ColBasis::LocalView;

    /// The type of the elements of the DOFVector
    using CoefficientType = T;

    /// The type of the matrix filled on an element with local contributions
    using ElementMatrix = FlatMatrix<CoefficientType>;

    /// Type of the sparsity pattern of the backend
    using SparsityPattern = typename Traits::SparsityPattern;

  public:
    /// Constructor. Stores the row and column basis in a local `shared_ptr` to const
    BiLinearForm(std::shared_ptr<RB> const& rowBasis, std::shared_ptr<CB> const& colBasis)
      : Super(*rowBasis, *colBasis)
      , ObserverSequence<event::adapt,2>(*rowBasis, *colBasis)
      , rowBasis_(rowBasis)
      , colBasis_(colBasis)
    {
      pattern_.init(*rowBasis_, *colBasis_);
      operators_.init(*rowBasis_, *colBasis_);

      auto const rowSize = rowBasis_->localView().maxSize();
      auto const colSize = colBasis_->localView().maxSize();
      elementMatrix_.resize(rowSize, colSize);
    }

    /// Wraps the passed global bases into (non-destroying) shared_ptr
    template <class RB_, class CB_,
      REQUIRES(Concepts::Similar<RB_,RB>),
      REQUIRES(Concepts::Similar<CB_,CB>)>
    BiLinearForm(RB_&& rowBasis, CB_&& colBasis)
      : BiLinearForm(Dune::wrap_or_move(FWD(rowBasis)), Dune::wrap_or_move(FWD(colBasis)))
    {}

    /// Constructor for rowBasis == colBasis
    template <class RB_ = RB, class CB_ = CB,
      REQUIRES(std::is_same<RB_,CB_>::value)>
    explicit BiLinearForm(std::shared_ptr<RB> const& rowBasis)
      : BiLinearForm(rowBasis, rowBasis)
    {}

    /// Wraps the passed row-basis into a (non-destroying) shared_ptr
    template <class RB_,
      REQUIRES(Concepts::Similar<RB_,RB>)>
    explicit BiLinearForm(RB_&& rowBasis)
      : BiLinearForm(Dune::wrap_or_move(FWD(rowBasis)))
    {}

    std::shared_ptr<RowBasis const> const& rowBasis() const { return rowBasis_; }
    std::shared_ptr<ColBasis const> const& colBasis() const { return colBasis_; }

    /// \brief Associate a local operator with this BiLinearForm
    /**
     * Stores an operator in a list that gets assembled during a call to \ref assemble().
     * The operator may be assigned to a specific context, i.e. either an element
     * operator, an intersection operator, or a boundary operator.
     * The \p row and \p col tree paths specify the sub-basis for test and trial
     * functions the operator is applied to.
     *
     * \tparam ContextTag  One of \ref tag::element_operator, \ref tag::intersection_operator
     *                     or \ref tag::boundary_operator indicating where to assemble
     *                     this operator.
     * \tparam Expr        An pre-operator that can be bound to a gridView, or a valid
     *                     GridOperator.
     * \tparam row         A tree-path for the RowBasis
     * \tparam col         A tree-path for the ColBasis
     *
     * [[expects: row is valid tree-path in RowBasis]]
     * [[expects: col is valid tree-path in ColBasis]]
     * @{
     **/
    template <class ContextTag, class Expr, class Row, class Col>
    void addOperator(ContextTag contextTag, Expr const& expr, Row row, Col col);

    // Add an operator to be assembled on the elements of the grid
    template <class Expr, class Row = RootTreePath, class Col = RootTreePath>
    void addOperator(Expr const& expr, Row row = {}, Col col = {})
    {
      using E = typename RowLocalView::Element;
      addOperator(tag::element_operator<E>{}, expr, row, col);
    }

    // Add an operator to be assembled on the intersections of the grid
    template <class Expr, class Row = RootTreePath, class Col = RootTreePath>
    void addIntersectionOperator(Expr const& expr, Row row = {}, Col col = {})
    {
      using I = typename RowLocalView::GridView::Intersection;
      addOperator(tag::intersection_operator<I>{}, expr, row, col);
    }
    /// @}

    /// Assemble the matrix operators on the bound element.
    void assemble(RowLocalView const& rowLocalView,
                  ColLocalView const& colLocalView);

    /// Assemble all matrix operators, TODO: incorporate boundary conditions
    void assemble(SymmetryStructure symmetry = SymmetryStructure::unknown);

    void init(SymmetryStructure symmetry = SymmetryStructure::unknown)
    {
      Super::init(pattern_, symmetry);
    }

    using Super::init;

    void updateImpl(event::adapt e, index_t<0> i) override { updateImpl2(e,i); }
    void updateImpl(event::adapt e, index_t<1> i) override { updateImpl2(e,i); }

  private:
    // Track for each basis wether updated and if both are, update the sparsity pattern
    template <std::size_t I>
    void updateImpl2(event::adapt, index_t<I>)
    {
      assert(!updateCounter_.test(I));
      updateCounter_.set(I);

      if (updateCounter_.all()) {
        pattern_.init(*rowBasis_, *colBasis_);
        updateCounter_.reset();
      }
    }

  protected:
    /// The structure of the non-zeros in the matrix
    SparsityPattern pattern_;

    /// Dense matrix to store coefficients during \ref assemble()
    ElementMatrix elementMatrix_;

    /// List of operators associated to row/col node
    MatrixOperators<RowBasis,ColBasis,ElementMatrix> operators_;

    std::shared_ptr<RowBasis const> rowBasis_;
    std::shared_ptr<ColBasis const> colBasis_;

  private:
    std::bitset<2> updateCounter_ = 0;
  };


#if DUNE_HAVE_CXX_CLASS_TEMPLATE_ARGUMENT_DEDUCTION
  template <class RB, class CB>
  BiLinearForm(RB&&, CB&&)
    -> BiLinearForm<Underlying_t<RB>, Underlying_t<CB>>;

  template <class RB>
  BiLinearForm(RB&&)
    -> BiLinearForm<Underlying_t<RB>, Underlying_t<RB>>;
#endif

  template <class T = double, class RB, class CB>
  auto makeBiLinearForm(RB&& rowBasis, CB&& colBasis)
  {
    using BLF = BiLinearForm<Underlying_t<RB>, Underlying_t<CB>, T>;
    return BLF{FWD(rowBasis), FWD(colBasis)};
  }

  template <class T = double, class RB>
  auto makeBiLinearForm(RB&& rowBasis)
  {
    using BLF = BiLinearForm<Underlying_t<RB>, Underlying_t<RB>, T>;
    return BLF{FWD(rowBasis)};
  }

} // end namespace AMDiS

#include <amdis/BiLinearForm.inc.hpp>