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

#ifdef HAVE_IPOPT
#undef HAVE_IPOPT
#endif

#ifndef srcdir
#error srcdir unset
#endif

#ifndef DIM
#error DIM unset
#endif

#ifndef HAVE_PYTHON
#error Python is required
#endif

#if !HAVE_ALUGRID
#error ALUGRID is required
#endif

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

#include <boost/format.hpp>

#include <Python.h>

#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/functionalassembler.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/assemblers/operatorassembler.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/solvers/loopsolver.hh>
#include <dune/solvers/solvers/solver.hh> // Solver::FULL

#include <dune/tectonic/myblockproblem.hh>
#include <dune/tectonic/myconvexproblem.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::ImplicitTwoStep:
      return Dune::make_shared<
          ImplicitTwoStep<VectorType, MatrixType, FunctionType, dims>>(
          stiffnessMatrix, u_initial, ud_initial, ignoreNodes,
          dirichletFunction);
    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);
  }
}
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, double L) {
  switch (model) {
    case Config::Dieterich:
      return Dune::make_shared<
          DieterichStateUpdater<SingletonVectorType, VectorType>>(
          alpha_initial, frictionalNodes, L);
    case Config::Ruina:
      return Dune::make_shared<
          RuinaStateUpdater<SingletonVectorType, VectorType>>(
          alpha_initial, frictionalNodes, L);
  }
}
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);
}

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

    typedef Dune::FieldVector<double, dims> SmallVector;
    typedef Dune::FieldMatrix<double, dims, dims> SmallMatrix;
    typedef Dune::BCRSMatrix<SmallMatrix> MatrixType;
    typedef Dune::BlockVector<SmallVector> VectorType;
    typedef Dune::BlockVector<Dune::FieldVector<double, 1>> SingletonVectorType;

    auto const E = parset.get<double>("body.E");
    auto const nu = parset.get<double>("body.nu");
    auto const mu = parset.get<double>("boundary.friction.mu");
    auto const a = parset.get<double>("boundary.friction.a");
    auto const b = parset.get<double>("boundary.friction.b");
    auto const V0 = parset.get<double>("boundary.friction.V0");
    auto const L = parset.get<double>("boundary.friction.L");

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

    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::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");

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

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

    // Set up bases
    typedef P0Basis<GridView, double> P0Basis;
    typedef P1NodalBasis<GridView, double> P1Basis;
    P0Basis const p0Basis(leafView);
    P1Basis const p1Basis(leafView);

    // 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
    typedef Dune::VirtualFunction<double, double> FunctionType;
    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;
    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;
        OperatorAssembler<P1Basis, P1Basis>(p1Basis, p1Basis)
            .assemble(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);
        FunctionalAssembler<P1Basis>(p1Basis)
            .assemble(gravityFunctionalAssembler, gravityFunctional, true);
      }
    }
    SingletonVectorType surfaceNormalStress(finestSize);
    surfaceNormalStress = normalStress;

    MatrixType stiffnessMatrix;
    {
      StVenantKirchhoffAssembler<GridType, P1Basis::LocalFiniteElement,
                                 P1Basis::LocalFiniteElement> const
      localStiffness(E, nu);
      OperatorAssembler<P1Basis, P1Basis>(p1Basis, p1Basis)
          .assemble(localStiffness, stiffnessMatrix);
    };

    auto const nodalIntegrals =
        assemble_frictional<GridType, GridView, SmallVector, P1Basis>(
            leafView, p1Basis, frictionalNodes);

    // {{{ Initial conditions
    VectorType u_initial(finestSize);
    u_initial = 0.0;
    VectorType ud_initial(finestSize);
    ud_initial = 0.0;
    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 solver_tolerance = parset.get<double>("solver.tolerance");
    typedef MyConvexProblem<MatrixType, VectorType> MyConvexProblemType;
    typedef MyBlockProblem<MyConvexProblemType> MyBlockProblemType;
    MySolver<dims, MatrixType, VectorType, GridType, MyBlockProblemType>
    mySolver(parset.sub("solver.tnnmg"), refinements, solver_tolerance, *grid,
             ignoreNodes);
    auto multigridStep = mySolver.getSolver();
    Solver::VerbosityMode const verbosity =
        parset.get<bool>("verbose") ? Solver::FULL : Solver::QUIET;

    std::fstream coordinate_writer("coordinates", std::fstream::out);
    for (size_t i = 0; i < frictionalNodes.size(); ++i)
      if (frictionalNodes[i][0])
        coordinate_writer << coordinates[i] << std::endl;
    coordinate_writer.close();

    std::fstream state_writer("states", std::fstream::out);
    std::fstream displacement_writer("displacements", std::fstream::out);
    std::fstream velocity_writer("velocities", std::fstream::out);
    std::fstream coefficient_writer("coefficients", std::fstream::out);
    ;
    std::fstream iteration_writer("iterations", 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, L);

    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<GridType, GridView, SmallVector, P1Basis>(
          leafView, p1Basis, neumannNodes, _ell, neumannFunction, _time);
      _ell += gravityFunctional;
    };

    VectorType ud = ud_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;
      SingletonVectorType alpha;

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

      double const time = tau * run;

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

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

      stateUpdater->setup(tau);
      timeSteppingScheme->setup(ell, tau, time, problem_rhs, problem_iterate,
                                problem_A);
      EnergyNorm<MatrixType, VectorType> energyNorm(problem_A);
      LoopSolver<VectorType> overallSolver(
          multigridStep, parset.get<size_t>("solver.tnnmg.maxiterations"),
          solver_tolerance, &energyNorm, verbosity, false); // absolute error

      size_t iterationCounter;
      auto solveDisplacementProblem = [&](VectorType &_problem_iterate,
                                          SingletonVectorType const &_alpha) {
        auto myGlobalNonlinearity =
            assemble_nonlinearity<MatrixType, VectorType>(
                parset.sub("boundary.friction"), *nodalIntegrals, _alpha,
                surfaceNormalStress);

        MyConvexProblemType const myConvexProblem(
            problem_A, *myGlobalNonlinearity, problem_rhs);
        MyBlockProblemType myBlockProblem(parset, myConvexProblem);
        multigridStep->setProblem(_problem_iterate, myBlockProblem);

        overallSolver.preprocess();
        overallSolver.solve();
        iterationCounter = overallSolver.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");
      for (size_t state_fpi = 1; state_fpi <= state_fpi_max; ++state_fpi) {
        stateUpdater->solve(ud);
        stateUpdater->extractState(alpha);

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

        iteration_writer << iterationCounter << " ";
        if (parset.get<bool>("printProgress")) {
          std::cerr << '.';
          std::cerr.flush();
        }
        if (state_fpi > 1) {
          double const correctionNorm = energyNorm.diff(u_saved, u);
          if (correctionNorm < fixedPointTolerance)
            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]) {
          state_writer << alpha[i][0] << " ";
          displacement_writer << u[i][0] << " ";
          velocity_writer << ud[i][0] << " ";
          coefficient_writer << mu + a *std::log(ud[i].two_norm() / V0) +
                                    b * (alpha[i] + std::log(V0 / L)) << " ";
        }

      state_writer << std::endl;
      displacement_writer << std::endl;
      velocity_writer << std::endl;
      coefficient_writer << std::endl;
      iteration_writer << std::endl;

      if (parset.get<bool>("writeVTK")) {
        SingletonVectorType vonMisesStress;
        VonMisesStressAssembler<GridType> localStressAssembler(
            E, nu,
            Dune::make_shared<BasisGridFunction<P1Basis, VectorType> const>(
                p1Basis, u));
        FunctionalAssembler<P0Basis>(p0Basis)
            .assemble(localStressAssembler, vonMisesStress, true);

        writeVtk<P1Basis, P0Basis, VectorType, SingletonVectorType, GridView>(
            p1Basis, u, alpha, p0Basis, vonMisesStress, leafView,
            (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;
    ;

    state_writer.close();
    displacement_writer.close();
    velocity_writer.close();
    coefficient_writer.close();
    iteration_writer.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;
  }
}