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

#include "backward_euler.hh"

#include "../../utils/debugutils.hh"

template <class Vector, class Matrix, class BoundaryFunctions, class BoundaryNodes>
BackwardEuler<Vector, Matrix, BoundaryFunctions, BoundaryNodes>::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 BoundaryNodes& _dirichletNodes,
        const BoundaryFunctions& _dirichletFunctions) :
    RateUpdater<Vector, Matrix, BoundaryFunctions, BoundaryNodes>(
        _matrices, _u_initial, _v_initial, _a_initial, _dirichletNodes,
        _dirichletFunctions) {}

template <class Vector, class Matrix, class BoundaryFunctions, class BoundaryNodes>
void BackwardEuler<Vector, Matrix, BoundaryFunctions, BoundaryNodes>::setup(
        const std::vector<Vector>& ell,
        double _tau,
        double relativeTime,
        std::vector<Vector>& rhs, std::vector<Vector>& iterate,
        std::vector<Matrix>& AM) {

    const size_t bodyCount = this->u.size();

    rhs.resize(bodyCount);
    iterate.resize(bodyCount);
    AM.resize(bodyCount);

    for (size_t i=0; i<bodyCount; i++) {
        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;

    for (size_t i=0; i<bodyCount; i++) {
        auto& bodyIterate = iterate[i];

        const auto& bodyDirichletNodes = this->dirichletNodes[i];
        const auto& bodyDirichletFunctions = this->dirichletFunctions[i];
        for (size_t bc=0; bc<bodyDirichletNodes.size(); ++bc) {
            const auto& bcDirichletNodes = *bodyDirichletNodes[bc];
            size_t dim = bcDirichletNodes[0].size();

            double dirichletValue;
            bodyDirichletFunctions[bc]->evaluate(relativeTime, dirichletValue);

            for (size_t k=0; k<bcDirichletNodes.size() ; ++k) {
                for (size_t j=0; j<dim; ++j) {
                    if (bcDirichletNodes[k][j]) {
                        bodyIterate[k][j] = dirichletValue;
                    }
                }
            }
        }
    }
}

template <class Vector, class Matrix, class BoundaryFunctions, class BoundaryNodes>
void BackwardEuler<Vector, Matrix, BoundaryFunctions, BoundaryNodes>::velocityObstacles(const Vector& u0, const Vector& uObstacles, const Vector& v0, Vector& v1Obstacles) {
    // v1 = 1/tau ( u1 - u0 )
    v1Obstacles.resize(u0.size());
    v1Obstacles = 0;

    Dune::MatrixVector::addProduct(v1Obstacles, 1.0 / this->tau, uObstacles);
    Dune::MatrixVector::subtractProduct(v1Obstacles, 1.0 / this->tau, u0);
}

template <class Vector, class Matrix, class BoundaryFunctions, class BoundaryNodes>
void BackwardEuler<Vector, Matrix, BoundaryFunctions, BoundaryNodes>::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 BoundaryFunctions, class BoundaryNodes>
auto BackwardEuler<Vector, Matrix, BoundaryFunctions, BoundaryNodes>::clone() const
-> std::shared_ptr<RateUpdater<Vector, Matrix, BoundaryFunctions, BoundaryNodes>> {

    return std::make_shared<BackwardEuler<Vector, Matrix, BoundaryFunctions, BoundaryNodes>>(*this);
}