#ifndef MULTILEVEL_PATCH_PRECONDITIONER_HH
#define MULTILEVEL_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 "nbodytransfer.hh"
#include "levelpatchpreconditioner.hh"
#include <dune/faultnetworks/dgmgtransfer.hh>

template <class ContactNetwork, class PatchSolver, class MatrixType, class VectorType>
class MultilevelPatchPreconditioner : public LinearIterationStep<MatrixType, VectorType> {
public:
    enum Mode {additive, multiplicative};

private:
    using LevelContactNetwork = typename ContactNetwork::LevelContactNetwork;
    using LevelPatchPreconditioner = LevelPatchPreconditioner<LevelContactNetwork, PatchSolver, MatrixType, VectorType>;
    using LevelGlobalPreconditioner = LevelPatchPreconditioner<LevelContactNetwork, PatchSolver, MatrixType, VectorType>;
    using BitVector = typename LinearIterationStep<MatrixType, VectorType>::BitVector;

    using MGTransfer = ;

    const ContactNetwork& contactNetwork_;
    const Dune::BitSetVector<1>& activeLevels_;

    const Mode mode_;

    std::vector<std::shared_ptr<LevelPatchPreconditioner>> levelPatchPreconditioners_;

    std::vector<MatrixType> levelMat_;
    std::vector<VectorType> levelX_;
    std::vector<VectorType> levelRhs_;
    std::vector<BitVector> levelIgnore_;

    PatchSolver coarseSolver_;

    std::vector<std::shared_ptr<MGTransfer>> mgTransfer_; //TODO

public:
    MultilevelPatchPreconditioner(const ContactNetwork& contactNetwork,
                                  const Dune::BitSetVector<1>& activeLevels,
                                  const Mode mode = MultilevelPatchPreconditioner::Mode::additive) :
          contactNetwork_(contactNetwork),
          activeLevels_(activeLevels),
          mode_(mode) {

        assert(activeLevels.size() == contactNetwork_.nLevels());

        // init level patch preconditioners and multigrid transfer
        levelPatchPreconditioners_.resize(0);
        mgTransfer_.resize(0);

        int maxLevel = -1;
        for (size_t i=0; i<activeLevels_.size(); i++) {
            if (activeLevels_[i][0]) {
                // init global problem on coarsest level
                const auto& levelNetwork = *contactNetwork_.level(i);
                coarseGlobalPreconditioner_ = std::make_shared<LevelGlobalPreconditioner>(levelNetwork);

                maxLevel = i;
                break;
            }
        }

        const int coarsestLevel = maxLevel;

        typename LevelPatchPreconditioner::Mode levelMode = LevelPatchPreconditioner::Mode::additive;
        if (mode_ != additive){
            levelMode = LevelPatchPreconditioner::Mode::multiplicative;
        }

        for (size_t i=maxLevel+1; i<activeLevels_.size(); i++) {
            if (activeLevels_[i][0]) {
                // init local patch preconditioners on each level
                const auto& fineNetwork = contactNetwork_.level(i);
                levelPatchPreconditioners_.push_back(std::make_shared<LevelPatchPreconditioner(*contactNetwork_.level(maxLevel), fineNetwork, levelMode));

                maxLevel = i;
            }
        }

        levelX_.resize(size());
        levelRhs_.resize(size());

        // init multigrid transfer
        for (size_t i=0; i<levelPatchPreconditioners_.size(); i++) {
            mgTransfer_.push_back(std::make_shared<MGTransfer>(...));
        }
        mgTransfer_.push_back(std::make_shared<MGTransfer>(coarseGlobalPreconditioner_->basis(), *allFaultBasis_));
    }


   void setProblem(const MatrixType& mat, VectorType& x, const VectorType& rhs) override {
        this->setProblem(mat, x, rhs);

        for (size_t i=0; i<levelPatchPreconditioners_.size(); i++) {
            const auto& transfer = *mgTransfer_[i];
            auto& _x = levelX_[i];
            auto& _rhs = levelRhs_[i];
            auto& _mat = levelMat_[i];

            transfer.restrict(x, _x);
            transfer.restrict(rhs, _rhs);
            transfer.galerkinRestrict(mat, _mat);

            levelPatchPreconditioners_[i]->setProblem(_mat, _x, _rhs);
        }

        // set coarse global problem
        const auto& coarseTransfer = *mgTransfer_.back();
        auto& coarseX = levelX_.back();
        auto& coarseRhs = levelRhs_.back();
        auto& coarseMat = levelMat_.back();

        coarseTransfer.restrict(x, coarseX);
        coarseTransfer.restrict(rhs, coarseRhs);
        coarseTransfer.galerkinRestrict(mat, coarseMat);

        coarseSolver_.getIterationStep().setProblem(coarseMat, coarseX, coarseRhs);
    }


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

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

        VectorType x;

        // solve coarse global problem
        auto& coarseIgnore = levelIgnore_.back();
        mgTransfer_.back()->restrict(this->ignore(), coarseIgnore);
        coarseSolver_.getIterationStep().setIgnore(coarseIgnore);

        coarseSolver_.check();
        coarseSolver_.preprocess();
        coarseSolver_.solve();

        mgTransfer_.back()->prolong(coarseSolver_.getSol(), x);

        for (size_t i=levelPatchPreconditioners_.size()-2; i>=0; i--) {
            const auto& transfer = *mgTransfer_[i];
            auto& preconditioner = *levelPatchPreconditioners_[i];

            auto& _ignore = levelIgnore_[i];
            transfer.restrict(this->ignore(), _ignore);
            preconditioner.setIgnore(_ignore);

            preconditioner.iterate();

            x += preconditioner.getSol();

            VectorType newX;
            transfer.prolong(x, newX);

            x = newX;
        }

        *(this->x_) = x;
    }


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

        /*
        VectorType x;

        for (int i=(levelPatchPreconditioners_.size()-1); i>=0; i--) {
            VectorType updatedRhs(*(this->rhs_));
            this->mat_->mmv(*(this->x_), updatedRhs);

            mgTransfer_[i]->restrict(*(this->x_), levelX_[i]);
            mgTransfer_[i]->restrict(updatedRhs, levelRhs_[i]);

            levelPatchPreconditioners_[i]->setProblem(*(this->mat_), levelX_[i], levelRhs_[i]);

            levelPatchPreconditioners_[i]->iterate();
            const VectorType& it = levelPatchPreconditioners_[i]->getSol();

            mgTransfer_[i]->prolong(it, x);

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

        VectorType updatedCoarseRhs(*(this->rhs_));
        this->mat_->mmv(*(this->x_), updatedCoarseRhs);

        size_t j = levelPatchPreconditioners_.size();
        mgTransfer_[j]->restrict(*(this->x_), levelX_[j]);
        mgTransfer_[j]->restrict(updatedCoarseRhs, levelRhs_[j]);

        coarseGlobalPreconditioner_->setProblem(*(this->mat_), levelX_[j], levelRhs_[j]);
        coarseGlobalPreconditioner_->iterate();
        const VectorType& it = coarseGlobalPreconditioner_->getSol();

        mgTransfer_[j]->prolong(it, x);

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

    void build() {
        coarseGlobalPreconditioner_->build();

        for (size_t i=0; i<levelPatchPreconditioners_.size(); i++) {
            levelPatchPreconditioners_[i]->build();
        }
    }

    void setPatchSolver(PatchSolver&& patchSolver) {
        for (size_t i=0; i<levelPatchPreconditioners_.size(); i++) {
            levelPatchPreconditioners_[i]->setPatchSolver(std::forward<PatchSolver>(patchSolver));
        }
        coarseGlobalPreconditioner_->setPatchSolver(std::forward<PatchSolver>(patchSolver));
    };

    void setPatchDepth(const size_t patchDepth = 0) {
        for (size_t i=0; i<levelPatchPreconditioners_.size(); i++) {
            levelPatchPreconditioners_[i]->setPatchDepth(patchDepth);
        }
    }

    size_t size() const {
        return levelPatchPreconditioners_.size()+1;
    }
};

#endif