#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/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/assemblers/nbodyassembler.hh>
//#include <dune/contact/assemblers/dualmortarcoupling.hh>
#include <dune/contact/projections/normalprojection.hh>

#include <dune/tectonic/geocoordinate.hh>
#include <dune/tectonic/myblockproblem.hh>
#include <dune/tectonic/globalfriction.hh>
#include <dune/fufem/hdf5/file.hh>

#include "assemblers.hh"
#include "diameter.hh"
#include "enumparser.hh"
#include "enums.hh"
#include "gridselector.hh"
#include "hdf5-writer.hh"
#include "hdf5/restart-io.hh"
#include "matrices.hh"
#include "program_state.hh"
#include "multi-body-problem-data/bc.hh"
#include "multi-body-problem-data/mybody.hh"
#include "multi-body-problem-data/cuboidgeometry.hh"
#include "multi-body-problem-data/myglobalfrictiondata.hh"
#include "multi-body-problem-data/mygrids.hh"
#include "multi-body-problem-data/weakpatch.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 "vtk.hh"

// for getcwd
#include <unistd.h>

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

#define USE_OLD_TNNMG

size_t const dims = MY_DIM;

template <class GridView>
void writeToVTK(const GridView& gridView, const std::string path, const std::string name) {
    Dune::VTKWriter<GridView> vtkwriter(gridView);

    //std::ofstream lStream( "garbage.txt" );
    std::streambuf* lBufferOld = std::cout.rdbuf();
    //std::cout.rdbuf(lStream.rdbuf());

    vtkwriter.pwrite(name, path, path);

    std::cout.rdbuf( lBufferOld );
}


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;
    }

    auto const parset = getParameters(argc, argv);

    using Vector = Dune::BlockVector<Dune::FieldVector<double, dims>>;
    using DeformedGridComplex = typename Dune::Contact::DeformedContinuaComplex<Grid, Vector>;
    using DeformedGrid = DeformedGridComplex::DeformedGridType;

    using DefLeafGridView = DeformedGrid::LeafGridView;
    using DefLevelGridView = DeformedGrid::LevelGridView;

    using MyAssembler = MyAssembler<DefLeafGridView, dims>;
    using Matrix = MyAssembler::Matrix;
    using LocalVector = Vector::block_type;
    using ScalarMatrix = MyAssembler::ScalarMatrix;
    using ScalarVector = MyAssembler::ScalarVector;
    using field_type = Matrix::field_type;

    // set up material properties of all bodies
    MyBody<dims> const body(parset);

    // set up cuboid geometries
    const size_t bodyCount = parset.get<int>("problem.bodyCount");
    std::vector<std::shared_ptr<CuboidGeometry>> cuboidGeometries(bodyCount);

    double const length = 1.00;
    double const width = 0.27;
    double const weakLength = 0.20;

#if MY_DIM == 3
    double const depth = 0.60;
    // TODO: replace with make_shared
    cuboidGeometries[0] = std::shared_ptr<CuboidGeometry>(new CuboidGeometry({0,0,0}, {0.2,width,0}, weakLength, length, width, depth));
    for (size_t i=1; i<bodyCount; i++) {
        cuboidGeometries[i] = std::shared_ptr<CuboidGeometry>(new CuboidGeometry(cuboidGeometries[i-1]->D, {0.6,i*width,0}, weakLength, length, width, depth));
    }
#else
    // TODO: replace with make_shared
    cuboidGeometries[0] = std::shared_ptr<CuboidGeometry>(new CuboidGeometry({0,0}, {0.2,width}, weakLength, length, width));
    for (size_t i=1; i<bodyCount; i++) {
        cuboidGeometries[i] = std::shared_ptr<CuboidGeometry>(new CuboidGeometry(cuboidGeometries[i-1]->D, {0.6,i*width}, weakLength, length, width));
    }
#endif

    // set up grids
    GridsConstructor<Grid> gridConstructor(cuboidGeometries);
    auto& grids = gridConstructor.getGrids();

    std::vector<ConvexPolyhedron<LocalVector>> weakPatches(grids.size());

    for (size_t i=0; i<grids.size(); i++) {
        const auto& cuboidGeometry = *cuboidGeometries[i];

        // define weak patch and refine grid
        auto& weakPatch = weakPatches[i];
        weakPatch = getWeakPatch<LocalVector>(parset.sub("boundary.friction.weakening"), cuboidGeometry);
        refine(*grids[i], weakPatch, parset.get<double>("boundary.friction.smallestDiameter"), cuboidGeometry.lengthScale);

        writeToVTK(grids[i]->leafGridView(), "", "grid" + std::to_string(i));

        // determine minDiameter and maxDiameter
        double minDiameter = std::numeric_limits<double>::infinity();
        double maxDiameter = 0.0;
        for (auto &&e : elements(grids[i]->leafGridView())) {
          auto const geometry = e.geometry();
          auto const diam = diameter(geometry);
          minDiameter = std::min(minDiameter, diam);
          maxDiameter = std::max(maxDiameter, diam);
        }
        std::cout << "Grid" << i << " min diameter: " << minDiameter << std::endl;
        std::cout << "Grid" << i << " max diameter: " << maxDiameter << std::endl;
    }

    // set up deformed grids
    DeformedGridComplex deformedGridComplex;
    for (size_t i=0; i<grids.size(); i++) {
        deformedGridComplex.addGrid(*grids[i], "body" + std::to_string(i));
    }
    const auto deformedGrids = deformedGridComplex.gridVector();

    std::vector<std::shared_ptr<DefLeafGridView>> leafViews(deformedGrids.size());
    std::vector<std::shared_ptr<DefLevelGridView>> levelViews(deformedGrids.size());
    std::vector<int> leafVertexCounts(deformedGrids.size());

    using LeafFaces = MyFaces<DefLeafGridView>;
    using LevelFaces = MyFaces<DefLevelGridView>;
    std::vector<std::shared_ptr<LeafFaces>> leafFaces(deformedGrids.size());
    std::vector<std::shared_ptr<LevelFaces>> levelFaces(deformedGrids.size());

    for (size_t i=0; i<deformedGrids.size(); i++) {
        leafViews[i] = std::make_shared<DefLeafGridView>(deformedGrids[i]->leafGridView());
        levelViews[i] = std::make_shared<DefLevelGridView>(deformedGrids[i]->levelGridView(0));

        leafVertexCounts[i] = leafViews[i]->size(dims);

        const auto& cuboidGeometry = *cuboidGeometries[i];
        leafFaces[i] = std::make_shared<LeafFaces>(*leafViews[i], cuboidGeometry);
        levelFaces[i] = std::make_shared<LevelFaces>(*levelViews[i], cuboidGeometry);
    }

    // set up contact couplings
    using LeafBoundaryPatch = BoundaryPatch<DefLeafGridView>;
    using LevelBoundaryPatch = BoundaryPatch<DefLevelGridView>;

    using CouplingPair = Dune::Contact::CouplingPair<DeformedGrid, DeformedGrid, field_type>;
    std::vector<std::shared_ptr<CouplingPair>> couplings(bodyCount-1);

    for (size_t i=0; i<couplings.size(); i++) {
      auto nonmortarPatch = std::make_shared<LevelBoundaryPatch>(levelFaces[i]->upper);
      auto mortarPatch = std::make_shared<LevelBoundaryPatch>(levelFaces[i+1]->lower);

      auto contactProjection = std::make_shared<Dune::Contact::NormalProjection<LeafBoundaryPatch>>();
      std::shared_ptr<CouplingPair::BackEndType> backend = nullptr;
      couplings[i]->set(i, i+1, nonmortarPatch, mortarPatch, 0.1, Dune::Contact::CouplingPairBase::STICK_SLIP, contactProjection, backend);
    }

    // set up dune-contact nBodyAssembler
    /*std::vector<const DeformedGrid*> const_grids(bodyCount);
    for (size_t i=0; i<const_grids.size(); i++) {
        const_grids[i] = &deformedGrids.grid("body" + std::to_string(i));
    }*/

    Dune::Contact::NBodyAssembler<DeformedGrid, Vector> nBodyAssembler(bodyCount, bodyCount-1);
    nBodyAssembler.setGrids(deformedGrids);

    for (size_t i=0; i<couplings.size(); i++) {
        nBodyAssembler.setCoupling(*couplings[i], i);
    }

    nBodyAssembler.assembleTransferOperator();
    nBodyAssembler.assembleObstacle();

    // set up boundary conditions
    std::vector<BoundaryPatch<DefLeafGridView>> neumannBoundaries(bodyCount);
    std::vector<BoundaryPatch<DefLeafGridView>> frictionalBoundaries(bodyCount);
    std::vector<BoundaryPatch<DefLeafGridView>> surfaces(bodyCount);

    // Dirichlet boundary
    std::vector<Dune::BitSetVector<dims>> noNodes(bodyCount);
    std::vector<Dune::BitSetVector<dims>> dirichletNodes(bodyCount);

    // set up functions for time-dependent boundary conditions
    using Function = Dune::VirtualFunction<double, double>;
    const Function& velocityDirichletFunction = VelocityDirichletCondition();
    const Function& neumannFunction = NeumannCondition();

    for (size_t i=0; i<grids.size(); i++) {
      const auto& leafVertexCount = leafVertexCounts[i];

      std::cout << "Grid" << i << " Number of DOFs: " << leafVertexCount << std::endl;

      // Neumann boundary
      auto& neumannBoundary = neumannBoundaries[i];
      neumannBoundary = BoundaryPatch<DefLeafGridView>(*leafViews[i]);

      // friction boundary
      auto& frictionalBoundary = frictionalBoundaries[i];
      if (i==0) {
        frictionalBoundary = BoundaryPatch<DefLeafGridView>(leafFaces[i]->upper);
      } else if (i==bodyCount-1) {
        frictionalBoundary = BoundaryPatch<DefLeafGridView>(leafFaces[i]->lower);
      } else {
          frictionalBoundary = BoundaryPatch<DefLeafGridView>(leafFaces[i]->lower);
          frictionalBoundary.addPatch(BoundaryPatch<DefLeafGridView>(leafFaces[i]->upper));
      }
      //surfaces[i] = BoundaryPatch<GridView>(myFaces.upper);

      // TODO: Dirichlet Boundary
      noNodes[i] = Dune::BitSetVector<dims>(leafVertexCount);
      dirichletNodes[i] = Dune::BitSetVector<dims>(leafVertexCount);
      auto& gridDirichletNodes = dirichletNodes[i];
      for (int j=0; j<leafVertexCount; j++) {
        if (leafFaces[i]->right.containsVertex(j))
          gridDirichletNodes[j][0] = true;

        if (leafFaces[i]->lower.containsVertex(j))
          gridDirichletNodes[j][0] = true;

        #if MY_DIM == 3
        if (leafFaces[i]->front.containsVertex(j) || leafFaces[i]->back.containsVertex(j))
            gridDirichletNodes[j][2] = true;
        #endif
      }
    }

    // set up individual assemblers for each body, assemble problem (matrices, forces, rhs)
    std::vector<std::shared_ptr<MyAssembler>> assemblers(bodyCount);
    Matrices<Matrix> matrices(bodyCount);

    std::vector<Vector> gravityFunctionals(bodyCount);
    std::vector<std::function<void(double, Vector&)>> externalForces(bodyCount);

    for (size_t i=0; i<assemblers.size(); i++) {
      auto& assembler = assemblers[i];
      assembler = std::make_shared<MyAssembler>(*leafViews[i]);

      assembler->assembleElasticity(body.getYoungModulus(), body.getPoissonRatio(), *matrices.elasticity[i]);
      assembler->assembleViscosity(body.getShearViscosityField(), body.getBulkViscosityField(), *matrices.damping[i]);
      assembler->assembleMass(body.getDensityField(), *matrices.mass[i]);

      ScalarMatrix relativeFrictionalBoundaryMass;
      assembler->assembleFrictionalBoundaryMass(frictionalBoundaries[i], relativeFrictionalBoundaryMass);
      relativeFrictionalBoundaryMass /= frictionalBoundaries[i].area();
      EnergyNorm<ScalarMatrix, ScalarVector> const stateEnergyNorm(relativeFrictionalBoundaryMass);

      // assemble forces
      assembler->assembleBodyForce(body.getGravityField(), gravityFunctionals[i]);

      // Problem formulation: right-hand side
      externalForces[i] =
            [&](double _relativeTime, Vector &_ell) {
              assembler->assembleNeumann(neumannBoundaries[i], _ell, neumannFunction,
                                          _relativeTime);
              _ell += gravityFunctionals[i];
      };
    }

    /* Jonny 2bcontact
    // make dirichlet bitfields containing dirichlet information for both grids
           int size = rhs[0].size() + rhs[1].size();

           Dune::BitSetVector<dims> totalDirichletNodes(size);

           for (size_t i=0; i<dirichletNodes[0].size(); i++)
               for (int j=0; j<dims; j++)
                   totalDirichletNodes[i][j] = dirichletNodes[0][i][j];

           int offset = rhs[0].size();
           for (size_t i=0; i<dirichletNodes[1].size(); i++)
               for (int j=0; j<dims; j++)
                   totalDirichletNodes[offset + i][j] = dirichletNodes[1][i][j];

           // assemble separate linear elasticity problems
           std::array<MatrixType,2> stiffnessMatrix;
           std::array<const MatrixType*, 2> submat;

           for (size_t i=0; i<2; i++) {
               OperatorAssembler<P1Basis,P1Basis> globalAssembler(*p1Basis[i],*p1Basis[i]);

               double s = (i==0) ? E0 : E1;
               StVenantKirchhoffAssembler<GridType, P1Basis::LocalFiniteElement, P1Basis::LocalFiniteElement> localAssembler(s, nu);

               globalAssembler.assemble(localAssembler, stiffnessMatrix[i]);

               submat[i] = &stiffnessMatrix[i];
           }

           MatrixType bilinearForm;
           contactAssembler.assembleJacobian(submat, bilinearForm);
 */

    using MyProgramState = ProgramState<Vector, ScalarVector>;
    MyProgramState programState(leafVertexCounts);
    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, leafVertexCounts)
                         : nullptr;*/
/*
    if (firstRestart > 0) // automatically adjusts the time and timestep
      restartIO->read(firstRestart, programState);
    else
     programState.setupInitialConditions(parset, nBodyAssembler, externalForces,
                                          matrices, assemblers, dirichletNodes,
                                          noNodes, frictionalBoundaries, body);

    */

    // assemble friction
    std::vector<std::shared_ptr<MyGlobalFrictionData<LocalVector>>> frictionInfo(weakPatches.size());
    std::vector<std::shared_ptr<GlobalFriction<Matrix, Vector>>> globalFriction(weakPatches.size());
    for (size_t i=0; i<frictionInfo.size(); i++) {
        frictionInfo[i] = std::make_shared<MyGlobalFrictionData<LocalVector>>(parset.sub("boundary.friction"), weakPatches[i]);
        globalFriction[i] = assemblers[i]->assembleFrictionNonlinearity(parset.get<Config::FrictionModel>("boundary.friction.frictionModel"), frictionalBoundaries[i], *frictionInfo[i], programState.weightedNormalStress[i]);
        globalFriction[i]->updateAlpha(programState.alpha[i]);
    }


    /*
    Vector vertexCoordinates(leafVertexCount);
    {
      Dune::MultipleCodimMultipleGeomTypeMapper<
          GridView, Dune::MCMGVertexLayout> const vertexMapper(leafView, Dune::mcmgVertexLayout());
      for (auto &&v : vertices(leafView))
        vertexCoordinates[vertexMapper.index(v)] = geoToPoint(v.geometry());
    }

    using MyVertexBasis = typename MyAssembler::VertexBasis;
    auto dataWriter =
        writeData ? std::make_unique<
                        HDF5Writer<MyProgramState, MyVertexBasis, GridView>>(
                        *dataFile, vertexCoordinates, myAssembler.vertexBasis,
                        surface, frictionalBoundary, weakPatch)
                  : nullptr;
    MyVTKWriter<MyVertexBasis, typename MyAssembler::CellBasis> const vtkWriter(
        myAssembler.cellBasis, myAssembler.vertexBasis, "obs");


    IterationRegister iterationCount;
    auto const report = [&](bool initial = false) {
      if (writeData) {
        dataWriter->reportSolution(programState, *myGlobalFriction);
        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")) {
        ScalarVector stress;
        myAssembler.assembleVonMisesStress(body.getYoungModulus(),
                                           body.getPoissonRatio(),
                                           programState.u, stress);
        vtkWriter.write(programState.timeStep, programState.u, programState.v,
                        programState.alpha, stress);
      }
    };
    report(true);
    */

/*
    // Set up TNNMG solver
    using NonlinearFactory = SolverFactory<
        dims,
        MyBlockProblem<ConvexProblem<GlobalFriction<Matrix, Vector>, Matrix>>,
        Grid>;
    NonlinearFactory factory(parset.sub("solver.tnnmg"), *grid, dirichletNodes);

    using MyUpdater = Updaters<RateUpdater<Vector, Matrix, Function, dims>,
                               StateUpdater<ScalarVector, Vector>>;
    MyUpdater current(
        initRateUpdater(parset.get<Config::scheme>("timeSteps.scheme"),
                        velocityDirichletFunction, dirichletNodes, matrices,
                        programState.u, programState.v, programState.a),
        initStateUpdater<ScalarVector, Vector>(
            parset.get<Config::stateModel>("boundary.friction.stateModel"),
            programState.alpha, *frictionalBoundary.getVertices(),
            parset.get<double>("boundary.friction.L"),
            parset.get<double>("boundary.friction.V0")));

    auto const refinementTolerance =
        parset.get<double>("timeSteps.refinementTolerance");
    auto const mustRefine = [&](MyUpdater &coarseUpdater,
                                MyUpdater &fineUpdater) {
      ScalarVector coarseAlpha;
      coarseUpdater.state_->extractAlpha(coarseAlpha);

      ScalarVector fineAlpha;
      fineUpdater.state_->extractAlpha(fineAlpha);

      return stateEnergyNorm.diff(fineAlpha, coarseAlpha) > refinementTolerance;
    };

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

    AdaptiveTimeStepper<NonlinearFactory, MyUpdater,
                        EnergyNorm<ScalarMatrix, ScalarVector>>
        adaptiveTimeStepper(factory, parset, myGlobalFriction, current,
                            programState.relativeTime, programState.relativeTau,
                            computeExternalForces, stateEnergyNorm, 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);

      report();

      if (terminationRequested) {
        std::cerr << "Terminating prematurely" << std::endl;
        break;
      }
    }
*/
  } catch (Dune::Exception &e) {
    Dune::derr << "Dune reported error: " << e << std::endl;
  } catch (std::exception &e) {
    std::cerr << "Standard exception: " << e.what() << std::endl;
  }
}