#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#ifdef HAVE_IPOPT
#undef HAVE_IPOPT
#endif

#include <atomic>
#include <cmath>
#include <csignal>
#include <exception>
#include <fstream>
#include <iostream>
#include <iomanip>
#include <memory>

#include <dune/common/bitsetvector.hh>
#include <dune/common/exceptions.hh>
#include <dune/common/fmatrix.hh>
#include <dune/common/function.hh>
#include <dune/common/fvector.hh>
#include <dune/common/parallel/mpihelper.hh>
#include <dune/common/parametertree.hh>
#include <dune/common/parametertreeparser.hh>

#include <dune/grid/common/mcmgmapper.hh>
#include <dune/istl/bcrsmatrix.hh>
#include <dune/istl/bvector.hh>

#include <dune/fufem/boundarypatch.hh>
#include <dune/fufem/formatstring.hh>
#include <dune/fufem/hdf5/file.hh>

#include <dune/solvers/norms/energynorm.hh>


/*
#include <dune/solvers/solvers/loopsolver.hh>
#include <dune/solvers/solvers/solver.hh>*/
#include <dune/tnnmg/problem-classes/convexproblem.hh>

#include <dune/contact/common/deformedcontinuacomplex.hh>
#include <dune/contact/common/couplingpair.hh>
#include <dune/contact/projections/normalprojection.hh>

#include <dune/tectonic/geocoordinate.hh>
#include <dune/tectonic/globalfriction.hh>


#include "assemblers.hh"
#include "gridselector.hh"
#include "explicitgrid.hh"
#include "explicitvectors.hh"

#include "data-structures/enumparser.hh"
#include "data-structures/enums.hh"
#include "data-structures/contactnetwork.hh"
#include "data-structures/matrices.hh"
#include "data-structures/program_state.hh"

#include "factories/threeblocksfactory.hh"

//#include "io/hdf5-levelwriter.hh"
#include "io/hdf5/restart-io.hh"
#include "io/vtk.hh"

#include "multi-body-problem-data/bc.hh"
#include "multi-body-problem-data/mybody.hh"
#include "multi-body-problem-data/grid/mygrids.hh"

#include "spatial-solving/tnnmg/functional.hh"
#include "spatial-solving/preconditioners/supportpatchfactory.hh"
#include "spatial-solving/tnnmg/localbisectionsolver.hh"
#include "spatial-solving/contact/nbodyassembler.hh"
#include "spatial-solving/solverfactory.hh"

#include "time-stepping/adaptivetimestepper.hh"
#include "time-stepping/rate.hh"
#include "time-stepping/state.hh"
#include "time-stepping/updaters.hh"

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

// for getcwd
#include <unistd.h>

//#include <tbb/tbb.h> //TODO multi threading preconditioner?
//#include <pthread.h>

size_t const dims = MY_DIM;

template <class SupportPatchFactory>
void testSuite(const SupportPatchFactory& supportPatchFactory, const size_t bodyID, const int patchDepth = 0) {
/*    const auto& coarseContactNetwork = supportPatchFactory.coarseContactNetwork();
    const auto& gridView = coarseContactNetwork.body(bodyID)->gridView();

    for (const auto& e : elements(gridView)) {

    }

    using Patch = typename SupportPatchFactory::Patch;
    Patch patch0, patch1, patch2, patch3;

    // (const size_t bodyID, const Element& coarseElement, const size_t localVertex, Patch& patchDofs, const int patchDepth = 0)

    // interior patch inside of one body
    supportPatchFactory.build(0, const Element& coarseElement, const size_t localVertex, patch0, patchDepth);

    // patch at friction interface with two bodies in contact
    supportPatchFactory.build(0, const Element& coarseElement, const size_t localVertex, patch1 patchDepth);

    // patch at friction interface with two bodies in contact and with dirichlet boundary
    supportPatchFactory.build(0, const Element& coarseElement, const size_t localVertex, patch2, patchDepth);

    // crosspoint patch, where all 3 bodies are in contact
    supportPatchFactory.build(0, const Element& coarseElement, const size_t localVertex, patch3, patchDepth);*/
}


Dune::ParameterTree getParameters(int argc, char *argv[]) {
  Dune::ParameterTree parset;
  Dune::ParameterTreeParser::readINITree("/home/mi/podlesny/software/dune/dune-tectonic/src/multi-body-problem.cfg", parset);
  Dune::ParameterTreeParser::readINITree(
      Dune::Fufem::formatString("/home/mi/podlesny/software/dune/dune-tectonic/src/multi-body-problem-%dD.cfg", dims), parset);
  Dune::ParameterTreeParser::readOptions(argc, argv, parset);
  return parset;
}

static std::atomic<bool> terminationRequested(false);
void handleSignal(int signum) { terminationRequested = true; }

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

    char buffer[256];
    char *val = getcwd(buffer, sizeof(buffer));
    if (val) {
        std::cout << buffer << std::endl;
        std::cout << argv[0] << std::endl;
    }

    std::ofstream out("../log.txt");
    std::streambuf *coutbuf = std::cout.rdbuf(); //save old buffer
    std::cout.rdbuf(out.rdbuf()); //redirect std::cout to log.txt

    auto const parset = getParameters(argc, argv);

    using Assembler = MyAssembler<DefLeafGridView, dims>;
    using field_type = Matrix::field_type;


    // ----------------------
    // set up contact network
    // ----------------------
    ThreeBlocksFactory<Grid, Vector> threeBlocksFactory(parset);
    using ContactNetwork = typename ThreeBlocksFactory<Grid, Vector>::ContactNetwork;
    threeBlocksFactory.build();

    ContactNetwork& contactNetwork = threeBlocksFactory.contactNetwork();
    const size_t bodyCount = contactNetwork.nBodies();

    for (size_t i=0; i<bodyCount; i++) {
       // printDofLocation(contactNetwork.body(i)->gridView());
     /* Vector def(contactNetwork.deformedGrids()[i]->size(dims));
      def = 1;
      deformedGridComplex.setDeformation(def, i);*/

        /*auto& levelViews = contactNetwork.levelViews(i);

        for (size_t j=0; j<levelViews.size(); j++) {
            writeToVTK(*levelViews[j], "", "body_" + std::to_string(i) + "_level_" + std::to_string(j));
        }

        writeToVTK(contactNetwork.leafView(i), "", "body_" + std::to_string(i) + "_leaf"); */
    }

    // ----------------------------
    // assemble contactNetwork
    // ----------------------------
    contactNetwork.assemble();

    const auto & coarseContactNetwork = *contactNetwork.level(0);
    const auto & fineContactNetwork = *contactNetwork.level(1);
    SupportPatchFactory<typename ContactNetwork::LevelContactNetwork> supportPatchFactory(coarseContactNetwork, fineContactNetwork);
    size_t bodyID = 2;
    size_t patchDepth = 0;

    std::cout << std::endl;

    // print coarse dofs
    for (size_t i=0; i<bodyCount; i++) {
        std::cout << "Coarse dofs body " << i << std::endl;
        const auto& gv = coarseContactNetwork.body(i)->gridView();
        printDofLocation(gv);

        ScalarVector dofs(gv.size(dims));
        for (size_t j=0; j<dofs.size(); j++) {
            dofs[j] = j;
        }
        writeToVTK(gv, dofs, "", "body_" + std::to_string(i) + "_coarse");
    }

    // print fine dofs
    for (size_t i=0; i<bodyCount; i++) {
        std::cout << "Fine dofs body " << i << std::endl;
        const auto& gv = fineContactNetwork.body(i)->gridView();
        printDofLocation(gv);

        ScalarVector dofs(gv.size(dims));
        for (size_t j=0; j<dofs.size(); j++) {
            dofs[j] = j;
        }
        writeToVTK(gv, dofs, "", "body_" + std::to_string(i) + "_fine");
    }

    using Patch = typename SupportPatchFactory<typename ContactNetwork::LevelContactNetwork>::Patch;
    Patch patch;

    const auto& gridView = coarseContactNetwork.body(bodyID)->gridView();

    Dune::PQkLocalFiniteElementCache<double, double, dims, 1> cache;
    Dune::BitSetVector<1> vertexVisited(gridView.size(dims));
    vertexVisited.unsetAll();

    for (const auto& e: elements(gridView)) {
        const auto& refElement = Dune::ReferenceElements<double, dims>::general(e.type());

        for (size_t i=0; i<refElement.size(dims); i++) {
            auto localIdx = cache.get(e.type()).localCoefficients().localKey(i).subEntity();
            auto globalIdx = gridView.indexSet().subIndex(e, i, dims);

            if (!vertexVisited[globalIdx][0]) {
                vertexVisited[globalIdx][0] = true;
                supportPatchFactory.build(bodyID, e, i, patch, patchDepth);

                print(patch, "patch:");

                size_t c = 0;
                for (size_t j=0; j<bodyCount; j++) {
                    const auto& gv = fineContactNetwork.body(j)->gridView();

                    ScalarVector patchVec(gv.size(dims));
                    for (size_t l=0; l<patchVec.size(); l++) {
                        if (patch[c++][0]) {
                            patchVec[l][0] = 1;
                        }
                    }

                    print(patchVec, "patchVec");

                    // output patch
                    writeToVTK(gv, patchVec, "", "patch_" + std::to_string(globalIdx) + "_body_" + std::to_string(j));
                }
            }
        }
    }

    std::cout << "Done! :)" << std::endl;
    return 1;

    //printMortarBasis<Vector>(contactNetwork.nBodyAssembler());

    // -----------------
    // init input/output
    // -----------------
    std::vector<size_t> nVertices(bodyCount);
    for (size_t i=0; i<bodyCount; i++) {
        nVertices[i] = contactNetwork.body(i)->nVertices();
    }

    using MyProgramState = ProgramState<Vector, ScalarVector>;
    MyProgramState programState(nVertices);
    auto const firstRestart = parset.get<size_t>("io.restarts.first");
    auto const restartSpacing = parset.get<size_t>("io.restarts.spacing");
    auto const writeRestarts = parset.get<bool>("io.restarts.write");
    auto const writeData = parset.get<bool>("io.data.write");
    bool const handleRestarts = writeRestarts or firstRestart > 0;


    auto dataFile =
        writeData ? std::make_unique<HDF5::File>("output.h5") : nullptr;

    auto restartFile = handleRestarts
                           ? std::make_unique<HDF5::File>(
                                 "restarts.h5",
                                 writeRestarts ? HDF5::Access::READWRITE
                                               : HDF5::Access::READONLY)
                           : nullptr;


    auto restartIO = handleRestarts
                         ? std::make_unique<RestartIO<MyProgramState>>(
                               *restartFile, nVertices)
                         : nullptr;

    if (firstRestart > 0) // automatically adjusts the time and timestep
      restartIO->read(firstRestart, programState);
    else
     programState.setupInitialConditions(parset, contactNetwork);


    auto& nBodyAssembler = contactNetwork.nBodyAssembler();
    for (size_t i=0; i<bodyCount; i++) {
      contactNetwork.body(i)->setDeformation(programState.u[i]);
    }
    nBodyAssembler.assembleTransferOperator();
    nBodyAssembler.assembleObstacle();

    // ------------------------
    // assemble global friction
    // ------------------------
    contactNetwork.assembleFriction(parset.get<Config::FrictionModel>("boundary.friction.frictionModel"), programState.weightedNormalStress);

    auto& globalFriction = contactNetwork.globalFriction();
    globalFriction.updateAlpha(programState.alpha);

    using MyVertexBasis = typename Assembler::VertexBasis;
    using MyCellBasis = typename Assembler::CellBasis;
    std::vector<Vector> vertexCoordinates(bodyCount);
    std::vector<const MyVertexBasis* > vertexBases(bodyCount);
    std::vector<const MyCellBasis* > cellBases(bodyCount);

    for (size_t i=0; i<bodyCount; i++) {
      const auto& body = contactNetwork.body(i);
      vertexBases[i] = &(body->assembler()->vertexBasis);
      cellBases[i] = &(body->assembler()->cellBasis);

      auto& vertexCoords = vertexCoordinates[i];
      vertexCoords.resize(nVertices[i]);

      Dune::MultipleCodimMultipleGeomTypeMapper<
          DefLeafGridView, Dune::MCMGVertexLayout> const vertexMapper(body->gridView(), Dune::mcmgVertexLayout());
      for (auto &&v : vertices(body->gridView()))
        vertexCoords[vertexMapper.index(v)] = geoToPoint(v.geometry());
    }

    //typename contactNetwork::BoundaryPatches frictionBoundaries;
    //contactNetwork.boundaryPatches("friction", frictionBoundaries);

    /*
    auto dataWriter =
        writeData ? std::make_unique<
                        HDF5Writer<MyProgramState, MyVertexBasis, DefLeafGridView>>(
                        *dataFile, vertexCoordinates, vertexBases,
                        frictionBoundaries) //, weakPatches)
                  : nullptr;*/

    const MyVTKWriter<MyVertexBasis, MyCellBasis> vtkWriter(cellBases, vertexBases, "/storage/mi/podlesny/software/dune/dune-tectonic/body");

    IterationRegister iterationCount;

    auto const report = [&](bool initial = false) {
      /*if (writeData) {
        dataWriter->reportSolution(programState, contactNetwork.globalFriction());
        if (!initial)
          dataWriter->reportIterations(programState, iterationCount);
        dataFile->flush();
      }

      if (writeRestarts and !initial and
          programState.timeStep % restartSpacing == 0) {
        restartIO->write(programState);
        restartFile->flush();
      }*/

      if (parset.get<bool>("io.printProgress"))
        std::cout << "timeStep = " << std::setw(6) << programState.timeStep
                  << ", time = " << std::setw(12) << programState.relativeTime
                  << ", tau = " << std::setw(12) << programState.relativeTau
                  << std::endl;

      if (parset.get<bool>("io.vtk.write")) {
        std::vector<ScalarVector> stress(bodyCount);

        for (size_t i=0; i<bodyCount; i++) {
          const auto& body = contactNetwork.body(i);
          body->assembler()->assembleVonMisesStress(body->data()->getYoungModulus(),
                                           body->data()->getPoissonRatio(),
                                           programState.u[i], stress[i]);

        }

        vtkWriter.write(programState.timeStep, programState.u, programState.v,
                        programState.alpha, stress);
      }
    };
    report(true);

    // -------------------
    // Set up TNNMG solver
    // -------------------

    BitVector totalDirichletNodes;
    contactNetwork.totalNodes("dirichlet", totalDirichletNodes);

    /*for (size_t i=0; i<totalDirichletNodes.size(); i++) {
        bool val = false;
        for (size_t d=0; d<dims; d++) {
            val = val || totalDirichletNodes[i][d];
        }

        totalDirichletNodes[i] = val;
        for (size_t d=0; d<dims; d++) {
            totalDirichletNodes[i][d] = val;
        }
    }*/

    print(totalDirichletNodes, "totalDirichletNodes:");

    using Functional = Functional<Matrix&, Vector&, GlobalFriction<Matrix, Vector>&, Vector&, Vector&, field_type>;
    using NonlinearFactory = SolverFactory<Functional, BitVector>;

    using BoundaryFunctions = typename ContactNetwork::BoundaryFunctions;
    using BoundaryNodes = typename ContactNetwork::BoundaryNodes;
    using Updaters = Updaters<RateUpdater<Vector, Matrix, BoundaryFunctions, BoundaryNodes>,
                               StateUpdater<ScalarVector, Vector>>;

    BoundaryFunctions velocityDirichletFunctions;
    contactNetwork.boundaryFunctions("dirichlet", velocityDirichletFunctions);

    BoundaryNodes dirichletNodes;
    contactNetwork.boundaryNodes("dirichlet", dirichletNodes);

    for (size_t i=0; i<dirichletNodes.size(); i++) {
        for (size_t j=0; j<dirichletNodes[i].size(); j++) {
        print(*dirichletNodes[i][j], "dirichletNodes_body_" + std::to_string(i) + "_boundary_" + std::to_string(j));
        }
    }

    std::vector<const Dune::BitSetVector<1>*> frictionNodes;
    contactNetwork.frictionNodes(frictionNodes);

    for (size_t i=0; i<frictionNodes.size(); i++) {
        print(*frictionNodes[i], "frictionNodes_body_" + std::to_string(i));
    }

    Updaters current(
        initRateUpdater(
            parset.get<Config::scheme>("timeSteps.scheme"),
            velocityDirichletFunctions,
            dirichletNodes,
            contactNetwork.matrices(),
            programState.u,
            programState.v,
            programState.a),
        initStateUpdater<ScalarVector, Vector>(
            parset.get<Config::stateModel>("boundary.friction.stateModel"),
            programState.alpha,
            nBodyAssembler.getContactCouplings(),
            contactNetwork.couplings())
            );


    auto const refinementTolerance = parset.get<double>("timeSteps.refinementTolerance");

    const auto& stateEnergyNorms = contactNetwork.stateEnergyNorms();

    auto const mustRefine = [&](Updaters &coarseUpdater,
                                Updaters &fineUpdater) {

        //return false;
      std::cout << "Step 1" << std::endl;

      std::vector<ScalarVector> coarseAlpha;
      coarseAlpha.resize(bodyCount);
      coarseUpdater.state_->extractAlpha(coarseAlpha);

      print(coarseAlpha, "coarseAlpha:");

      std::vector<ScalarVector> fineAlpha;
      fineAlpha.resize(bodyCount);
      fineUpdater.state_->extractAlpha(fineAlpha);

      print(fineAlpha, "fineAlpha:");

      std::cout << "Step 3" << std::endl;

      ScalarVector::field_type energyNorm = 0;
      for (size_t i=0; i<stateEnergyNorms.size(); i++) {
          std::cout << "for " << i << std::endl;

          std::cout << not stateEnergyNorms[i] << std::endl;

          if (coarseAlpha[i].size()==0 || fineAlpha[i].size()==0)
              continue;

          energyNorm += stateEnergyNorms[i]->diff(fineAlpha[i], coarseAlpha[i]);
      }
      std::cout << "energy norm: " << energyNorm << " tol: " << refinementTolerance <<  std::endl;
      std::cout << "must refine: " << (energyNorm > refinementTolerance) <<  std::endl;
      return energyNorm > refinementTolerance;
    };


    std::signal(SIGXCPU, handleSignal);
    std::signal(SIGINT, handleSignal);
    std::signal(SIGTERM, handleSignal);

    typename ContactNetwork::ExternalForces externalForces;
    contactNetwork.externalForces(externalForces);

    AdaptiveTimeStepper<NonlinearFactory, std::decay_t<decltype(nBodyAssembler)>, Updaters, std::decay_t<decltype(stateEnergyNorms)>>
        adaptiveTimeStepper(parset, nBodyAssembler, totalDirichletNodes, globalFriction, frictionNodes, current,
                            programState.relativeTime, programState.relativeTau,
                            externalForces, stateEnergyNorms, mustRefine);

    while (!adaptiveTimeStepper.reachedEnd()) {
      programState.timeStep++;

      iterationCount = adaptiveTimeStepper.advance();

      programState.relativeTime = adaptiveTimeStepper.relativeTime_;
      programState.relativeTau = adaptiveTimeStepper.relativeTau_;
      current.rate_->extractDisplacement(programState.u);
      current.rate_->extractVelocity(programState.v);
      current.rate_->extractAcceleration(programState.a);
      current.state_->extractAlpha(programState.alpha);

      print(programState.u, "current u:");
      print(programState.v, "current v:");
      print(programState.a, "current a:");
      print(programState.alpha, "current alpha:");

      contactNetwork.setDeformation(programState.u);

      report();

      if (programState.timeStep==50) {
        std::cout << "limit of timeSteps reached!" << std::endl;
        break; // TODO remove after debugging
      }

      if (terminationRequested) {
        std::cerr << "Terminating prematurely" << std::endl;
        break;
      }
    }

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

  } catch (Dune::Exception &e) {
    Dune::derr << "Dune reported error: " << e << std::endl;
  } catch (std::exception &e) {
    std::cerr << "Standard exception: " << e.what() << std::endl;
  }
}