#pragma once

#include <cmath>

#include <dune/common/timer.hh>

#include <amdis/OperatorList.hpp>
#include <amdis/common/Math.hpp>
#include <amdis/typetree/MultiIndex.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 Backend  A linear-algebra backend for the matrix storage
   **/
  template <class RowBasisType, class ColBasisType, class Backend>
  class DOFMatrixBase
  {
  public:
    /// The type of the finite element space / basis of the row
    using RowBasis = RowBasisType;
    using RowLocalView = typename RowBasis::LocalView;

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

    using Element = typename RowLocalView::Element;
    using Geometry = typename Element::Geometry;

    /// The index/size - type
    using size_type = typename RowBasis::size_type;

    using BaseMatrix = typename Backend::BaseMatrix;
    using ElementMatrix = Dune::DynamicMatrix<double>;

  public:
    /// Constructor. Constructs new BaseVector.
    DOFMatrixBase(RowBasis const& rowBasis, ColBasis const& colBasis)
      : rowBasis_(&rowBasis)
      , colBasis_(&colBasis)
    {
      operators_.init(rowBasis, colBasis);
    }

    DOFMatrixBase(DOFMatrixBase const&) = delete;
    DOFMatrixBase(DOFMatrixBase&&) = delete;

    /// Return the row-basis \ref rowBasis of the matrix
    RowBasis const& rowBasis() const
    {
      return *rowBasis_;
    }

    /// Return the col-basis \ref colBasis of the matrix
    ColBasis const& colBasis() const
    {
      return *colBasis_;
    }

    /// Return the data-matrix
    BaseMatrix const& matrix() const
    {
      return backend_.matrix();
    }

    /// Return the data-matrix
    BaseMatrix& matrix()
    {
      return backend_.matrix();
    }

    /// Return the size of the \ref rowBasis_
    size_type rows() const
    {
      return rowBasis_->dimension();
    }

    /// Return the size of the \ref colBasis_
    size_type cols() const
    {
      return colBasis_->dimension();
    }

    /// Initialize the matrix for insertion, e.g. allocate the non-zero pattern
    /// If \p setToZero is true, the matrix is set to 0
    void init(bool asmMatrix);

    /// Finish the matrix insertion, e.g. cleanup or final insertion
    void finish(bool asmMatrix);

    /// Insert a block of values into the matrix (add to existing values)
    void insert(RowLocalView const& rowLocalView,
                ColLocalView const& colLocalView,
                ElementMatrix const& elementMatrix);

    /// Insert a single value into the matrix (add to existing value)
    template <class RowIndex, class ColIndex>
    void insert(RowIndex row, ColIndex col, typename Backend::value_type const& value)
    {
      backend_.insert(flatMultiIndex(row), flatMultiIndex(col), value);
    }

    /// \brief Associate a local operator with this DOFMatrix
    /**
     * 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 PreOperator  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 Operator,
              class RowTreePath = RootTreePath, class ColTreePath = RootTreePath>
    void addOperator(ContextTag contextTag, Operator const& preOp,
                     RowTreePath row = {}, ColTreePath col = {});

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

    /// Assemble all matrix operators, TODO: incooperate boundary conditions
    void assemble();

    /// \brief Deletes all rows with \p dirichletNodes != 0 and if \p apply is true, adds
    /// a one on the diagonal of the matrix.
    auto applyDirichletBC(std::vector<char> const& dirichletNodes)
    {
      return backend_.applyDirichletBC(dirichletNodes);
    }

    size_type nnz() const
    {
      return backend_.nnz();
    }

  protected:
    /// The finite element space / basis associated with the rows/columns
    RowBasis const* rowBasis_;
    ColBasis const* colBasis_;

    /// Data backend
    Backend backend_;

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

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

} // end namespace AMDiS

#include "DOFMatrixBase.inc.hpp"