#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>
#include <dune/grid/alugrid.hh>
#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/tectonic/myblockproblem.hh>

#include "assemblers.hh"
#include "mysolver.hh"
#include "vtk.hh"

#include "enums.hh"
#include "enum_parser.cc"
#include "enum_state_model.cc"
#include "enum_scheme.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 &dirichletFunction,
                Dune::BitSetVector<dims> const &ignoreNodes,
                MatrixType const &massMatrix, MatrixType const &stiffnessMatrix,
                VectorType const &u_initial, VectorType const &ud_initial,
                VectorType const &udd_initial) {
  switch (scheme) {
    case Config::ImplicitEuler:
      return Dune::make_shared<
          ImplicitEuler<VectorType, MatrixType, FunctionType, dims>>(
          stiffnessMatrix, u_initial, ud_initial, ignoreNodes,
          dirichletFunction);
    case Config::Newmark:
      return Dune::make_shared<
          Newmark<VectorType, MatrixType, FunctionType, dims>>(
          stiffnessMatrix, massMatrix, u_initial, ud_initial, udd_initial,
          ignoreNodes, dirichletFunction);
    case Config::EulerPair:
      return Dune::make_shared<
          EulerPair<VectorType, MatrixType, FunctionType, dims>>(
          stiffnessMatrix, massMatrix, u_initial, ud_initial, ignoreNodes,
          dirichletFunction);
    default:
      assert(false);
  }
}

template <class SingletonVectorType, class VectorType>
Dune::shared_ptr<StateUpdater<SingletonVectorType, VectorType>>
initStateUpdater(Config::state_model 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);
  }
}

template <class FunctionMap> void initPython(FunctionMap &functions) {
  Python::start();

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

  Python::import("one-body-sample")
      .get("Functions")
      .toC<typename FunctionMap::Base>(functions);
}

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

    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> ignoreNodes(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;
          ignoreNodes[id][1] = true;
        }

        // upper face
        else if (localCoordinates[1] == upperRight[1])
          ignoreNodes[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>;
    SharedPointerMap<std::string, FunctionType> functions;
    initPython(functions);
    auto const &dirichletFunction = functions.get("dirichletCondition");
    auto const &neumannFunction = functions.get("neumannCondition");

    // Set up normal stress, mass matrix, and gravity functional
    double normalStress;
    MatrixType massMatrix;
    EnergyNorm<MatrixType, VectorType> const massMatrixNorm(massMatrix);
    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, massMatrix);
        massMatrix *= 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 stiffnessMatrix;
    EnergyNorm<MatrixType, VectorType> const stiffnessMatrixNorm(
        stiffnessMatrix);
    {
      StVenantKirchhoffAssembler<GridType, P1Basis::LocalFiniteElement,
                                 P1Basis::LocalFiniteElement> const
      localStiffness(E, nu);
      p1Assembler.assembleOperator(localStiffness, stiffnessMatrix);
    }

    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 velocityEnergyNorm(1.0, stiffnessMatrixNorm, 1.0,
                                                 massMatrixNorm);

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

    // {{{ Initial conditions
    VectorType u_initial(finestSize);
    u_initial = 0.0;
    VectorType ud_initial(finestSize);
    ud_initial = 0.0;
    // FIXME: This only happens to be correct
    VectorType udd_initial(finestSize);
    udd_initial = 0.0;

    SingletonVectorType alpha_initial(finestSize);
    alpha_initial =
        std::log(parset.get<double>("boundary.friction.initial_state"));
    // }}}

    // Set up TNNMG solver
    auto const solverTolerance = parset.get<double>("solver.tolerance");
    MySolver<dims, MatrixType, VectorType, GridType> solverHost(
        parset.sub("solver.tnnmg"), refinements, solverTolerance, *grid,
        ignoreNodes);
    auto multigridStep = solverHost.getSolver();
    Solver::VerbosityMode const verbosity =
        parset.get<bool>("verbose") ? Solver::FULL : Solver::QUIET;

    {
      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>("timeSteppingScheme"),
                        dirichletFunction, ignoreNodes, massMatrix,
                        stiffnessMatrix, u_initial, ud_initial, udd_initial);
    auto stateUpdater = initStateUpdater<SingletonVectorType, VectorType>(
        parset.get<Config::state_model>("boundary.friction.state_model"),
        alpha_initial, frictionalNodes, frictionData);

    auto const timesteps = parset.get<size_t>("timeSteps");
    auto const tau = parset.get<double>("endOfTime") / timesteps;
    auto const createRHS = [&](double _time, VectorType &_ell) {
      assemble_neumann(leafView, p1Assembler, neumannNodes, _ell,
                       neumannFunction, _time);
      _ell += gravityFunctional;
    };

    VectorType ud = ud_initial;
    SingletonVectorType alpha = alpha_initial;
    auto const state_fpi_max =
        parset.get<size_t>("solver.tnnmg.fixed_point_iterations");
    for (size_t run = 1; run <= timesteps; ++run) {
      VectorType u;
      double lastCorrection;

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

      double const time = tau * run;

      VectorType ell(finestSize);
      createRHS(time, ell);

      MatrixType problem_AB;
      VectorType problem_rhs(finestSize);
      VectorType problem_iterate(finestSize);

      stateUpdater->setup(tau);
      timeSteppingScheme->setup(ell, tau, time, problem_rhs, problem_iterate,
                                problem_AB);

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

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

        using ConvexProblemType = ConvexProblem<
            Dune::GlobalNonlinearity<MatrixType, VectorType>, MatrixType>;
        // FIXME: Do we really need to pass u here?
        ConvexProblemType const myConvexProblem(1.0, problem_AB,
                                                *myGlobalNonlinearity,
                                                problem_rhs, _problem_iterate);
        MyBlockProblem<ConvexProblemType> velocityProblem(parset,
                                                          myConvexProblem);
        multigridStep->setProblem(_problem_iterate, velocityProblem);

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

      // Since the velocity explodes in the quasistatic case, use the
      // displacement as a convergence criterion
      VectorType u_saved;
      double const fixedPointTolerance =
          parset.get<double>("solver.tnnmg.fixed_point_tolerance");
      double const damping = parset.get<double>("solver.damping");
      double const minimalCorrectionReduction =
          parset.get<double>("solver.minimal_correction_reduction");

      for (size_t state_fpi = 1; state_fpi <= state_fpi_max; ++state_fpi) {
        stateUpdater->solve(ud);
        {
          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 < minimalCorrectionReduction * lastCorrection) {
            alpha = computed_state;
            dampingWriter << "N ";
          } else { // alpha is still the old time step here
            alpha *= damping;
            alpha.axpy(1.0 - damping, computed_state);
            dampingWriter << "Y ";
          }
          lastCorrection = correction;
        }

        solveVelocityProblem(problem_iterate, alpha);
        timeSteppingScheme->postProcess(problem_iterate);
        timeSteppingScheme->extractDisplacement(u);
        timeSteppingScheme->extractVelocity(ud);

        iterationWriter << iterationCounter << " ";
        if (parset.get<bool>("printProgress")) {
          std::cerr << '.';
          std::cerr.flush();
        }
        if (state_fpi > 1) {
          double const correctionNorm = velocityEnergyNorm.diff(u_saved, u);
          if (correctionNorm < fixedPointTolerance) {
            if (parset.get<bool>("printProgress")) {
              std::cerr << '#';
              std::cerr.flush();
            }
            break;
          }
        }
        if (state_fpi == state_fpi_max)
          DUNE_THROW(Dune::Exception, "FPI failed to converge");

        u_saved = u;
      }
      if (parset.get<bool>("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 << ud[i][0] << " ";
          coefficientWriter << computeCOF(frictionData, ud[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>("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>("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;
  }
}