#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
*/

#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/bcrsmatrix.hh>
#include <dune/istl/matrix.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/matrixreader.hh>
#include <dune/faultnetworks/utils/vectorreader.hh>
#include <dune/faultnetworks/utils/levelinterfacenetworkwriter.hh>
#include <dune/faultnetworks/utils/debugutils.hh>

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

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

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

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

#include <dune/faultnetworks/dgmgtransfer.hh>
#include <dune/faultnetworks/compressedmultigridtransfer.hh>

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

#include <dune/faultnetworks/faultfactories/multileveluniformfaultfactory.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("/home/mi/podlesny/software/dune/faultnetworks/src/uniformfaultnetworks/multilevelpatches.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 faultCount = problemParameters.get<int>("faultCount");

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

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

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


    std::ofstream out(resultPath + oscDataFile + "_" + std::to_string(coarseResolution) + "_" + std::to_string(fineResolution) + "_" + std::to_string(exactResolution) + "_" + std::to_string(faultCount) + ".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::round(std::log2(oscGridN)) - coarseResolution;

    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;

    std::vector<int> faultToLevel(faultCount);
    for (size_t i=0; i<faultToLevel.size(); ) {
        for (int j=fineLevelIdx; j>0 && i<faultToLevel.size(); j--) {
            faultToLevel[i] = j;
            i++;
        }

        for (int j=0; j<fineLevelIdx && i<faultToLevel.size(); j++) {
            faultToLevel[i] = j;
            i++;
        }
    }

    print(faultToLevel, "faultToLevel: ");

    MultilevelUniformFaultFactory<GridType> faultFactory(faultCount, faultToLevel, coarseResolution, exactLevelIdx);
    faultFactory.prolongToAll();

    const GridType& grid = faultFactory.grid();
    const InterfaceNetwork<GridType>& interfaceNetwork = faultFactory.interfaceNetwork();

    int writeLevel = fineLevelIdx;
    LevelInterfaceNetworkWriter networkWriter(resultPath + "levelinterfacenetwork_" + std::to_string(writeLevel) + ".tikz");
    networkWriter.write(interfaceNetwork.levelInterfaceNetwork(writeLevel));

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


    // ----------------------------------------------------------------
    // ---              compute initial iterate                     ---
    // ----------------------------------------------------------------
    std::cout << std::endl;
    std::cout << "Computing initial iterate!" << std::endl;

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

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

    std::vector<std::shared_ptr<LocalInterfaceAssembler>> coarseLocalInterfaceAssemblers(coarseLevelInterfaceNetwork.size());
    for (size_t i=0; i<coarseLocalInterfaceAssemblers.size(); i++) {
        //const double pen = penaltyFactor;
        const double pen = penaltyFactor*std::pow(2.0, coarseResolution);
        coarseLocalInterfaceAssemblers[i] = std::make_shared<LocalInterfaceAssembler>(B, 2, pen);
    }

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

    const VectorType& coarseSol = coarseGlobalAssembler.getSol();
    const DGBasis& coarseBasis = coarseGlobalAssembler.basis();

    // ----------------------------------------------------------------
    // ---              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 double pen = penaltyFactor;
        const double pen = penaltyFactor*std::pow(2.0, coarseResolution + exactLevelInterfaceNetwork.getIntersectionsLevels().at(i));
        exactLocalInterfaceAssemblers[i] = std::make_shared<LocalInterfaceAssembler>(B, 2, pen);
    }

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

    const VectorType& exactSol = exactGlobalAssembler.getSol();
    const 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 double pen = penaltyFactor;
        const double pen = penaltyFactor*std::pow(2.0, coarseResolution + fineLevelInterfaceNetwork.getIntersectionsLevels().at(i));
        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();
    const DGBasis& fineBasis = fineGlobalAssembler.basis();

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

    // ---------------------------------------------------------------
    // multilevel transfer
    // ---------------------------------------------------------------
    DGMGTransfer<DGBasis> coarseFineTransfer(coarseBasis, fineBasis);

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

    BitVector activeLevels(fineLevelIdx+1, true);
    //activeLevels[0][0] = true;
    //activeLevels[1][0] = false;
    //activeLevels[2][0] = false;
    //activeLevels[fineLevelIdx][0] = 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 double pen = penaltyFactor;
            const double pen = penaltyFactor*std::pow(2.0, coarseResolution + levelInterfaceNetwork.getIntersectionsLevels().at(j));
            levelLocalIntersectionAssemblers[j] = std::make_shared<LocalInterfaceAssembler>(B, 2, pen);
        }
    }

    typedef MultilevelPatchPreconditioner<DGBasis, LocalOscAssembler, LocalInterfaceAssembler, MatrixType, VectorType> MultilevelPatchPreconditioner;
    typedef typename MultilevelPatchPreconditioner::LevelPatchPreconditionerType LevelPatchPreconditioner;
    //typedef typename MultilevelPatchPreconditioner::LevelGlobalPreconditionerType LevelGlobalPreconditioner;
    MultilevelPatchPreconditioner::Mode preconditionerMode = MultilevelPatchPreconditioner::Mode::additive;

    MultilevelPatchPreconditioner preconditioner(interfaceNetwork, activeLevels, localAssemblers, localIntersectionAssemblers, preconditionerMode);
    for (size_t i=0; i<preconditioner.size()-1; i++) {
        LevelPatchPreconditioner& 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;
    
    /*Dune::BlockVector<Dune::FieldVector<double,1>> initialError(fineX.size());
    initialError = initialX;
    initialError -= fineX;
    std::cout << "Initial error in energy norm: " << energyNorm.operator()(initialError) << std::endl;
    std::cout << "Initial error in two norm: " << initialError.two_norm() << std::endl << std::endl;*/


    VectorType initialX;
    coarseFineTransfer.prolong(coarseSol, initialX);

    VectorType initialXCopy(initialX);
    VectorType fineRhsCopy(fineRhs);
    VectorType rhs(fineRhs);

    //print(initialX, "initialX");
    //print(fineRhs, "fineRhs");
    //print(exactSol, "exactSol");
    //print(fineSol, "fineSol");
    
    /* Basis fineContBasis(fineGridView);
    std::cout << "fineContBasis.size(): " << fineContBasis.size() << std::endl;
    DGMGTransfer<Basis> contMGTransfer(coarseBasis, fineContBasis);

    Basis exactContBasis(exactGridView);
    std::cout << "exactContBasis.size(): " << exactContBasis.size() << std::endl; */

    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.check();
    solver.preprocess();
    solver.solve();


    /*
    VectorType finalError(fineSol.size());
    finalError = initialX;
    finalError -= fineSol;
    std::cout << "Final error in energy norm: " << fineEnergyNorm.operator()(finalError) << std::endl;
    std::cout << "Final error in two norm: " << finalError.two_norm() << std::endl << std::endl;*/

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

    // regular unpreconditioned cg solver for comparisson
/*
    std::cout << "Setup complete, starting cg iteration!" << std::endl;
    std::cout << std::endl;
    
    // solve
    OscCGSolver<MatrixType, VectorType, DGBasis >
    regSolver(fineBasis, &fineGlobalAssembler.matrix(), &initialXCopy, &fineRhsCopy, nullptr,
                    maxIterations, solverTolerance, &exactEnergyNorm,
                    Solver::FULL, &exactSol, &fineSol, discMGTransfer, 1.0, true);
    regSolver.check();
    regSolver.preprocess();
    regSolver.solve();
*/
    // TODO: use once debugging finished
    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;

}