#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

#ifndef DIM
#error DIM unset
#endif

#if !HAVE_ALUGRID
#error ALUGRID is required
#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>

#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/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/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 "assemblers.hh"
#include "tobool.hh"
#include "enumparser.hh"
#include "enums.hh"
#include "friction_writer.hh"
#include "sand-wedge-data/mybody.hh"
#include "sand-wedge-data/mygeometry.hh"
#include "sand-wedge-data/mygeometry.cc" // FIXME
#include "sand-wedge-data/myglobalfrictiondata.hh"
#include "sand-wedge-data/mygrid.cc" // FIXME
#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 = DIM;

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

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

template <class Factory, class StateUpdater, class VelocityUpdater>
class FixedPointIterator {
  using ScalarVector = typename StateUpdater::ScalarVector;
  using Vector = typename Factory::Vector;
  using Matrix = typename Factory::Matrix;
  using ConvexProblem = typename Factory::ConvexProblem;
  using BlockProblem = typename Factory::BlockProblem;
  using Nonlinearity = typename ConvexProblem::NonlinearityType;

public:
  FixedPointIterator(Factory &factory, Dune::ParameterTree const &parset,
                     std::shared_ptr<Nonlinearity> globalFriction)
      : factory_(factory),
        parset_(parset),
        globalFriction_(globalFriction),
        fixedPointMaxIterations_(parset.get<size_t>("v.fpi.maximumIterations")),
        fixedPointTolerance_(parset.get<double>("v.fpi.tolerance")),
        lambda_(parset.get<double>("v.fpi.lambda")),
        velocityMaxIterations_(
            parset.get<size_t>("v.solver.maximumIterations")),
        velocityTolerance_(parset.get<double>("v.solver.tolerance")),
        verbosity_(parset.get<Solver::VerbosityMode>("v.solver.verbosity")) {}

  int run(std::shared_ptr<StateUpdater> stateUpdater,
          std::shared_ptr<VelocityUpdater> velocityUpdater,
          Matrix const &velocityMatrix, Norm<Vector> const &velocityMatrixNorm,
          Vector const &velocityRHS, Vector &velocityIterate) {
    auto multigridStep = factory_.getSolver();

    LoopSolver<Vector> velocityProblemSolver(
        multigridStep, velocityMaxIterations_, velocityTolerance_,
        &velocityMatrixNorm, verbosity_, false); // absolute error

    Vector previousVelocityIterate = velocityIterate;

    size_t fixedPointIteration;
    for (fixedPointIteration = 1;
         fixedPointIteration <= fixedPointMaxIterations_;
         ++fixedPointIteration) {
      Vector v_m;
      velocityUpdater->extractOldVelocity(v_m);
      v_m *= 1.0 - lambda_;
      Arithmetic::addProduct(v_m, lambda_, velocityIterate);

      // solve a state problem
      stateUpdater->solve(v_m);
      ScalarVector alpha;
      stateUpdater->extractAlpha(alpha);

      // solve a velocity problem
      globalFriction_->updateAlpha(alpha);
      ConvexProblem convexProblem(1.0, velocityMatrix, *globalFriction_,
                                  velocityRHS, velocityIterate);
      BlockProblem velocityProblem(parset_, convexProblem);
      multigridStep->setProblem(velocityIterate, velocityProblem);
      velocityProblemSolver.preprocess();
      velocityProblemSolver.solve();

      if (velocityMatrixNorm.diff(previousVelocityIterate, velocityIterate) <
          fixedPointTolerance_)
        break;

      previousVelocityIterate = velocityIterate;
    }
    if (fixedPointIteration == fixedPointMaxIterations_)
      DUNE_THROW(Dune::Exception, "FPI failed to converge");

    velocityUpdater->postProcess(velocityIterate);
    velocityUpdater->postProcessRelativeQuantities();

    return fixedPointIteration;
  }

private:
  Factory &factory_;
  Dune::ParameterTree const &parset_;
  std::shared_ptr<Nonlinearity> globalFriction_;

  size_t fixedPointMaxIterations_;
  double fixedPointTolerance_;
  double lambda_;
  size_t velocityMaxIterations_;
  double velocityTolerance_;
  Solver::VerbosityMode verbosity_;
};

template <class Factory, class StateUpdater, class VelocityUpdater>
class CoupledTimeStepper {
  using Vector = typename Factory::Vector;
  using Matrix = typename Factory::Matrix;
  using ConvexProblem = typename Factory::ConvexProblem;
  using Nonlinearity = typename ConvexProblem::NonlinearityType;

public:
  CoupledTimeStepper(double finalTime, Factory &factory,
                     Dune::ParameterTree const &parset,
                     std::shared_ptr<Nonlinearity> globalFriction,
                     std::shared_ptr<StateUpdater> stateUpdater,
                     std::shared_ptr<VelocityUpdater> velocityUpdater,
                     std::function<void(double, Vector &)> externalForces)
      : finalTime_(finalTime),
        factory_(factory),
        parset_(parset),
        globalFriction_(globalFriction),
        stateUpdater_(stateUpdater),
        velocityUpdater_(velocityUpdater),
        externalForces_(externalForces) {}

  int step(double relativeTime, double relativeTau) {
    stateUpdater_->nextTimeStep();
    velocityUpdater_->nextTimeStep();

    auto const newRelativeTime = relativeTime + relativeTau;
    Vector ell;
    externalForces_(newRelativeTime, ell);

    Matrix velocityMatrix;
    Vector velocityRHS;
    Vector velocityIterate;

    auto const tau = relativeTau * finalTime_;
    stateUpdater_->setup(tau);
    velocityUpdater_->setup(ell, tau, newRelativeTime, velocityRHS,
                            velocityIterate, velocityMatrix);
    EnergyNorm<Matrix, Vector> const velocityMatrixNorm(velocityMatrix);

    FixedPointIterator<Factory, StateUpdater, VelocityUpdater>
    fixedPointIterator(factory_, parset_, globalFriction_);
    auto const iterations = fixedPointIterator.run(
        stateUpdater_, velocityUpdater_, velocityMatrix, velocityMatrixNorm,
        velocityRHS, velocityIterate);
    return iterations;
  }

private:
  double finalTime_;
  Factory &factory_;
  Dune::ParameterTree const &parset_;
  std::shared_ptr<Nonlinearity> globalFriction_;
  std::shared_ptr<StateUpdater> stateUpdater_;
  std::shared_ptr<VelocityUpdater> velocityUpdater_;
  std::function<void(double, Vector &)> externalForces_;
};

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
    using Grid = Dune::ALUGrid<dims, dims, Dune::simplex, Dune::nonconforming>;
    GridConstructor<Grid> gridConstructor;
    auto grid = gridConstructor.getGrid();

    auto const refinements = parset.get<size_t>("grid.refinements");
    grid->globalRefine(refinements);

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

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

    auto velocityUpdater = initTimeStepper(
        parset.get<Config::scheme>("timeSteps.scheme"),
        velocityDirichletFunction, dirichletNodes, M, A, C, u_initial,
        ur_initial, v_initial, vr_initial, a_initial);
    auto stateUpdater = 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"));

    auto const finalTime = parset.get<double>("problem.finalTime");
    double relativeTime = 0.0;

    double relativeTau = 1e-6; // FIXME (not really important, though)

    using MyStateUpdater = StateUpdater<ScalarVector, Vector>;
    using MyTimeStepper = TimeSteppingScheme<Vector, Matrix, Function, dims>;

    auto const refinementTolerance =
        parset.get<double>("timeSteps.refinementTolerance");

    size_t timeStep = 1;

    std::fstream iterationWriter("iterations", std::fstream::out);

    auto stateUpdaterR1 = stateUpdater->clone();
    auto velocityUpdaterR1 = velocityUpdater->clone();
    {
      CoupledTimeStepper<NonlinearFactory, MyStateUpdater, MyTimeStepper>
      coupledTimeStepper(finalTime, factory, parset, myGlobalFriction,
                         stateUpdaterR1, velocityUpdaterR1,
                         computeExternalForces);
      iterationWriter << "R1 ";
      auto const iterations =
          coupledTimeStepper.step(relativeTime, relativeTau);
      iterationWriter << iterations << std::endl;
    }

    std::shared_ptr<MyStateUpdater> stateUpdaterR2 = nullptr;
    std::shared_ptr<MyTimeStepper> velocityUpdaterR2 = nullptr;

    while (relativeTime < 1.0 - 1e-10) {
      bool didCoarsen = false;

      while (true) {
        stateUpdaterR2 = stateUpdaterR1->clone();
        velocityUpdaterR2 = velocityUpdaterR1->clone();

        ScalarVector alphaR2;
        {
          CoupledTimeStepper<NonlinearFactory, MyStateUpdater, MyTimeStepper>
          coupledTimeStepper(finalTime, factory, parset, myGlobalFriction,
                             stateUpdaterR2, velocityUpdaterR2,
                             computeExternalForces);
          iterationWriter << "R2 ";
          auto const iterations =
              coupledTimeStepper.step(relativeTime + relativeTau, relativeTau);
          iterationWriter << iterations << " " << std::flush;
        }
        stateUpdaterR2->extractAlpha(alphaR2);

        auto stateUpdaterC = stateUpdater->clone();
        auto velocityUpdaterC = velocityUpdater->clone();

        ScalarVector alphaC;
        {
          CoupledTimeStepper<NonlinearFactory, MyStateUpdater, MyTimeStepper>
          coupledTimeStepper(finalTime, factory, parset, myGlobalFriction,
                             stateUpdaterC, velocityUpdaterC,
                             computeExternalForces);
          iterationWriter << "C ";
          auto const iterations =
              coupledTimeStepper.step(relativeTime, 2.0 * relativeTau);
          iterationWriter << iterations << " " << std::flush;
        }
        stateUpdaterC->extractAlpha(alphaC);

        auto const coarseningError = stateEnergyNorm.diff(alphaC, alphaR2);

        if (coarseningError < refinementTolerance) {
          stateUpdaterR2 = nullptr;
          velocityUpdaterR2 = nullptr;

          stateUpdaterR1 = stateUpdaterC;
          velocityUpdaterR1 = velocityUpdaterC;
          relativeTau *= 2.0;
          didCoarsen = true;
        } else {
          break;
        }
      }

      if (!didCoarsen) {
        ScalarVector alphaR1;
        while (true) {
          auto stateUpdaterF2 = stateUpdater->clone();
          auto velocityUpdaterF2 = velocityUpdater->clone();

          std::shared_ptr<MyStateUpdater> stateUpdaterF1;
          std::shared_ptr<MyTimeStepper> velocityUpdaterF1;

          ScalarVector alphaF2;
          {
            CoupledTimeStepper<NonlinearFactory, MyStateUpdater, MyTimeStepper>
            coupledTimeStepper(finalTime, factory, parset, myGlobalFriction,
                               stateUpdaterF2, velocityUpdaterF2,
                               computeExternalForces);
            iterationWriter << "F1 ";
            auto const iterationsF1 =
                coupledTimeStepper.step(relativeTime, relativeTau / 2.0);
            iterationWriter << iterationsF1 << " " << std::flush;

            stateUpdaterF1 = stateUpdaterF2->clone();
            velocityUpdaterF1 = velocityUpdaterF2->clone();

            iterationWriter << "F2 ";
            auto const iterationsF2 = coupledTimeStepper.step(
                relativeTime + relativeTau / 2.0, relativeTau / 2.0);
            iterationWriter << iterationsF2 << " " << std::flush;
          }
          stateUpdaterF2->extractAlpha(alphaF2);
          stateUpdaterR1->extractAlpha(alphaR1);
          auto const refinementError = stateEnergyNorm.diff(alphaR1, alphaF2);

          if (refinementError < refinementTolerance) {
            break;
          } else {
            stateUpdaterR1 = stateUpdaterF1;
            velocityUpdaterR1 = velocityUpdaterF1;
            stateUpdaterR2 = stateUpdaterF2;
            velocityUpdaterR2 = velocityUpdaterF2;
            relativeTau /= 2.0;
          }
        }
      }
      iterationWriter << std::endl;

      reportTimeStep(relativeTime, relativeTau);

      stateUpdater = stateUpdaterR1;
      velocityUpdater = velocityUpdaterR1;
      stateUpdaterR1 = stateUpdaterR2;
      velocityUpdaterR1 = velocityUpdaterR2;
      relativeTime += relativeTau;

      Vector u, ur, vr;
      ScalarVector alpha;
      velocityUpdater->extractDisplacement(u);
      velocityUpdater->extractRelativeDisplacement(ur);
      velocityUpdater->extractRelativeVelocity(vr);
      stateUpdater->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();
    iterationWriter.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;
  }
}