#ifndef SRC_PROGRAM_STATE_HH
#define SRC_PROGRAM_STATE_HH

#include <dune/common/parametertree.hh>

#include <dune/fufem/boundarypatch.hh>
#include <dune/tnnmg/nonlinearities/zerononlinearity.hh>
#include <dune/tnnmg/problem-classes/blocknonlineartnnmgproblem.hh>

#include <dune/tectonic/body.hh>

#include "assemblers.hh"
#include "matrices.hh"
#include "solverfactory.hh"

template <class VectorTEMPLATE, class ScalarVectorTEMPLATE> class ProgramState {
public:
  using Vector = VectorTEMPLATE;
  using ScalarVector = ScalarVectorTEMPLATE;

  ProgramState(int leafVertexCount)
      : u(leafVertexCount),
        v(leafVertexCount),
        a(leafVertexCount),
        alpha(leafVertexCount),
        weightedNormalStress(leafVertexCount) {}

  // Set up initial conditions
  template <class Matrix, class GridView>
  void setupInitialConditions(
      Dune::ParameterTree const &parset,
      std::function<void(double, Vector &)> externalForces,
      Matrices<Matrix> const matrices,
      MyAssembler<GridView, Vector::block_type::dimension> const &myAssembler,
      Dune::BitSetVector<Vector::block_type::dimension> const &dirichletNodes,
      Dune::BitSetVector<Vector::block_type::dimension> const &noNodes,
      BoundaryPatch<GridView> const &frictionalBoundary,
      Body<Vector::block_type::dimension> const &body) {

    using LocalVector = typename Vector::block_type;
    using LocalMatrix = typename Matrix::block_type;
    auto const dims = LocalVector::dimension;

    // Solving a linear problem with a multigrid solver
    auto const solveLinearProblem = [&](
        Dune::BitSetVector<dims> const &_dirichletNodes, Matrix const &_matrix,
        Vector const &_rhs, Vector &_x,
        Dune::ParameterTree const &_localParset) {

      using LinearFactory = SolverFactory<
          dims, BlockNonlinearTNNMGProblem<ConvexProblem<
                    ZeroNonlinearity<LocalVector, LocalMatrix>, Matrix>>,
          typename GridView::Grid>;
      ZeroNonlinearity<LocalVector, LocalMatrix> zeroNonlinearity;

      LinearFactory factory(parset.sub("solver.tnnmg"), // FIXME
                            myAssembler.gridView.grid(), _dirichletNodes);

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

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

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

    timeStep = 0;
    relativeTime = 0.0;
    relativeTau = 1e-6;

    Vector ell0(u.size());
    externalForces(relativeTime, ell0);

    // Initial velocity
    v = 0.0;

    // Initial displacement: Start from a situation of minimal stress,
    // which is automatically attained in the case [v = 0 = a].
    // Assuming dPhi(v = 0) = 0, we thus only have to solve Au = ell0
    solveLinearProblem(dirichletNodes, matrices.elasticity, ell0, u,
                       parset.sub("u0.solver"));

    // Initial acceleration: Computed in agreement with Ma = ell0 - Au
    // (without Dirichlet constraints), again assuming dPhi(v = 0) = 0
    Vector accelerationRHS = ell0;
    Arithmetic::subtractProduct(accelerationRHS, matrices.elasticity, u);
    solveLinearProblem(noNodes, matrices.mass, accelerationRHS, a,
                       parset.sub("a0.solver"));

    // Initial state
    alpha = parset.get<double>("boundary.friction.initialAlpha");

    // Initial normal stress
    myAssembler.assembleWeightedNormalStress(
        frictionalBoundary, weightedNormalStress, body.getYoungModulus(),
        body.getPoissonRatio(), u);
  }

public:
  Vector u;
  Vector v;
  Vector a;
  ScalarVector alpha;
  ScalarVector weightedNormalStress;
  double relativeTime;
  double relativeTau;
  size_t timeStep;
};

#endif