#ifndef SRC_PROGRAM_STATE_HH
#define SRC_PROGRAM_STATE_HH

#include <dune/common/parametertree.hh>

#include <dune/matrix-vector/axpy.hh>

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

#include <dune/contact/assemblers/nbodyassembler.hh>

#include <dune/tectonic/body.hh>

#include "assemblers.hh"
#include "matrices.hh"
#include "spatial-solving/solverfactory.hh"

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

  ProgramState(const std::vector<int>& leafVertexCounts)
    : bodyCount(leafVertexCounts.size()) {
    u.resize(bodyCount);
    v.resize(bodyCount);
    a.resize(bodyCount);
    alpha.resize(bodyCount);
    weightedNormalStress.resize(bodyCount);

    for (size_t i=0; i<bodyCount; i++) {
        size_t leafVertexCount = leafVertexCounts[i];
        u[i].resize(leafVertexCount);
        v[i].resize(leafVertexCount);
        a[i].resize(leafVertexCount);
        alpha[i].resize(leafVertexCount);
        weightedNormalStress[i].resize(leafVertexCount);
    }
  }

  // Set up initial conditions
  template <class Matrix, class GridView>
  void setupInitialConditions(
      const Dune::ParameterTree& parset,
      const Dune::Contact::NBodyAssembler<typename GridView::Grid, Vector>& nBodyAssembler,
      std::vector<std::function<void(double, Vector &)>> externalForces,
      const Matrices<Matrix>& matrices,
      const std::vector<std::shared_ptr<MyAssembler<GridView, Vector::block_type::dimension>>>& assemblers,
      const std::vector<Dune::BitSetVector<Vector::block_type::dimension>>& dirichletNodes,
      const std::vector<Dune::BitSetVector<Vector::block_type::dimension>>& noNodes,
      const std::vector<BoundaryPatch<GridView>>& frictionalBoundaries,
      const Body<Vector::block_type::dimension>& body) {

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

    /*
    // Solving a linear problem with a multigrid solver
    auto const solveLinearProblem = [&](
        Dune::BitSetVector<dims> const &_dirichletNodes, const std::vector<std::shared_ptr<Matrix>>& _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
                            assemblers.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);
      //multigridStep->setProblem(_x);
      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();

      Vector totalX = multigridStep->getSol();

      // cleanup
      delete(multigridStep);

      nBodyAssembler.postprocess(totalX, x);
    };
    */

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

    std::vector<Vector> ell0(bodyCount);
    for (size_t i=0; i<bodyCount; i++) {
      // Initial velocity
      v[i] = 0.0;

      ell0[i].resize(u[i].size());
      externalForces[i](relativeTime, ell0[i]);
    }

    // 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
    std::vector<Vector> accelerationRHS = ell0;
    for (size_t i=0; i<bodyCount; i++) {
      // Initial state
      alpha[i] = parset.get<double>("boundary.friction.initialAlpha");

      // Initial normal stress
      assemblers[i]->assembleWeightedNormalStress(
        frictionalBoundaries[i], weightedNormalStress[i], body.getYoungModulus(),
        body.getPoissonRatio(), u[i]);

      Dune::MatrixVector::subtractProduct(accelerationRHS[i], *matrices.elasticity[i], u[i]);
    }

    solveLinearProblem(noNodes, matrices.mass, accelerationRHS, a,
                       parset.sub("a0.solver"));
  }

public:
  std::vector<Vector> u;
  std::vector<Vector> v;
  std::vector<Vector> a;
  std::vector<ScalarVector> alpha;
  std::vector<ScalarVector> weightedNormalStress;
  double relativeTime;
  double relativeTau;
  size_t timeStep;

private:
  const size_t bodyCount;
};

#endif