#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/common/parametertree.hh>

#include <dune/solvers/norms/energynorm.hh>
#include <dune/solvers/solvers/loopsolver.hh>
#include <dune/solvers/solvers/umfpacksolver.hh>
#include <dune/solvers/iterationsteps/blockgssteps.hh>
//#include <dune/solvers/iterationsteps/blockgsstep.hh>
#include <dune/solvers/iterationsteps/cgstep.hh>
#include <dune/solvers/iterationsteps/lineariterationstep.hh>
#include <dune/solvers/iterationsteps/truncatedblockgsstep.hh>

#include "nbodycontacttransfer.hh"
#include "levelpatchpreconditioner.hh"

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

template <class ContactNetwork, class MatrixType, class VectorType>
class MultilevelPatchPreconditioner : public LinearIterationStep<MatrixType, VectorType> {
private:
    using Base = LinearIterationStep<MatrixType, VectorType>;

    using PatchSolver = LoopSolver<Vector, BitVector>;
    using PatchSolverStep = TruncatedBlockGSStep<MatrixType, VectorType>;
    using CoarseSolver = Dune::Solvers::UMFPackSolver<MatrixType, VectorType>;

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

    using MGTransfer = NBodyContactTransfer<ContactNetwork, VectorType>;

    static const int dim = ContactNetwork::dim;

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

    const MPPMode mode_;

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

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

    std::vector<BitVector> ignoreHierarchy_;

    std::vector<std::shared_ptr<EnergyNorm<MatrixType, VectorType>>> levelErrorNorms_;
    std::vector<std::shared_ptr<LinearIterationStep<MatrixType, VectorType>>> levelItSteps_;
    std::vector<std::shared_ptr<PatchSolver>> levelSolvers_;
    CoarseSolver coarseSolver_;

    //std::vector<BitVector> recompute_;
    std::vector<std::shared_ptr<MGTransfer>> mgTransfer_;

public:
    MultilevelPatchPreconditioner(const Dune::ParameterTree& parset,
                                  const ContactNetwork& contactNetwork,
                                  const Dune::BitSetVector<1>& activeLevels) :
          contactNetwork_(contactNetwork),
          activeLevels_(activeLevels),
          mode_(parset.get<MPPMode>("mode")){

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

        // init level patch preconditioners and multigrid transfer
        levelPatchPreconditioners_.resize(1);
        levelPatchPreconditioners_[0] = nullptr; // blank level representing global coarse problem to make further indexing easier

        int maxLevel = -1;
        for (size_t i=0; i<activeLevels_.size(); i++) {
            if (activeLevels_[i][0]) {
                maxLevel = i;
                break;
            }
        }

        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, mode_));
                levelPatchPreconditioners_.back()->setPatchDepth(parset.get<size_t>("patchDepth"));
                maxLevel = i;
            }
        }

        levelMat_.resize(size()-1);
        levelX_.resize(size());
        levelRhs_.resize(size());

        // init patch solvers
        levelSolvers_.resize(size());
        levelItSteps_.resize(size());
        levelErrorNorms_.resize(size());

        for (size_t i=1; i<levelSolvers_.size(); i++) {
            //auto gsStep = Dune::Solvers::BlockGSStepFactory<MatrixType, VectorType>::create(Dune::Solvers::BlockGS::LocalSolvers::direct(0.0));
            levelItSteps_[i] = std::make_shared<TruncatedBlockGSStep<MatrixType, VectorType>>(); //Dune::Solvers::BlockGSStepFactory<MatrixType, VectorType>::createPtr(Dune::Solvers::BlockGS::LocalSolvers::direct(0.0));
            levelErrorNorms_[i] = std::make_shared<EnergyNorm<MatrixType, VectorType>>(*levelItSteps_[i].get());
            levelSolvers_[i] = std::make_shared<PatchSolver>(*levelItSteps_[i].get(),
                                       parset.get<size_t>("maximumIterations"),
                                       parset.get<double>("tolerance"),
                                       *levelErrorNorms_[i].get(),
                                       parset.get<Solver::VerbosityMode>("verbosity"));
            levelPatchPreconditioners_[i]->setPatchSolver(levelSolvers_[i]);
        }

        // init multigrid transfer
        mgTransfer_.resize(levelPatchPreconditioners_.size()-1);
        for (size_t i=0; i<mgTransfer_.size(); i++) {
            mgTransfer_[i] = std::make_shared<MGTransfer>();
            mgTransfer_[i]->setup(contactNetwork_, levelPatchPreconditioners_[i+1]->level(), levelPatchPreconditioners_[i+1]->fineLevel());
        }

        // from dune/contact/solvers/nsnewtonmgstep.cc
        /*// Remove any recompute filed so that initially the full transferoperator is assembled
        for (size_t i=0; i<mgTransfer_.size(); i++)
            std::dynamic_pointer_cast<TruncatedMGTransfer<VectorType> >(mgTransfer_[i])->setRecomputeBitField(nullptr);

        // specify the subset of entries to be reassembled at each iteration
        recompute_.resize(size());
        recompute_[recompute_.size()-1] = contactNetwork_.nBodyAssembler().totalHasObstacle_;
        for (int i=mgTransfer_.size(); i>=1; i--) {
            std::dynamic_pointer_cast<TruncatedMGTransfer<VectorType> >(mgTransfer_[i-1])->restrict(recompute_[i], recompute_[i-1]);
          //  std::dynamic_pointer_cast<DenseMultigridTransfer<VectorType, BitVector, MatrixType> >(mgTransfer_[i])->restrict(recompute_[i+1], recompute_[i]);

            //print(recompute_[i], "recompute: ");
        }

        for (size_t i=0; i<this->mgTransfer_.size(); i++)
            std::dynamic_pointer_cast<TruncatedMGTransfer<VectorType> >(mgTransfer_[i])->setRecomputeBitField(&recompute_[i]); */
    }

   void setMatrix(const MatrixType& mat) override {
       Base::setMatrix(mat);

       ignoreHierarchy_.resize(size());
       ignoreHierarchy_.back() = Base::ignore();
       for (int i=mgTransfer_.size()-1; i>=0; i--) {
           std::dynamic_pointer_cast<TruncatedMGTransfer<VectorType> >(mgTransfer_[i])->restrictToFathers(ignoreHierarchy_[i+1], ignoreHierarchy_[i]);
       }

        size_t maxLevel = levelPatchPreconditioners_.size()-1;
        levelPatchPreconditioners_[maxLevel]->setMatrix(mat);

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

            transfer.galerkinRestrictSetOccupation(*levelPatchPreconditioners_[i+1]->getMatrix(), levelMat_[i]);
            transfer.galerkinRestrict(*levelPatchPreconditioners_[i+1]->getMatrix(), levelMat_[i]);
            levelPatchPreconditioners_[i]->setMatrix(levelMat_[i]);

            // Set solution vector sizes for the lower levels
            Dune::MatrixVector::resize(levelX_[i], levelMat_[i]);
            //Dune::MatrixVector::resize(levelRhs_[i], levelMat_[i]);
        }

        // set coarse global problem
        const auto& coarseTransfer = *mgTransfer_[0];
        const auto& fineMat = *levelPatchPreconditioners_[1]->getMatrix();
        coarseTransfer.galerkinRestrictSetOccupation(fineMat, levelMat_[0]);
        coarseTransfer.galerkinRestrict(fineMat, levelMat_[0]);

        Dune::MatrixVector::resize(levelX_[0], levelMat_[0]);
   }

    void iterate() {
        size_t maxLevel = levelPatchPreconditioners_.size()-1;
        levelX_[maxLevel] = *this->getIterate();
        levelRhs_[maxLevel] = *Base::rhs_;

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

            transfer.restrict(levelX_[i+1], levelX_[i]);
            transfer.restrict(levelRhs_[i+1], levelRhs_[i]);
        }

        // set coarse global problem
        const auto& coarseTransfer = *mgTransfer_[0];
        coarseTransfer.restrict(levelX_[1], levelX_[0]);
        coarseTransfer.restrict(levelRhs_[1], levelRhs_[0]);


        if (mode_ == additive)
            iterateAdd();
        else
            iterateMult();
    }

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

        VectorType x;

        // solve coarse global problem
        LocalProblem<MatrixType, VectorType> localProblem(levelMat_[0], levelRhs_[0], ignoreHierarchy_[0]);
        Vector newR;
        localProblem.getLocalRhs(levelX_[0], newR);

        coarseSolver_.setProblem(localProblem.getMat(), levelX_[0], newR);
        coarseSolver_.preprocess();
        coarseSolver_.solve();

        mgTransfer_[0]->prolong(levelX_[0], x);

        for (size_t i=1; i<levelPatchPreconditioners_.size()-1; i++) {
            const auto& transfer = *mgTransfer_[i];
            auto& preconditioner = *levelPatchPreconditioners_[i];

            preconditioner.setIgnore(ignoreHierarchy_[i]);
            preconditioner.apply(levelX_[i], levelRhs_[i]);

            x += preconditioner.getSol();

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

            x = newX;
        }

        auto& preconditioner = *levelPatchPreconditioners_.back();

        preconditioner.setIgnore(ignoreHierarchy_.back());
        preconditioner.apply(levelX_.back(), levelRhs_.back());

        x += preconditioner.getSol();

        *(this->x_) = x;
    }


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

        DUNE_THROW(Dune::Exception, "Not implemented!");

        /*
        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() {
        for (size_t i=1; i<levelPatchPreconditioners_.size(); i++) {
            levelPatchPreconditioners_[i]->build();
        }
    }

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

    auto contactNetwork() const {
        return contactNetwork_;
    }

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

#endif