#include <dune/matrix-vector/axpy.hh>
#include <dune/istl/matrixindexset.hh>

#include "backward_euler.hh"

template <class Vector, class Matrix, class Function, size_t dim>
BackwardEuler<Vector, Matrix, Function, dim>::BackwardEuler(
    const Matrices<Matrix,2>& _matrices, const std::vector<Vector>& _u_initial,
    const std::vector<Vector>& _v_initial, const std::vector<Vector>& _a_initial,
    const std::vector<Dune::BitSetVector<dim>>& _dirichletNodes,
    const std::vector<const Function*>& _dirichletFunctions)
    : RateUpdater<Vector, Matrix, Function, dim>(
          _matrices, _u_initial, _v_initial, _a_initial, _dirichletNodes,
          _dirichletFunctions) {}

template <class Vector, class Matrix, class Function, size_t dim>
void BackwardEuler<Vector, Matrix, Function, dim>::setup(const std::vector<Vector>& ell,
                                                         double _tau,
                                                         double relativeTime,
                                                         std::vector<Vector>& rhs, std::vector<Vector>& iterate,
                                                         std::vector<Matrix>& AM) {
  for (size_t i=0; i<this->u.size(); i++) {
    this->dirichletFunctions[i]->evaluate(relativeTime, this->dirichletValue);
    this->tau = _tau;

    /* We start out with the formulation

           M a + C v + A u = ell

         Backward Euler means

           a1 = 1.0/tau ( v1 - v0 )
           u1 = tau v1 + u0

         in summary, we get at time t=1

           M [1.0/tau ( v1 - v0 )] + C v1
           + A [tau v1 + u0] = ell

         or

           1.0/tau M v1 + C v1 + tau A v1
           = [1.0/tau M + C + tau A] v1
           = ell + 1.0/tau M v0 - A u0
    */

    // set up LHS (for fixed tau, we'd only really have to do this once)
    Matrix& LHS = AM[i];
    {
      Dune::MatrixIndexSet indices(this->matrices.elasticity[i]->N(),
                                   this->matrices.elasticity[i]->M());
      indices.import(*this->matrices.elasticity[i]);
      indices.import(*this->matrices.mass[i]);
      indices.import(*this->matrices.damping[i]);
      indices.exportIdx(LHS);
    }
    LHS = 0.0;
    Dune::MatrixVector::addProduct(LHS, 1.0 / this->tau, *this->matrices.mass[i]);
    Dune::MatrixVector::addProduct(LHS, 1.0, *this->matrices.damping[i]);
    Dune::MatrixVector::addProduct(LHS, this->tau, *this->matrices.elasticity[i]);

    // set up RHS
    {
      Vector& rhss = rhs[i];
      rhss = ell[i];
      Dune::MatrixVector::addProduct(rhss, 1.0 / this->tau, *this->matrices.mass[i],
                           this->v_o[i]);
      Dune::MatrixVector::subtractProduct(rhss, *this->matrices.elasticity[i], this->u_o[i]);
    }

    iterate = this->v_o;

    const Dune::BitSetVector<dim>& dirichletNodess = this->dirichletNodes[i];
    for (size_t k = 0; k < dirichletNodess.size(); ++k)
      for (size_t j = 0; j < dim; ++j)
        if (dirichletNodess[k][j])
          iterate[k][j] = (j == 0) ? this->dirichletValue : 0;
  }
}

template <class Vector, class Matrix, class Function, size_t dim>
void BackwardEuler<Vector, Matrix, Function, dim>::postProcess(
    const std::vector<Vector>& iterate) {
  this->postProcessCalled = true;

  this->v = iterate;

  this->u = this->u_o;

  for (size_t i=0; i<this->u.size(); i++) {
    Dune::MatrixVector::addProduct(this->u[i], this->tau, this->v[i]);

    Vector& ai = this->a[i];
    ai = this->v[i];
    ai -= this->v_o[i];
    ai /= this->tau;
  }
}

template <class Vector, class Matrix, class Function, size_t dim>
std::shared_ptr<RateUpdater<Vector, Matrix, Function, dim>>
BackwardEuler<Vector, Matrix, Function, dim>::clone() const {
  return std::make_shared<BackwardEuler<Vector, Matrix, Function, dim>>(*this);
}