#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 srcdir
#error srcdir unset
#endif

#ifndef DIM
#error DIM unset
#endif

#if !HAVE_ALUGRID
#error ALUGRID is required
#endif

#include <cmath>
#include <exception>
#include <iostream>

#include <boost/format.hpp>

#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/common/shared_ptr.hh>
#include <dune/common/timer.hh>

#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wignored-qualifiers"
#include <dune/grid/alugrid.hh>
#pragma clang diagnostic pop

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

#include <dune/fufem/assemblers/assembler.hh>
#include <dune/fufem/assemblers/localassemblers/boundarymassassembler.hh>
#include <dune/fufem/assemblers/localassemblers/l2functionalassembler.hh>
#include <dune/fufem/assemblers/localassemblers/massassembler.hh>
#include <dune/fufem/assemblers/localassemblers/stvenantkirchhoffassembler.hh>
#include <dune/fufem/assemblers/localassemblers/vonmisesstressassembler.hh>
#include <dune/fufem/boundarypatch.hh>
#include <dune/fufem/dunepython.hh>
#include <dune/fufem/functions/basisgridfunction.hh>
#include <dune/fufem/functions/constantfunction.hh>
#include <dune/fufem/functionspacebases/p0basis.hh>
#include <dune/fufem/functionspacebases/p1nodalbasis.hh>
#include <dune/fufem/sharedpointermap.hh>
#include <dune/solvers/norms/energynorm.hh>
#include <dune/solvers/norms/sumnorm.hh>
#include <dune/solvers/solvers/loopsolver.hh>
#include <dune/solvers/solvers/solver.hh> // Solver::FULL
#include <dune/tnnmg/problem-classes/convexproblem.hh>
#include <dune/tnnmg/nonlinearities/zerononlinearity.hh>
#include <dune/tnnmg/problem-classes/blocknonlineartnnmgproblem.hh>

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

#include "assemblers.hh"
#include "solverfactory.hh"
#include "vtk.hh"

#include "enums.hh"
#include "enum_parser.cc"
#include "enum_state_model.cc"
#include "enum_scheme.cc"
#include "enum_verbosity.cc"

#include "timestepping.hh"

#include "state/stateupdater.hh"
#include "state/ruinastateupdater.hh"
#include "state/dieterichstateupdater.hh"

int const dims = DIM;

template <class VectorType, class MatrixType, class FunctionType, int dims>
Dune::shared_ptr<TimeSteppingScheme<VectorType, MatrixType, FunctionType, dims>>
initTimeStepper(Config::scheme scheme,
                FunctionType const &velocityDirichletFunction,
                Dune::BitSetVector<dims> const &velocityDirichletNodes,
                MatrixType const &massMatrix, MatrixType const &stiffnessMatrix,
                VectorType const &u_initial, VectorType const &v_initial,
                VectorType const &a_initial) {
  switch (scheme) {
    case Config::Newmark:
      return Dune::make_shared<
          Newmark<VectorType, MatrixType, FunctionType, dims>>(
          stiffnessMatrix, massMatrix, u_initial, v_initial, a_initial,
          velocityDirichletNodes, velocityDirichletFunction);
    case Config::BackwardEuler:
      return Dune::make_shared<
          BackwardEuler<VectorType, MatrixType, FunctionType, dims>>(
          stiffnessMatrix, massMatrix, u_initial, v_initial,
          velocityDirichletNodes, velocityDirichletFunction);
    default:
      assert(false);
  }
}

template <class SingletonVectorType, class VectorType>
Dune::shared_ptr<StateUpdater<SingletonVectorType, VectorType>>
initStateUpdater(Config::stateModel model,
                 SingletonVectorType const &alpha_initial,
                 Dune::BitSetVector<1> const &frictionalNodes,
                 FrictionData const &fd) {
  switch (model) {
    case Config::Dieterich:
      return Dune::make_shared<
          DieterichStateUpdater<SingletonVectorType, VectorType>>(
          alpha_initial, frictionalNodes, fd.L);
    case Config::Ruina:
      return Dune::make_shared<
          RuinaStateUpdater<SingletonVectorType, VectorType>>(
          alpha_initial, frictionalNodes, fd.L);
    default:
      assert(false);
  }
}

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

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

double computeCOF(FrictionData const &fd, double V, double log_state) {
  double const mu = fd.mu0 + fd.a * std::log(V / fd.V0) +
                    fd.b * (log_state + std::log(fd.V0 / fd.L));
  return (mu > 0) ? mu : 0;
}

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

    Dune::ParameterTree parset;
    Dune::ParameterTreeParser::readINITree(srcdir "/one-body-sample.parset",
                                           parset);
    Dune::ParameterTreeParser::readOptions(argc, argv, parset);

    using SmallVector = Dune::FieldVector<double, dims>;
    using SmallMatrix = Dune::FieldMatrix<double, dims, dims>;
    using SmallSingletonMatrix = Dune::FieldMatrix<double, 1, 1>;
    using MatrixType = Dune::BCRSMatrix<SmallMatrix>;
    using SingletonMatrixType = Dune::BCRSMatrix<SmallSingletonMatrix>;
    using VectorType = Dune::BlockVector<SmallVector>;
    using SingletonVectorType = Dune::BlockVector<Dune::FieldVector<double, 1>>;
    using NonlinearityType = Dune::GlobalNonlinearity<MatrixType, VectorType>;

    auto const E = parset.get<double>("body.E");
    auto const nu = parset.get<double>("body.nu");

    // {{{ Set up grid
    using GridType =
        Dune::ALUGrid<dims, dims, Dune::simplex, Dune::nonconforming>;

    Dune::FieldVector<typename GridType::ctype, dims> lowerLeft(0);
    Dune::FieldVector<typename GridType::ctype, dims> upperRight(1);
    upperRight[0] = parset.get<size_t>("body.width");
    upperRight[1] = parset.get<size_t>("body.height");

    Dune::shared_ptr<GridType> grid;
    {
      Dune::array<unsigned int, dims> elements;
      std::fill(elements.begin(), elements.end(), 1);
      elements[0] = parset.get<size_t>("body.width");
      elements[1] = parset.get<size_t>("body.height");

      grid = Dune::StructuredGridFactory<GridType>::createSimplexGrid(
          lowerLeft, upperRight, elements);
    }

    auto const refinements = parset.get<size_t>("grid.refinements");
    grid->globalRefine(refinements);
    size_t const finestSize = grid->size(grid->maxLevel(), dims);

    using GridView = GridType::LeafGridView;
    GridView const leafView = grid->leafView();
    // }}}

    // Set up bases
    using P0Basis = P0Basis<GridView, double>;
    using P1Basis = P1NodalBasis<GridView, double>;
    P0Basis const p0Basis(leafView);
    P1Basis const p1Basis(leafView);
    Assembler<P0Basis, P0Basis> const p0Assembler(p0Basis, p0Basis);
    Assembler<P1Basis, P1Basis> const p1Assembler(p1Basis, p1Basis);

    // Set up the boundary
    Dune::BitSetVector<dims> velocityDirichletNodes(finestSize, false);
    Dune::BitSetVector<dims> const &displacementDirichletNodes =
        velocityDirichletNodes;
    Dune::BitSetVector<dims> accelerationDirichletNodes(finestSize, false);
    Dune::BitSetVector<1> neumannNodes(finestSize, false);
    Dune::BitSetVector<1> frictionalNodes(finestSize, false);

    VectorType coordinates(finestSize);
    {
      Dune::MultipleCodimMultipleGeomTypeMapper<
          GridView, Dune::MCMGVertexLayout> const myVertexMapper(leafView);

      for (auto it = leafView.begin<dims>(); it != leafView.end<dims>(); ++it) {
        assert(it->geometry().corners() == 1);
        size_t const id = myVertexMapper.map(*it);
        coordinates[id] = it->geometry().corner(0);
        auto const &localCoordinates = coordinates[id];

        // lower face
        if (localCoordinates[1] == lowerLeft[1]) {
          frictionalNodes[id] = true;
          velocityDirichletNodes[id][1] = true;
        }

        // upper face
        else if (localCoordinates[1] == upperRight[1])
          velocityDirichletNodes[id] = true;

        // right face (except for both corners)
        else if (localCoordinates[0] == upperRight[0])
          ;

        // left face (except for both corners)
        else if (localCoordinates[0] == lowerLeft[0])
          ;
      }
    }

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

    // Set up normal stress, mass matrix, and gravity functional
    double normalStress;
    MatrixType M;
    EnergyNorm<MatrixType, VectorType> const MNorm(M);
    VectorType gravityFunctional;
    {
      double const gravity = 9.81;
      double const density = parset.get<double>("body.density");
      {
        MassAssembler<GridType, P1Basis::LocalFiniteElement,
                      P1Basis::LocalFiniteElement,
                      Dune::ScaledIdentityMatrix<double, dims>> const localMass;
        p1Assembler.assembleOperator(localMass, M);
        M *= density;
      }
      {
        double volume = 1.0;
        for (int i = 0; i < dims; ++i)
          volume *= (upperRight[i] - lowerLeft[i]);

        double area = 1.0;
        for (int i = 0; i < dims; ++i)
          if (i != 1)
            area *= (upperRight[i] - lowerLeft[i]);

        // volume * gravity * density / area    = normal stress
        // V      * g       * rho     / A       = sigma_n
        // m^d    * N/kg    * kg/m^d  / m^(d-1) = N/m^(d-1)
        normalStress = volume * gravity * density / area;
      }
      {
        SmallVector weightedGravitationalDirection(0);
        weightedGravitationalDirection[1] = -density * gravity;
        ConstantFunction<SmallVector, SmallVector> const gravityFunction(
            weightedGravitationalDirection);
        L2FunctionalAssembler<GridType, SmallVector> gravityFunctionalAssembler(
            gravityFunction);
        p1Assembler.assembleFunctional(gravityFunctionalAssembler,
                                       gravityFunctional);
      }
    }
    FrictionData const frictionData(parset.sub("boundary.friction"),
                                    normalStress);

    MatrixType A;
    EnergyNorm<MatrixType, VectorType> const ANorm(A);
    {
      StVenantKirchhoffAssembler<GridType, P1Basis::LocalFiniteElement,
                                 P1Basis::LocalFiniteElement> const
      localStiffness(E, nu);
      p1Assembler.assembleOperator(localStiffness, A);
    }

    SingletonMatrixType frictionalBoundaryMassMatrix;
    EnergyNorm<SingletonMatrixType, SingletonVectorType> const stateEnergyNorm(
        frictionalBoundaryMassMatrix);
    {
      BoundaryPatch<GridView> const frictionalBoundary(leafView,
                                                       frictionalNodes);

      BoundaryMassAssembler<
          GridType, BoundaryPatch<GridView>, P1Basis::LocalFiniteElement,
          P1Basis::LocalFiniteElement, SmallSingletonMatrix> const
      frictionalBoundaryMassAssembler(frictionalBoundary);
      p1Assembler.assembleOperator(frictionalBoundaryMassAssembler,
                                   frictionalBoundaryMassMatrix);
    }
    // Q: Does it make sense to weigh them in this manner?
    SumNorm<VectorType> const AMNorm(1.0, ANorm, 1.0, MNorm);

    auto const nodalIntegrals = assemble_frictional<GridView, SmallVector>(
        leafView, p1Assembler, frictionalNodes);
    auto myGlobalNonlinearity = assemble_nonlinearity<MatrixType, VectorType>(
        frictionalNodes, *nodalIntegrals, frictionData);

    // Problem formulation: right-hand side
    auto const createRHS = [&](double _time, VectorType &_ell) {
      assemble_neumann(leafView, p1Assembler, neumannNodes, _ell,
                       neumannFunction, _time);
      _ell += gravityFunctional;
    };
    VectorType ell(finestSize);
    createRHS(0.0, ell);

    // {{{ Initial conditions
    SingletonVectorType alpha_initial(finestSize);
    alpha_initial =
        std::log(parset.get<double>("boundary.friction.initialState"));

    using LinearFactoryType = SolverFactory<
        dims, BlockNonlinearTNNMGProblem<ConvexProblem<
                  ZeroNonlinearity<SmallVector, SmallMatrix>, MatrixType>>,
        GridType>;
    ZeroNonlinearity<SmallVector, SmallMatrix> zeroNonlinearity;

    // Solve the stationary problem
    VectorType u_initial(finestSize);
    u_initial = 0.0;
    {
      LinearFactoryType displacementFactory(parset.sub("solver.tnnmg"), // FIXME
                                            refinements, *grid,
                                            displacementDirichletNodes);
      auto multigridStep = displacementFactory.getSolver();
      auto const verbosity = parset.get<Solver::VerbosityMode>("io.verbosity");

      typename LinearFactoryType::ConvexProblemType myConvexProblem(
          1.0, A, zeroNonlinearity, ell, u_initial);
      typename LinearFactoryType::BlockProblemType initialDisplacementProblem(
          parset, myConvexProblem);

      multigridStep->setProblem(u_initial, initialDisplacementProblem);
      LoopSolver<VectorType> initialDisplacementProblemSolver(
          multigridStep,
          parset.get<size_t>("solver.tnnmg.maxiterations"), // FIXME
          1e-12,                                            // FIXME
          &ANorm, verbosity, false);                        // absolute error

      initialDisplacementProblemSolver.preprocess();
      initialDisplacementProblemSolver.solve();
    }
    VectorType v_initial(finestSize);
    {
      // Prescribe a homogeneous velocity field in the x-direction
      // This way, the initial condition for v and the Dirichlet
      // condition match up at t=0
      v_initial = 0.0;
      double v_initial_const;
      velocityDirichletFunction.evaluate(0.0, v_initial_const);
      for (size_t i = 0; i < v_initial.size(); ++i)
        v_initial[i][0] = v_initial_const;
    }
    VectorType a_initial(finestSize);
    a_initial = 0.0;
    {
      /* We solve Au + Cv + Ma + Psi(v) = ell, thus
                                     Ma = - (Au + Cv + Psi(v) - ell)
      */
      MatrixType const &C = A;
      double const wc = 0.0; // watch out -- do not choose 1.0 here!
      VectorType accelerationRHS(finestSize);
      {
        accelerationRHS = 0.0;
        Arithmetic::addProduct(accelerationRHS, A, u_initial);
        Arithmetic::addProduct(accelerationRHS, wc, C, v_initial);
        myGlobalNonlinearity->updateState(alpha_initial);
        // NOTE: We assume differentiability of Psi at v0 here!
        myGlobalNonlinearity->addGradient(v_initial, accelerationRHS);
        accelerationRHS -= ell;
        // instead of multiplying M by (1.0 - wc), we divide the RHS
        accelerationRHS *= -1.0 / (1.0 - wc);
      }
      LinearFactoryType accelerationFactory(parset.sub("solver.tnnmg"),
                                            refinements, *grid,
                                            accelerationDirichletNodes);
      auto multigridStep = accelerationFactory.getSolver();
      auto const verbosity = parset.get<Solver::VerbosityMode>("io.verbosity");

      typename LinearFactoryType::ConvexProblemType myConvexProblem(
          1.0, M, zeroNonlinearity, accelerationRHS, a_initial);
      typename LinearFactoryType::BlockProblemType initialAccelerationProblem(
          parset, myConvexProblem);

      multigridStep->setProblem(a_initial, initialAccelerationProblem);
      LoopSolver<VectorType> initialAccelerationProblemSolver(
          multigridStep,
          parset.get<size_t>("solver.tnnmg.maxiterations"), // FIXME
          1e-12,                                            // FIXME
          &MNorm, verbosity, false);                        // absolute error

      initialAccelerationProblemSolver.preprocess();
      initialAccelerationProblemSolver.solve();
    }
    // }}}

    // Set up TNNMG solver
    auto const solverTolerance = parset.get<double>("solver.tolerance");
    using NonlinearFactoryType = SolverFactory<
        dims,
        MyBlockProblem<ConvexProblem<
            Dune::GlobalNonlinearity<MatrixType, VectorType>, MatrixType>>,
        GridType>;
    NonlinearFactoryType factory(parset.sub("solver.tnnmg"), refinements, *grid,
                                 velocityDirichletNodes);
    auto multigridStep = factory.getSolver();
    auto const verbosity = parset.get<Solver::VerbosityMode>("io.verbosity");

    {
      std::fstream coordinateWriter("coordinates", std::fstream::out);
      for (size_t i = 0; i < frictionalNodes.size(); ++i)
        if (frictionalNodes[i][0])
          coordinateWriter << coordinates[i] << std::endl;
      coordinateWriter.close();
    }
    std::fstream stateWriter("states", std::fstream::out);
    std::fstream displacementWriter("displacements", std::fstream::out);
    std::fstream velocityWriter("velocities", std::fstream::out);
    std::fstream coefficientWriter("coefficients", std::fstream::out);
    std::fstream iterationWriter("iterations", std::fstream::out);
    std::fstream dampingWriter("damping", std::fstream::out);

    auto timeSteppingScheme =
        initTimeStepper(parset.get<Config::scheme>("timeSteps.scheme"),
                        velocityDirichletFunction, velocityDirichletNodes, M, A,
                        u_initial, v_initial, a_initial);
    auto stateUpdater = initStateUpdater<SingletonVectorType, VectorType>(
        parset.get<Config::stateModel>("boundary.friction.stateModel"),
        alpha_initial, frictionalNodes, frictionData);

    auto const timesteps = parset.get<size_t>("timeSteps.number");
    auto const tau = parset.get<double>("problem.finalTime") / timesteps;

    VectorType v = v_initial;
    SingletonVectorType alpha = alpha_initial;
    auto const state_fpi_max =
        parset.get<size_t>("solver.tnnmg.fixed_point_iterations");
    auto const fixedPointTolerance =
        parset.get<double>("solver.tnnmg.fixed_point_tolerance");
    auto const damping = parset.get<double>("solver.damping");
    auto const requiredReduction =
        parset.get<double>("solver.requiredReduction");
    auto const printProgress = parset.get<bool>("io.printProgress");
    for (size_t run = 1; run <= timesteps; ++run) {
      VectorType u;
      double lastCorrection;

      stateUpdater->nextTimeStep();
      timeSteppingScheme->nextTimeStep();

      double const time = tau * run;
      createRHS(time, ell);

      MatrixType velocityMatrix;
      VectorType velocityRHS(finestSize);
      VectorType velocityIterate(finestSize);

      stateUpdater->setup(tau);
      timeSteppingScheme->setup(ell, tau, time, velocityRHS, velocityIterate,
                                velocityMatrix);

      LoopSolver<VectorType> velocityProblemSolver(
          multigridStep, parset.get<size_t>("solver.tnnmg.maxiterations"),
          solverTolerance, &AMNorm, verbosity, false); // absolute error

      size_t iterationCounter;
      auto solveVelocityProblem = [&](VectorType &_velocityIterate,
                                      SingletonVectorType const &_alpha) {
        myGlobalNonlinearity->updateState(_alpha);

        // NIT: Do we really need to pass u here?
        typename NonlinearFactoryType::ConvexProblemType const myConvexProblem(
            1.0, velocityMatrix, *myGlobalNonlinearity, velocityRHS,
            _velocityIterate);
        typename NonlinearFactoryType::BlockProblemType velocityProblem(
            parset, myConvexProblem);
        multigridStep->setProblem(_velocityIterate, velocityProblem);

        velocityProblemSolver.preprocess();
        velocityProblemSolver.solve();
        iterationCounter = velocityProblemSolver.getResult().iterations;
      };

      // Since the velocity explodes in the quasistatic case, use the
      // displacement as a convergence criterion
      // Q: is this reasonable?
      VectorType u_saved;
      for (size_t state_fpi = 1; state_fpi <= state_fpi_max; ++state_fpi) {
        stateUpdater->solve(v);
        {
          SingletonVectorType computed_state;
          stateUpdater->extractState(computed_state);
          double const correction = stateEnergyNorm.diff(computed_state, alpha);
          if (state_fpi <= 2 // Let the first two steps pass through unchanged
              || correction < requiredReduction * lastCorrection) {
            alpha = computed_state;
            dampingWriter << "N ";
          } else { // alpha is still the old time step here
            alpha *= damping;
            Arithmetic::addProduct(alpha, 1.0 - damping, computed_state);
            dampingWriter << "Y ";
          }
          lastCorrection = correction;
        }

        solveVelocityProblem(velocityIterate, alpha);
        timeSteppingScheme->postProcess(velocityIterate);
        timeSteppingScheme->extractDisplacement(u);
        timeSteppingScheme->extractVelocity(v);

        iterationWriter << iterationCounter << " ";
        if (printProgress)
          (std::cerr << '.').flush();

        if (state_fpi > 1) {
          double const correctionNorm = AMNorm.diff(u_saved, u);
          if (correctionNorm < fixedPointTolerance) {
            if (printProgress)
              (std::cerr << '#').flush();

            break;
          }
        }
        if (state_fpi == state_fpi_max)
          DUNE_THROW(Dune::Exception, "FPI failed to converge");

        u_saved = u;
      }
      if (printProgress)
        std::cerr << std::endl;

      for (size_t i = 0; i < frictionalNodes.size(); ++i)
        if (frictionalNodes[i][0]) {
          stateWriter << alpha[i][0] << " ";
          displacementWriter << u[i][0] << " ";
          velocityWriter << v[i][0] << " ";
          coefficientWriter << computeCOF(frictionData, v[i].two_norm(),
                                          alpha[i]) << " ";
        }

      stateWriter << std::endl;
      displacementWriter << std::endl;
      velocityWriter << std::endl;
      coefficientWriter << std::endl;
      iterationWriter << std::endl;
      dampingWriter << std::endl;

      if (parset.get<bool>("io.writeVTK")) {
        SingletonVectorType vonMisesStress;
        auto const gridDisplacement =
            Dune::make_shared<BasisGridFunction<P1Basis, VectorType> const>(
                p1Basis, u);
        VonMisesStressAssembler<GridType> localStressAssembler(
            E, nu, gridDisplacement);
        p0Assembler.assembleFunctional(localStressAssembler, vonMisesStress);

        writeVtk(p1Basis, u, alpha, p0Basis, vonMisesStress,
                 (boost::format("obs%d") % run).str());
      }
    }
    if (parset.get<bool>("io.enableTimer"))
      std::cerr << std::endl << "Making " << timesteps << " time steps took "
                << timer.elapsed() << "s" << std::endl;

    stateWriter.close();
    displacementWriter.close();
    velocityWriter.close();
    coefficientWriter.close();
    iterationWriter.close();
    dampingWriter.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;
  }
}