#include <Python.h>

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

#ifndef HAVE_PYTHON
#error Python is required
#endif

#ifdef HAVE_IPOPT
#undef HAVE_IPOPT
#endif

#ifndef datadir
#error datadir unset
#endif

#include <cmath>
#include <exception>
#include <fstream>
#include <iostream>
#include <iomanip>

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

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

#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wdeprecated-declarations"
#include <dune/fufem/boundarypatch.hh>
#pragma clang diagnostic pop

#include <dune/fufem/dunepython.hh>
#include <dune/fufem/functions/basisgridfunction.hh>
#include <dune/fufem/sharedpointermap.hh>
#include <dune/solvers/norms/energynorm.hh>
#include <dune/solvers/norms/sumnorm.hh>
#include <dune/solvers/norms/twonorm.hh>
#include <dune/solvers/solvers/loopsolver.hh>
#include <dune/solvers/solvers/solver.hh>
#include <dune/tnnmg/nonlinearities/zerononlinearity.hh>
#include <dune/tnnmg/problem-classes/blocknonlineartnnmgproblem.hh>
#include <dune/tnnmg/problem-classes/convexproblem.hh>

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

#include "adaptivetimestepper.hh" // FIXME
#include "assemblers.hh"
#include "tobool.hh"
#include "coupledtimestepper.hh"
#include "enumparser.hh"
#include "enums.hh"
#include "fixedpointiterator.hh"
#include "friction_writer.hh"
#include "gridselector.hh"
#include "sand-wedge-data/mybody.hh"
#include "sand-wedge-data/mygeometry.hh"
#include "sand-wedge-data/myglobalfrictiondata.hh"
#include "sand-wedge-data/mygrid.hh"
#include "sand-wedge-data/special_writer.hh"
#include "solverfactory.hh"
#include "state.hh"
#include "timestepping.hh"
#include "vtk.hh"

size_t const dims = MY_DIM;

void initPython() {
  Python::start();

  Python::run("import sys");
  Python::run("sys.path.append('" datadir "')");
}

template <class Geometry> double diameter(Geometry const &geometry) {
  auto const numCorners = geometry.corners();
  std::vector<typename Geometry::GlobalCoordinate> corners(numCorners);
  for (int i = 0; i < numCorners; ++i)
    corners[i] = geometry.corner(i);

  TwoNorm<typename Geometry::GlobalCoordinate> twoNorm;

  double diameter = 0.0;
  for (int i = 0; i < numCorners; ++i)
    for (int j = 0; j < i; ++j)
      diameter = std::max(diameter, twoNorm.diff(corners[i], corners[j]));
  return diameter;
}

#include <dune/tectonic/closestpointprojection.hh>

// we assume the weakening region to be a line segment for now
template <class Geometry> double distanceToWeakeningRegion(Geometry const &g) {
  TwoNorm<typename Geometry::GlobalCoordinate> twoNorm;

  BoundarySegment<typename Geometry::GlobalCoordinate> bs;
  bs.nVertices = 2;
  bs.vertices.resize(bs.nVertices);
  bs.vertices[0] = MyGeometry::X;
  bs.vertices[1] = MyGeometry::Y;

  double distance = std::numeric_limits<double>::infinity();
  auto const numCorners = g.corners();
  for (int i = 0; i < numCorners; ++i) {
    auto const corner = g.corner(i);
    auto const projection = computeClosestPoint(corner, bs, 10); // FIXME
    distance = std::min(distance, twoNorm.diff(corner, projection));
  }
  return distance;
}

double computeAdmissibleDiameter(double distance, double smallestDiameter) {
  return (distance / 0.0125 + 1.0) * smallestDiameter;
}

int main(int argc, char *argv[]) {
  try {
    Dune::ParameterTree parset;
    Dune::ParameterTreeParser::readINITree(datadir "/parset.cfg", parset);
    Dune::ParameterTreeParser::readOptions(argc, argv, parset);

    MyGeometry::render();
    MyGeometry::write();

    // {{{ Set up grid
    GridConstructor<Grid> gridConstructor;
    auto grid = gridConstructor.getGrid();

    auto const smallestDiameter =
        parset.get<double>("boundary.friction.smallestDiameter");
    bool needRefine = true;

    while (true) {
      needRefine = false;
      for (auto it = grid->template leafbegin<0>();
           it != grid->template leafend<0>(); ++it) {
        auto const geometry = it->geometry();

        auto const weakeningRegionDistance =
            distanceToWeakeningRegion(geometry);
        auto const admissibleDiameter = computeAdmissibleDiameter(
            weakeningRegionDistance, smallestDiameter);

        if (diameter(geometry) <= admissibleDiameter)
          continue;

        needRefine = true;
        grid->mark(1, *it);
      }
      if (!needRefine)
        break;

      grid->preAdapt();
      grid->adapt();
      grid->postAdapt();
    }
    auto const refinements = grid->maxLevel();

    double minDiameter = std::numeric_limits<double>::infinity();
    double maxDiameter = 0.0;
    for (auto it = grid->template leafbegin<0>();
         it != grid->template leafend<0>(); ++it) {
      auto const geometry = it->geometry();
      auto const diam = diameter(geometry);
      minDiameter = std::min(minDiameter, diam);
      maxDiameter = std::max(maxDiameter, diam);
    }
    std::cout << "min diameter: " << minDiameter << std::endl;
    std::cout << "max diameter: " << maxDiameter << std::endl;

    using GridView = Grid::LeafGridView;
    GridView const leafView = grid->leafGridView();
    size_t const leafVertexCount = leafView.size(dims);

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

    auto myFaces = gridConstructor.constructFaces(leafView);

    // Neumann boundary
    BoundaryPatch<GridView> const neumannBoundary(leafView);

    // Frictional Boundary
    BoundaryPatch<GridView> const &frictionalBoundary = myFaces.lower;
    Dune::BitSetVector<1> frictionalNodes(leafVertexCount);
    frictionalBoundary.getVertices(frictionalNodes);

    // Surface
    BoundaryPatch<GridView> const &surface = myFaces.upper;
    Dune::BitSetVector<1> surfaceNodes(leafVertexCount);
    surface.getVertices(surfaceNodes);

    // Dirichlet Boundary
    Dune::BitSetVector<dims> noNodes(leafVertexCount);
    Dune::BitSetVector<dims> dirichletNodes(leafVertexCount);
    for (size_t i = 0; i < leafVertexCount; ++i) {
      if (myFaces.right.containsVertex(i))
        dirichletNodes[i][0] = true;

      if (myFaces.lower.containsVertex(i))
        dirichletNodes[i][1] = true;

#if MY_DIM == 3
      if (myFaces.front.containsVertex(i) || myFaces.back.containsVertex(i))
        dirichletNodes[i][2] = true;
#endif
    }

    // Set up functions for time-dependent boundary conditions
    using Function = Dune::VirtualFunction<double, double>;
    using FunctionMap = SharedPointerMap<std::string, Function>;
    FunctionMap functions;
    {
      initPython();
      Python::import("boundaryconditions")
          .get("Functions")
          .toC<typename FunctionMap::Base>(functions);
    }
    auto const &velocityDirichletFunction =
                   functions.get("velocityDirichletCondition"),
               &neumannFunction = functions.get("neumannCondition");

    std::fstream timeStepWriter("timeSteps", std::fstream::out);
    timeStepWriter << std::fixed << std::setprecision(10);
    auto const reportTimeStep = [&](double _relativeTime, double _relativeTau) {
      timeStepWriter << _relativeTime << " " << _relativeTau << std::endl;
    };

    using MyAssembler = MyAssembler<GridView, dims>;
    using Matrix = MyAssembler::Matrix;
    using LocalMatrix = Matrix::block_type;
    using Vector = MyAssembler::Vector;
    using LocalVector = Vector::block_type;
    using ScalarMatrix = MyAssembler::ScalarMatrix;
    using ScalarVector = MyAssembler::ScalarVector;
    using LocalScalarVector = ScalarVector::block_type;

    MyAssembler myAssembler(leafView);

    MyBody<dims> const body(parset);

    Matrix A, C, M;
    myAssembler.assembleElasticity(body.getYoungModulus(),
                                   body.getPoissonRatio(), A);
    myAssembler.assembleViscosity(body.getShearViscosityField(),
                                  body.getBulkViscosityField(), C);
    myAssembler.assembleMass(body.getDensityField(), M);
    EnergyNorm<Matrix, Vector> const ANorm(A), MNorm(M);

    ScalarMatrix frictionalBoundaryMass;
    myAssembler.assembleFrictionalBoundaryMass(frictionalBoundary,
                                               frictionalBoundaryMass);
    EnergyNorm<ScalarMatrix, ScalarVector> const stateEnergyNorm(
        frictionalBoundaryMass);

    // Assemble forces
    Vector gravityFunctional;
    myAssembler.assembleBodyForce(body.getGravityField(), gravityFunctional);

    // Problem formulation: right-hand side
    std::function<void(double, Vector &)> computeExternalForces = [&](
        double _relativeTime, Vector &_ell) {
      myAssembler.assembleNeumann(neumannBoundary, _ell, neumannFunction,
                                  _relativeTime);
      _ell += gravityFunctional;
    };
    Vector ell0(leafVertexCount);
    computeExternalForces(0.0, ell0);

    // {{{ Initial conditions
    using LinearFactory = SolverFactory<
        dims, BlockNonlinearTNNMGProblem<ConvexProblem<
                  ZeroNonlinearity<LocalVector, LocalMatrix>, Matrix>>,
        Grid>;
    ZeroNonlinearity<LocalVector, LocalMatrix> zeroNonlinearity;

    auto const solveLinearProblem = [&](
        Dune::BitSetVector<dims> const &_dirichletNodes, Matrix const &_matrix,
        Vector const &_rhs, Vector &_x, EnergyNorm<Matrix, Vector> const &_norm,
        Dune::ParameterTree const &_localParset) {
      LinearFactory factory(parset.sub("solver.tnnmg"), // FIXME
                            refinements, *grid, _dirichletNodes);

      typename LinearFactory::ConvexProblem convexProblem(
          1.0, _matrix, zeroNonlinearity, _rhs, _x);
      typename LinearFactory::BlockProblem problem(parset, convexProblem);

      auto multigridStep = factory.getSolver();
      multigridStep->setProblem(_x, problem);
      LoopSolver<Vector> solver(
          multigridStep, _localParset.get<size_t>("maximumIterations"),
          _localParset.get<double>("tolerance"), &_norm,
          _localParset.get<Solver::VerbosityMode>("verbosity"),
          false); // absolute error

      solver.preprocess();
      solver.solve();
    };

    // Solve the stationary problem
    Vector u_initial(leafVertexCount);
    u_initial = 0.0;
    solveLinearProblem(dirichletNodes, A, ell0, u_initial, ANorm,
                       parset.sub("u0.solver"));
    Vector ur_initial(leafVertexCount);
    ur_initial = 0.0;

    ScalarVector alpha_initial(leafVertexCount);
    alpha_initial = parset.get<double>("boundary.friction.initialAlpha");
    ScalarVector normalStress(leafVertexCount);
    myAssembler.assembleNormalStress(frictionalBoundary, normalStress,
                                     body.getYoungModulus(),
                                     body.getPoissonRatio(), u_initial);
    MyGlobalFrictionData<dims> frictionInfo(parset.sub("boundary.friction"));
    auto myGlobalFriction = myAssembler.assembleFrictionNonlinearity(
        parset.get<Config::FrictionModel>("boundary.friction.frictionModel"),
        frictionalBoundary, frictionInfo, normalStress);
    myGlobalFriction->updateAlpha(alpha_initial);

    Vector v_initial(leafVertexCount);
    v_initial = 0.0;
    {
      double v_initial_const;
      velocityDirichletFunction.evaluate(0.0, v_initial_const);
      assert(std::abs(v_initial_const) < 1e-14);
    }
    Vector vr_initial(leafVertexCount);
    vr_initial = 0.0;

    Vector a_initial(leafVertexCount);
    a_initial = 0.0;
    {
      // We solve Ma = ell0 - [Au + Cv + DPsi(v)]
      Vector accelerationRHS(leafVertexCount);
      {
        accelerationRHS = 0.0;
        Arithmetic::addProduct(accelerationRHS, A, u_initial);
        Arithmetic::addProduct(accelerationRHS, C, v_initial);
        // NOTE: We assume differentiability of Psi at 0 here!
        myGlobalFriction->addGradient(v_initial, accelerationRHS);
        accelerationRHS *= -1.0;
        accelerationRHS += ell0;
      }
      solveLinearProblem(noNodes, M, accelerationRHS, a_initial, MNorm,
                         parset.sub("a0.solver"));
    }
    // }}}

    Vector vertexCoordinates(leafVertexCount);
    {
      Dune::MultipleCodimMultipleGeomTypeMapper<
          GridView, Dune::MCMGVertexLayout> const vertexMapper(leafView);
      for (auto it = leafView.begin<dims>(); it != leafView.end<dims>(); ++it) {
        auto const geometry = it->geometry();
        assert(geometry.corners() == 1);
        vertexCoordinates[vertexMapper.map(*it)] = geometry.corner(0);
      }
    }

    FrictionWriter<ScalarVector, Vector> frictionWriter(
        vertexCoordinates, frictionalNodes, "friction",
        MyGeometry::horizontalProjection);
    BoundaryWriter<ScalarVector, Vector> verticalSurfaceWriter(
        vertexCoordinates, surfaceNodes, "verticalSurface",
        MyGeometry::verticalProjection);
    BoundaryWriter<ScalarVector, Vector> horizontalSurfaceWriter(
        vertexCoordinates, surfaceNodes, "horizontalSurface",
        MyGeometry::horizontalProjection);

    auto const report = [&](Vector const &_u, Vector const &_v,
                            ScalarVector const &_alpha) {
      horizontalSurfaceWriter.writeKinetics(_u, _v);
      verticalSurfaceWriter.writeKinetics(_u, _v);

      ScalarVector c;
      myGlobalFriction->coefficientOfFriction(_v, c);
      frictionWriter.writeKinetics(_u, _v);
      frictionWriter.writeOther(c, _alpha);
    };
    report(ur_initial, vr_initial, alpha_initial);

    MyVTKWriter<typename MyAssembler::VertexBasis,
                typename MyAssembler::CellBasis> const
    vtkWriter(myAssembler.cellBasis, myAssembler.vertexBasis, "obs");

    if (parset.get<bool>("io.writeVTK")) {
      ScalarVector stress;
      myAssembler.assembleVonMisesStress(
          body.getYoungModulus(), body.getPoissonRatio(), u_initial, stress);
      vtkWriter.write(0, ur_initial, vr_initial, alpha_initial, stress);
    }

    SpecialWriter<GridView, dims> specialVelocityWriter("specialVelocities",
                                                        leafView);
    SpecialWriter<GridView, dims> specialDisplacementWriter(
        "specialDisplacements", leafView);

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

    using UpdaterPair = std::pair<
        std::shared_ptr<StateUpdater<ScalarVector, Vector>>,
        std::shared_ptr<TimeSteppingScheme<Vector, Matrix, Function, dims>>>;
    UpdaterPair current(
        initStateUpdater<ScalarVector, Vector>(
            parset.get<Config::stateModel>("boundary.friction.stateModel"),
            alpha_initial, frictionalNodes,
            parset.get<double>("boundary.friction.L"),
            parset.get<double>("boundary.friction.V0")),
        initTimeStepper(parset.get<Config::scheme>("timeSteps.scheme"),
                        velocityDirichletFunction, dirichletNodes, M, A, C,
                        u_initial, ur_initial, v_initial, vr_initial,
                        a_initial));

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

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

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

    size_t timeStep = 1;

    AdaptiveTimeStepper<NonlinearFactory, UpdaterPair> adaptiveTimeStepper(
        factory, parset, myGlobalFriction, current, computeExternalForces,
        mustRefine);
    while (!adaptiveTimeStepper.reachedEnd()) {
      adaptiveTimeStepper.advance();
      reportTimeStep(adaptiveTimeStepper.getRelativeTime(),
                     adaptiveTimeStepper.getRelativeTau());

      Vector u, ur, vr;
      ScalarVector alpha;
      current.second->extractDisplacement(u);
      current.second->extractRelativeDisplacement(ur);
      current.second->extractRelativeVelocity(vr);
      current.first->extractAlpha(alpha);

      report(ur, vr, alpha);
      {
        BasisGridFunction<typename MyAssembler::VertexBasis, Vector>
        relativeVelocity(myAssembler.vertexBasis, vr);
        BasisGridFunction<typename MyAssembler::VertexBasis, Vector>
        relativeDisplacement(myAssembler.vertexBasis, ur);
        specialVelocityWriter.write(relativeVelocity);
        specialDisplacementWriter.write(relativeDisplacement);
      }

      if (parset.get<bool>("io.writeVTK")) {
        ScalarVector stress;
        myAssembler.assembleVonMisesStress(body.getYoungModulus(),
                                           body.getPoissonRatio(), u, stress);
        vtkWriter.write(timeStep, ur, vr, alpha, stress);
      }
      timeStep++;
    }
    timeStepWriter.close();
    Python::stop();
  }
  catch (Dune::Exception &e) {
    Dune::derr << "Dune reported error: " << e << std::endl;
  }
  catch (std::exception &e) {
    std::cerr << "Standard exception: " << e.what() << std::endl;
  }
}