// Based on dune/tnnmg/problem-classes/blocknonlineartnnmgproblem.hh

#ifndef MY_BLOCK_PROBLEM_HH
#define MY_BLOCK_PROBLEM_HH

#include <dune/common/bitsetvector.hh>
#include <dune/common/nullptr.hh>
#include <dune/common/parametertree.hh>

#include <dune/fufem/arithmetic.hh>
#include <dune/solvers/common/interval.hh>
#include <dune/solvers/computeenergy.hh>
#include <dune/tnnmg/problem-classes/bisection.hh>

#include "ellipticenergy.hh"
#include "globalnonlinearity.hh"
#include "minimisation.hh"
#include "mydirectionalconvexfunction.hh"

/** \brief Base class for problems where each block can be solved with a
 * modified gradient method */
template <class ConvexProblem> class MyBlockProblem {
public:
  using ConvexProblemType = ConvexProblem;
  using VectorType = typename ConvexProblem::VectorType;
  using MatrixType = typename ConvexProblem::MatrixType;
  using LocalVector = typename ConvexProblem::LocalVectorType;
  using LocalMatrix = typename ConvexProblem::LocalMatrixType;
  size_t static const block_size = ConvexProblem::block_size;
  size_t static const coarse_block_size = block_size;

  /** \brief Solves one local system using a modified gradient method */
  class IterateObject;

  struct Linearization {
    size_t static const block_size = coarse_block_size;

    using LocalMatrix = typename MyBlockProblem<ConvexProblem>::LocalMatrix;
    using MatrixType = Dune::BCRSMatrix<typename Linearization::LocalMatrix>;
    using VectorType =
        Dune::BlockVector<Dune::FieldVector<double, Linearization::block_size>>;
    using BitVectorType = Dune::BitSetVector<Linearization::block_size>;

    typename Linearization::MatrixType A;
    typename Linearization::VectorType b;
    typename Linearization::BitVectorType ignore;

    Dune::BitSetVector<Linearization::block_size> truncation;
  };

  MyBlockProblem(Dune::ParameterTree const &parset,
                 ConvexProblem const &problem)
      : parset(parset), problem(problem), localBisection() // NOTE: defaults
  {}

  std::string getOutput(bool header = false) const {
    if (header) {
      outStream.str("");
      for (size_t j = 0; j < block_size; ++j)
        outStream << "  trunc" << std::setw(2) << j;
    }
    std::string s = outStream.str();
    outStream.str("");
    return s;
  }

  double computeEnergy(const VectorType &v) const {
    return 0.0; // FIXME
    // return ::computeEnergy(problem_.A, v, problem_.f) + problem_.phi(v);
  }

  void projectCoarseCorrection(VectorType const &u,
                               typename Linearization::VectorType const &v,
                               VectorType &projected_v,
                               Linearization const &linearization) const {
    projected_v = v;
    for (size_t i = 0; i < v.size(); ++i)
      for (size_t j = 0; j < block_size; ++j)
        if (linearization.truncation[i][j])
          projected_v[i][j] = 0;
  }

  double computeDampingParameter(VectorType const &u,
                                 VectorType const &projected_v) const {
    VectorType v = projected_v;

    double const vnorm = v.two_norm();
    if (vnorm <= 0)
      return 1.0;

    v /= vnorm; // Rescale for numerical stability

    MyDirectionalConvexFunction<
        GlobalNonlinearity<MatrixType, VectorType>> const
    psi(computeDirectionalA(problem.A, v),
        computeDirectionalb(problem.A, problem.f, u, v), problem.phi, u, v);

    Dune::Solvers::Interval<double> D;
    psi.subDiff(0, D);
    // NOTE: Numerical instability can actually get us here
    if (D[1] > 0)
      return 0;

    int bisectionsteps = 0;
    Bisection const globalBisection; // NOTE: defaults
    return globalBisection.minimize(psi, vnorm, 0.0, bisectionsteps) / vnorm;
  }

  void assembleTruncate(VectorType const &u, Linearization &linearization,
                        Dune::BitSetVector<block_size> const &ignore) const {
    // we can just copy the ignore information
    linearization.ignore = ignore;

    // determine truncation pattern
    linearization.truncation.resize(u.size());
    linearization.truncation.unsetAll();
    for (size_t i = 0; i < u.size(); ++i) {
      if (problem.phi.regularity(i, u[i]) > 1e8) { // TODO: Make customisable
        linearization.truncation[i] = true;
        continue;
      }

      for (size_t j = 0; j < block_size; ++j)
        if (linearization.ignore[i][j])
          linearization.truncation[i][j] = true;
    }

    // construct sparsity pattern for linearization
    Dune::MatrixIndexSet indices(problem.A.N(), problem.A.M());
    indices.import(problem.A);
    problem.phi.addHessianIndices(indices);

    // construct matrix from pattern and initialize it
    indices.exportIdx(linearization.A);
    linearization.A = 0.0;

    // compute quadratic part of hessian (linearization.A += problem.A)
    for (size_t i = 0; i < problem.A.N(); ++i) {
      auto const end = problem.A[i].end();
      for (auto it = problem.A[i].begin(); it != end; ++it)
        linearization.A[i][it.index()] += *it;
    }

    // compute nonlinearity part of hessian
    problem.phi.addHessian(u, linearization.A);

    // compute quadratic part of gradient
    linearization.b.resize(u.size());
    problem.A.mv(u, linearization.b);
    linearization.b -= problem.f;

    // compute nonlinearity part of gradient
    problem.phi.addGradient(u, linearization.b);

    // -grad is needed for Newton step
    linearization.b *= -1.0;

    // apply truncation to stiffness matrix and rhs
    for (size_t row = 0; row < linearization.A.N(); ++row) {
      auto const col_end = linearization.A[row].end();
      for (auto col_it = linearization.A[row].begin(); col_it != col_end;
           ++col_it) {
        size_t const col = col_it.index();
        for (size_t i = 0; i < col_it->N(); ++i) {
          auto const blockEnd = (*col_it)[i].end();
          for (auto blockIt = (*col_it)[i].begin(); blockIt != blockEnd;
               ++blockIt)
            if (linearization.truncation[row][i] ||
                linearization.truncation[col][blockIt.index()])
              *blockIt = 0.0;
        }
      }

      for (size_t j = 0; j < block_size; ++j)
        if (linearization.truncation[row][j])
          linearization.b[row][j] = 0.0;
    }

    for (size_t j = 0; j < block_size; ++j)
      outStream << std::setw(9) << linearization.truncation.countmasked(j);
  }

  /** \brief Constructs and returns an iterate object */
  IterateObject getIterateObject() {
    return IterateObject(parset, localBisection, problem);
  }

private:
  Dune::ParameterTree const &parset;

  // problem data
  ConvexProblem const &problem;

  Bisection const localBisection;

  mutable std::ostringstream outStream;
};

/** \brief Solves one local system using a scalar Gauss-Seidel method */
template <class ConvexProblem>
class MyBlockProblem<ConvexProblem>::IterateObject {
  friend class MyBlockProblem;

protected:
  /** \brief Constructor, protected so only friends can instantiate it
   * \param bisection The class used to do a scalar bisection
   * \param problem The problem including quadratic part and nonlinear part
   */
  IterateObject(Dune::ParameterTree const &parset, Bisection const &bisection,
                ConvexProblem const &problem)
      : problem(problem),
        bisection_(bisection),
        localsteps(parset.get<size_t>("localsolver.steps")) {}

public:
  /** \brief Set the current iterate */
  void setIterate(VectorType &u) {
    this->u = u;
    return;
  }

  /** \brief Update the i-th block of the current iterate */
  void updateIterate(LocalVector const &ui, size_t i) {
    u[i] = ui;
    return;
  }

  /** \brief Minimise a local problem using a modified gradient method
   * \param[out] ui The solution
   * \param m Block number
   * \param ignore Set of degrees of freedom to leave untouched
   */
  void solveLocalProblem(
      LocalVector &ui, size_t m,
      typename Dune::BitSetVector<block_size>::const_reference ignore) {
    {
      LocalMatrix const *localA = nullptr;
      LocalVector localb(problem.f[m]);

      auto const end = problem.A[m].end();
      for (auto it = problem.A[m].begin(); it != end; ++it) {
        size_t const j = it.index();
        if (j == m)
          localA = &(*it); // localA = A[m][m]
        else
          Arithmetic::subtractProduct(localb, *it, u[j]);
      }
      assert(localA != nullptr);

      auto const phi = problem.phi.restriction(m);
      EllipticEnergy<block_size> localJ(*localA, localb, phi, ignore);
      minimise(localJ, ui, localsteps, bisection_);
    }
  }

private:
  // problem data
  ConvexProblem const &problem;

  Bisection const bisection_;

  // state data for smoothing procedure used by:
  // setIterate, updateIterate, solveLocalProblem
  VectorType u;

  size_t const localsteps;
};

#endif