#ifndef LEVEL_PATCH_PRECONDITIONER_HH
#define LEVEL_PATCH_PRECONDITIONER_HH

#include <string>

#include <dune/common/timer.hh>
#include <dune/common/fvector.hh>
#include <dune/common/bitsetvector.hh>

#include <dune/solvers/iterationsteps/lineariterationstep.hh>

#include <dune/faultnetworks/assemblers/globalfaultassembler.hh>
#include <dune/faultnetworks/localproblem.hh>
#include <dune/faultnetworks/levelinterfacenetwork.hh>
#include <dune/faultnetworks/utils/debugutils.hh>

#include <dune/fufem/boundarypatch.hh>
#include <dune/fufem/functiontools/boundarydofs.hh>

template <class BasisType, class LocalAssembler, class LocalInterfaceAssembler, class MatrixType, class VectorType>
class LevelPatchPreconditioner : public LinearIterationStep<MatrixType, VectorType> {

public:
    enum Mode {ADDITIVE, MULTIPLICATIVE};
    enum BoundaryMode {homogeneous, fromIterate};
    enum Direction {FORWARD, BACKWARD, SYMMETRIC};

protected:
    typedef typename BasisType::GridView GridView;
    typedef typename GridView::Grid GridType;

    const LevelInterfaceNetwork<GridView>& levelInterfaceNetwork_;
    const BasisType& patchLevelBasis_;
    const LocalAssembler& localAssembler_;
    const std::vector<std::shared_ptr<LocalInterfaceAssembler>>& localInterfaceAssemblers_;

    const Mode mode_;
    Direction multDirection_;

    const GridType& grid_;
    const int level_;
    const BasisType basis_;

    size_t patchDepth_;
    BoundaryMode boundaryMode_;

    MatrixType matrix_;
    std::vector<std::shared_ptr<OscLocalProblem<MatrixType, VectorType>> > localProblems_;

public:
    LevelPatchPreconditioner(const LevelInterfaceNetwork<GridView>& levelInterfaceNetwork,
                             const BasisType& patchLevelBasis,
                             const LocalAssembler& localAssembler,
                             const std::vector<std::shared_ptr<LocalInterfaceAssembler>>& localInterfaceAssemblers,
                             const Mode mode = LevelPatchPreconditioner::Mode::ADDITIVE) :
          levelInterfaceNetwork_(levelInterfaceNetwork),
          patchLevelBasis_(patchLevelBasis),
          localAssembler_(localAssembler),
          localInterfaceAssemblers_(localInterfaceAssemblers),
          mode_(mode),
          grid_(levelInterfaceNetwork_.grid()),
          level_(levelInterfaceNetwork_.level()),
          basis_(levelInterfaceNetwork_)
    {
        assert(localInterfaceAssemblers_.size() == levelInterfaceNetwork_.size());
        setPatchDepth();
        setBoundaryMode();
        setDirection();
        setup();
    }

    // has to setup localProblems_
    virtual void build() = 0;

    void setPatchDepth(const size_t patchDepth = 0) {
        patchDepth_ = patchDepth;
    }

    void setBoundaryMode(const BoundaryMode boundaryMode = LevelPatchPreconditioner::BoundaryMode::homogeneous) {
        boundaryMode_ = boundaryMode;
    }

    virtual void setProblem(const MatrixType& mat, VectorType& x, const VectorType& rhs) {
        this->x_ = &x;
        this->rhs_ = &rhs;
        this->mat_ = Dune::stackobject_to_shared_ptr(mat);

        for (size_t i=0; i<localProblems_.size(); i++) {
            if (boundaryMode_ == BoundaryMode::homogeneous)
                localProblems_[i]->updateRhs(rhs);
            else
                localProblems_[i]->updateRhsAndBoundary(rhs, x);
        }
    }

    void setDirection(Direction dir = SYMMETRIC) {
        multDirection_ = dir;
    }

    virtual void iterate() {
        if (mode_ == ADDITIVE)
            iterateAdd();
        else
            iterateMult();
    }

    const BasisType& basis() const {
        return basis_;
    }

    const GridView& gridView() const {
        return basis_.getGridView();
    }

    const LevelInterfaceNetwork<GridView>& levelInterfaceNetwork() const {
        return levelInterfaceNetwork_;
    }

    size_t size() const {
        return localProblems_.size();
    }

private:
    void setup() {
        // assemble stiffness matrix for entire level including all faults
        GlobalFaultAssembler<BasisType, BasisType> globalFaultAssembler(basis_, basis_, levelInterfaceNetwork_);
        globalFaultAssembler.assembleOperator(localAssembler_, localInterfaceAssemblers_, matrix_);

        // set boundary conditions
        Dune::BitSetVector<1> globalBoundaryDofs;
        BoundaryPatch<GridView> boundaryPatch(levelInterfaceNetwork_.levelGridView(), true);
        constructBoundaryDofs(boundaryPatch, basis_, globalBoundaryDofs);

        typedef typename MatrixType::row_type RowType;
        typedef typename RowType::ConstIterator ColumnIterator;

        for(size_t i=0; i<globalBoundaryDofs.size(); i++) {
            if(!globalBoundaryDofs[i][0])
                continue;

            RowType& row = matrix_[i];

            ColumnIterator cIt    = row.begin();
            ColumnIterator cEndIt = row.end();

            for(; cIt!=cEndIt; ++cIt) {
                row[cIt.index()] = 0;
            }
            row[i] = 1;
        }
    }

    void iterateAdd() {
        //*(this->x_) = 0;

        VectorType it, x;
        for (size_t i=0; i<localProblems_.size(); i++) {
            localProblems_[i]->solve(it);
            localProblems_[i]->prolong(it, x);

            /*if (i==5) {
                writeToVTK(basis_, x, "/storage/mi/podlesny/data/faultnetworks/iterates/", "exactvertexdata_patchDepth_"+std::to_string(patchDepth_));
            }*/

            *(this->x_) += x;
        }
    }

    void iterateMult() {
        if (multDirection_ != BACKWARD) {
            for (size_t i=0; i<localProblems_.size(); i++)
                iterateStep(i);
        }

        if (multDirection_ != Direction::FORWARD)
            for (size_t i=localProblems_.size()-1; i>=0 && i<localProblems_.size(); i--)
                iterateStep(i);
    }

    void iterateStep(size_t i) {
        //*(this->x_) = 0;

        auto& localProblem = *localProblems_[i];

        VectorType v, v1, v2, v3;
        localProblem.restrict(*(this->x_), v);

        v1.resize(localProblem.getLocalMat().N());
        localProblem.getLocalMat().mv(v, v1);
        v3.resize(matrix_.N());
        matrix_.mv(*(this->x_), v3);
        localProblem.restrict(v3, v2);

        //print(v1, "v1");
        //print(v2, "v2");

        VectorType r1;
        localProblem.restrict(*(this->rhs_), r1);
        Dune::MatrixVector::subtractProduct(r1, localProblem.getLocalMat(), v);

        /*localProblem.updateLocalRhs(r);

        VectorType x;
        localProblem.solve(x);
        v += x;

        localProblem.updateVector(v, *(this->x_));*/

        VectorType localR;
        VectorType r = *(this->rhs_);
        Dune::MatrixVector::subtractProduct(r, matrix_, *(this->x_));
        localProblem.restrict(r, localR);

       // print(r1, "r1");
       // print(localR, "localR");

        localProblem.updateLocalRhs(localR);

        VectorType x;
        localProblem.solve(x);

        //VectorType v;
        localProblem.prolong(x, v);
        *(this->x_) += v;
    }
};

#endif