#include <config.h>

#include <cmath>
#include <iostream>
#include <sstream>
#include <fstream>
#include <string>

/*
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-parameter"
#include <dune/grid/uggrid/ugwrapper.hh>
#pragma GCC diagnostic pop
*/

// dune-common includes
#include <dune/common/fvector.hh>
#include <dune/common/fmatrix.hh>
#include <dune/common/shared_ptr.hh>
#include <dune/common/parametertree.hh>
#include <dune/common/parametertreeparser.hh>
#include <dune/common/timer.hh>

// dune-istl includes
#include <dune/istl/matrix.hh>
#include <dune/istl/bcrsmatrix.hh>

// dune-fufem includes
#include <dune/fufem/assemblers/localassemblers/l2functionalassembler.hh>

// dune-grid includes
#include <dune/grid/uggrid.hh>
#include <dune/grid/io/file/vtk/vtkwriter.hh>

// dune-solver includes
#include <dune/solvers/norms/energynorm.hh>
#include <dune/solvers/norms/twonorm.hh>

// dune-faultnetworks includes
#include <dune/faultnetworks/utils/debugutils.hh>
#include <dune/faultnetworks/utils/matrixreader.hh>
#include <dune/faultnetworks/utils/vectorreader.hh>
#include <dune/faultnetworks/utils/levelinterfacenetworkwriter.hh>

#include <dune/faultnetworks/preconditioners/multilevelpatchpreconditioner.hh>
#include <dune/faultnetworks/solvers/osccgsolver.hh>

#include <dune/faultnetworks/assemblers/intersectionjumpmassassembler.hh>
#include <dune/faultnetworks/assemblers/osclocalassembler.hh>

#include <dune/faultnetworks/oscrhs.hh>
#include <dune/faultnetworks/oscdata.hh>
#include <dune/faultnetworks/oscdatahierarchy.hh>

#include <dune/faultnetworks/faultp1nodalbasis.hh>
#include <dune/faultnetworks/interfacenetwork.hh>
#include <dune/faultnetworks/levelinterfacenetwork.hh>
#include <dune/faultnetworks/levelinterfacenetworkproblem.hh>
#include <dune/faultnetworks/dgmgtransfer.hh>

#include <dune/faultnetworks/faultfactories/cantorfaultfactory.hh>

#include <dune/solvers/solvers/loopsolver.hh>

int main(int argc, char** argv) { try
{
    MPIHelper::instance(argc, argv);

    const int dim = 2;
    //typedef double ctype;


    typedef typename Dune::FieldVector<double, dim> WorldVectorType;
    typedef typename Dune::BlockVector<Dune::FieldVector<double,1>> VectorType;
    typedef typename Dune::BitSetVector<1> BitVector;
    typedef typename Dune::BCRSMatrix<Dune::FieldMatrix<double, 1, 1>> MatrixType;


    #if HAVE_UG
    typedef typename Dune::UGGrid<dim> GridType;
    #else
    #error No UG!
    #endif

    typedef typename GridType::LevelGridView GV;
    typedef FaultP1NodalBasis<GV, double> DGBasis;
    typedef typename DGBasis::LocalFiniteElement DGFE;

    typedef OscLocalAssembler<GridType, DGFE, DGFE> LocalOscAssembler;
    typedef IntersectionJumpMassAssembler<GV, DGFE, DGFE> LocalInterfaceAssembler;

  // parse parameter file
    ParameterTree parameterSet;
    if (argc==2)
        ParameterTreeParser::readINITree(argv[1], parameterSet);
    else
        ParameterTreeParser::readINITree("cantorfaultnetwork.parset", parameterSet);

    const std::string path = parameterSet.get<std::string>("path");
    const std::string resultPath = parameterSet.get<std::string>("resultPath");

    int problemCount = 0;
    while (parameterSet.hasSub("problem" + std::to_string(problemCount))) {
	
    ParameterTree& problemParameters = parameterSet.sub("problem" + std::to_string(problemCount));

    const std::string oscDataFile = problemParameters.get<std::string>("oscDataFile");

    const int coarseResolution = problemParameters.get<int>("coarseResolution");
    const int fineResolution = problemParameters.get<int>("fineResolution");
    const int exactResolution = problemParameters.get<int>("exactResolution");
    const int minCantorResolution = problemParameters.get<int>("minCantorResolution");

    const double c = problemParameters.get<double>("penaltyFactor");

    const int maxIterations = problemParameters.get<int>("maxIterations");
    const double solverTolerance = problemParameters.get<double>("tol");
    const bool nestedIteration = problemParameters.get<bool>("nestedIteration");

    const int fineGridN = std::pow(2, fineResolution);

    std::ofstream out(resultPath + oscDataFile + "_" + std::to_string(coarseResolution) + "_" + std::to_string(fineResolution) + "_" + std::to_string(exactResolution) + ".log");
    std::streambuf *coutbuf = std::cout.rdbuf(); //save old buf
    std::cout.rdbuf(out.rdbuf()); //redirect std::cout to out.txt


    std::cout << std::endl;
    std::cout << "Parameter file read successfully!" << std::endl;

    typedef Dune::Matrix<Dune::FieldMatrix<double,1,1>, std::allocator<Dune::FieldMatrix<double,1,1>>> OscMatrixType;
    OscMatrixType matrix;

    // interface integral jump penalty, B:
    OscMatrixType B(2,2);
    B[0][0] = 1;
    B[1][1] = 1;
    B[0][1] = -1;
    B[1][0] = -1;

    MatrixReader<OscMatrixType> matrixReader(path + oscDataFile);
    matrixReader.read(matrix);

    const int oscGridN = matrix.N();
    const int oscGridLevelIdx = std::max(std::round(std::log2(oscGridN)) - coarseResolution, 0.0);

    if (oscGridN>fineGridN)
        DUNE_THROW(Dune::Exception, "Provided oscData too large!");

    std::cout << "OscData file read successfully!" << std::endl << std::endl;

    const int fineLevelIdx = fineResolution - coarseResolution;
    const int exactLevelIdx = exactResolution - coarseResolution;

    size_t maxCantorLevel = exactLevelIdx+minCantorResolution;
    std::map<size_t, size_t> levelToCantorLevel;
    for (int i=0; i<=exactLevelIdx; i++) {
        levelToCantorLevel[i] = i+minCantorResolution;
    }

    // build multilevel cantor fault network
    CantorFaultFactory<GridType> faultFactory(levelToCantorLevel, coarseResolution, exactLevelIdx, maxCantorLevel);
    const InterfaceNetwork<GridType>& interfaceNetwork = faultFactory.interfaceNetwork();

    /*if (problemCount==0) {
        std::vector<int> writeLevels {1, 2, 8};
        for (size_t i=0; i<writeLevels.size(); i++) {
            int writeLevel = writeLevels[i];
            bool writeGrid = !(writeLevel==8);

            LevelInterfaceNetworkWriter networkWriter(resultPath + "levelinterfacenetwork_" + std::to_string(writeLevel) + ".tikz");
            networkWriter.write(interfaceNetwork.levelInterfaceNetwork(writeLevel), writeGrid);
        }
    }*/

    const GridType& grid = faultFactory.grid();

    OscData<GridType, OscMatrixType> oscData(grid, oscGridLevelIdx);
    oscData.set(matrix);
    OscDataHierarchy<GridType, OscMatrixType> oscDataHierarchy(oscData);
    
    Timer totalTimer;
    totalTimer.reset();
    totalTimer.start();


    // ----------------------------------------------------------------
    // ---              compute initial iterate                     ---
    // ----------------------------------------------------------------
    VectorType coarseSol;
    std::shared_ptr<DGBasis> coarseBasis;

    OscRhs<WorldVectorType, Dune::FieldVector<double,1>> f;
    L2FunctionalAssembler<GridType, DGFE> l2FunctionalAssembler(f);

    std::cout << std::endl;
    std::cout << "Computing initial iterate!" << std::endl;

    if (!nestedIteration) {
        const LevelInterfaceNetwork<GV>& coarseLevelInterfaceNetwork = interfaceNetwork.levelInterfaceNetwork(0);
        LevelInterfaceNetworkProblem<DGBasis, MatrixType, VectorType> coarseGlobalAssembler(coarseLevelInterfaceNetwork);
        LocalOscAssembler coarseLocalAssembler(oscDataHierarchy[0]->data(), oscDataHierarchy[0]->mapper());

        std::vector<std::shared_ptr<LocalInterfaceAssembler>> coarseLocalInterfaceAssemblers(coarseLevelInterfaceNetwork.size());
        for (size_t i=0; i<coarseLocalInterfaceAssemblers.size(); i++) {
            const int k = minCantorResolution + coarseLevelInterfaceNetwork.getIntersectionsLevels().at(i);
            const double pen = (1+1.0/c) * std::pow(1.0+c, k) * std::pow(2.0, k-1);
            coarseLocalInterfaceAssemblers[i] = std::make_shared<LocalInterfaceAssembler>(B, 2, pen);
        }

        coarseGlobalAssembler.assemble(coarseLocalAssembler, coarseLocalInterfaceAssemblers, l2FunctionalAssembler);
        coarseGlobalAssembler.solve();

        coarseSol = coarseGlobalAssembler.getSol();
        coarseBasis = coarseGlobalAssembler.basis();
    } else {
        const LevelInterfaceNetwork<GV>& nLevelInterfaceNetwork = interfaceNetwork.levelInterfaceNetwork(fineLevelIdx-1);
        coarseBasis = std::make_shared<DGBasis>(nLevelInterfaceNetwork);

        coarseSol.resize(coarseBasis->size());

        std::stringstream solFilename;
        solFilename << resultPath << "solutions/level_" << std::to_string(fineResolution-1);

        std::ifstream file(solFilename.str().c_str(), std::ios::in|std::ios::binary);
        if (not(file))
            DUNE_THROW(SolverError, "Couldn't open file " << solFilename.str() << " for reading");

        Dune::MatrixVector::Generic::readBinary(file, coarseSol);

        file.close();
    }

    // ----------------------------------------------------------------
    // ---              compute exact solution                      ---
    // ----------------------------------------------------------------
    std::cout << std::endl;
    std::cout << "Computing exact solution!" << std::endl;

    const LevelInterfaceNetwork<GV>& exactLevelInterfaceNetwork = interfaceNetwork.levelInterfaceNetwork(exactLevelIdx);
    LevelInterfaceNetworkProblem<DGBasis, MatrixType, VectorType> exactGlobalAssembler(exactLevelInterfaceNetwork);
    LocalOscAssembler exactLocalAssembler(oscDataHierarchy[exactLevelIdx]->data(), oscDataHierarchy[exactLevelIdx]->mapper());

    std::vector<std::shared_ptr<LocalInterfaceAssembler>> exactLocalInterfaceAssemblers(exactLevelInterfaceNetwork.size());
    for (size_t i=0; i<exactLocalInterfaceAssemblers.size(); i++) {
        const int k = minCantorResolution + exactLevelInterfaceNetwork.getIntersectionsLevels().at(i);
        const double pen = (1+1.0/c) * std::pow(1.0+c, k) * std::pow(2.0, k-1);
        exactLocalInterfaceAssemblers[i] = std::make_shared<LocalInterfaceAssembler>(B, 2, pen);
    }

    exactGlobalAssembler.assemble(exactLocalAssembler, exactLocalInterfaceAssemblers, l2FunctionalAssembler);
    exactGlobalAssembler.solve();

    const VectorType& exactSol = exactGlobalAssembler.getSol();
    std::shared_ptr<DGBasis> exactBasis = exactGlobalAssembler.basis();

    EnergyNorm<MatrixType, VectorType> exactEnergyNorm(exactGlobalAssembler.matrix());

    // ----------------------------------------------------------------
    // ---                set up cg solver                          ---
    // ----------------------------------------------------------------
    std::cout << std::endl;
    std::cout << "---------------------------------------------" << std::endl; 
    std::cout << "Setting up the fine level CG solver:" << std::endl; 
    std::cout << "---------------------------------------------" << std::endl << std::endl;

    const LevelInterfaceNetwork<GV>& fineLevelInterfaceNetwork = interfaceNetwork.levelInterfaceNetwork(fineLevelIdx);

    LevelInterfaceNetworkProblem<DGBasis, MatrixType, VectorType> fineGlobalAssembler(fineLevelInterfaceNetwork);
    LocalOscAssembler fineLocalAssembler(oscDataHierarchy[fineLevelIdx]->data(), oscDataHierarchy[fineLevelIdx]->mapper());

    std::vector<std::shared_ptr<LocalInterfaceAssembler>> fineLocalInterfaceAssemblers(fineLevelInterfaceNetwork.size());
    for (size_t i=0; i<fineLocalInterfaceAssemblers.size(); i++) {
        const int k = minCantorResolution + fineLevelInterfaceNetwork.getIntersectionsLevels().at(i);
        const double pen = (1+1.0/c) * std::pow(1.0+c, k) * std::pow(2.0, k-1);
        fineLocalInterfaceAssemblers[i] = std::make_shared<LocalInterfaceAssembler>(B, 2, pen);
    }

    fineGlobalAssembler.assemble(fineLocalAssembler, fineLocalInterfaceAssemblers, l2FunctionalAssembler);
    fineGlobalAssembler.solve();

    const VectorType& fineSol = fineGlobalAssembler.getSol();
    const VectorType& fineRhs = fineGlobalAssembler.rhs();
    std::shared_ptr<DGBasis> fineBasis = fineGlobalAssembler.basis();

    EnergyNorm<MatrixType, VectorType> fineEnergyNorm(fineGlobalAssembler.matrix());

    std::cout << "Setting up the preconditioner!" << std::endl;

    BitVector activeLevels(fineLevelIdx+1, true);

    std::vector<std::shared_ptr<LocalOscAssembler>> localAssemblers;
    std::vector<std::vector<std::shared_ptr<LocalInterfaceAssembler>>> localIntersectionAssemblers;

    localAssemblers.resize(activeLevels.size());
    localIntersectionAssemblers.resize(activeLevels.size());

    for (size_t i=0; i<activeLevels.size(); i++) {
        localAssemblers[i] = std::make_shared<LocalOscAssembler>(oscDataHierarchy[i]->data(), oscDataHierarchy[i]->mapper());

        const LevelInterfaceNetwork<GV>& levelInterfaceNetwork = interfaceNetwork.levelInterfaceNetwork(i);
        std::vector<std::shared_ptr<LocalInterfaceAssembler>>& levelLocalIntersectionAssemblers = localIntersectionAssemblers[i];
        levelLocalIntersectionAssemblers.resize(levelInterfaceNetwork.size());
        for (size_t j=0; j<levelLocalIntersectionAssemblers.size(); j++) {
            const int k = minCantorResolution + levelInterfaceNetwork.getIntersectionsLevels().at(j);
            const double pen = (1+1.0/c) * std::pow(1.0+c, k) * std::pow(2.0, k-1);
            levelLocalIntersectionAssemblers[j] = std::make_shared<LocalInterfaceAssembler>(B, 2, pen);
        }
    }

    const auto& preconditionerParset = parameterSet.sub("preconditioner");
    using MultilevelPatchPreconditioner = MultilevelPatchPreconditioner<DGBasis, LocalOscAssembler, LocalInterfaceAssembler, MatrixType, VectorType>;

    MultilevelPatchPreconditioner preconditioner(interfaceNetwork, activeLevels, localAssemblers, localIntersectionAssemblers, preconditionerParset);
    /*for (size_t i=0; i<preconditioner.size()-1; i++) {
        auto& levelFaultPreconditioner = *preconditioner.levelPatchPreconditioner(i);

        levelFaultPreconditioner.setPatchDepth(patchDepth);
        levelFaultPreconditioner.setBoundaryMode(LevelPatchPreconditioner::BoundaryMode::homogeneous);
    }*/
    preconditioner.build();

    std::cout << "Setup complete, starting preconditioned cg iteration!" << std::endl;
    std::cout << std::endl << std::endl;

    // set initial iterate
    VectorType initialX;
    DGMGTransfer<DGBasis> coarseFineTransfer(*coarseBasis, *fineBasis);
    coarseFineTransfer.prolong(coarseSol, initialX);

    VectorType rhs(fineRhs);

    DGMGTransfer<DGBasis> discMGTransfer(*fineBasis, *exactBasis);

    // solve
    OscCGSolver<MatrixType, VectorType, DGBasis>
    solver(*fineBasis, &fineGlobalAssembler.matrix(), &initialX, &rhs, &preconditioner,
                    maxIterations, solverTolerance, &exactEnergyNorm,
                    Solver::FULL, &exactSol, &fineSol, discMGTransfer, 1.0, true); //((oscGridN+0.0)/fineGridN)


    //solver.historyBuffer_ = resultPath + "solutions/level_" + std::to_string(fineResolution) + ".vec";


    /*VectorType x = initialX;
    preconditioner.setProblem(fineGlobalAssembler.matrix(), x, rhs);

    Dune::Solvers::LoopSolver<VectorType>
    solver (preconditioner, maxIterations, solverTolerance,
                   fineEnergyNorm,
                   Solver::FULL,
                   true, &fineSol);*/

    solver.check();
    solver.preprocess();
    solver.solve();

    std::cout << std::endl;
    std::cout << "Total time: " << totalTimer.elapsed() << " seconds." << std::endl; 

    std::cout.rdbuf(coutbuf); //reset to standard output again

    std::cout << "Problem " << problemCount << " done!" << std::endl;
    problemCount++;
    }
} catch (Dune::Exception e) {

    std::cout << e << std::endl;

}
}