// Based on dune/tnnmg/problem-classes/blocknonlineargsproblem.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/tnnmg/problem-classes/bisection.hh>

#include "globalnonlinearity.hh"
#include "localnonlinearity.hh"
#include "mydirectionalconvexfunction.hh"
#include "samplefunctional.hh"

/* Just for debugging */
template <int dim, class VectorType, class MatrixType>
double computeEnergy(
    MatrixType const &A, VectorType const &b,
    Dune::GlobalNonlinearity<VectorType, MatrixType> const &phi,
    VectorType const &x) {
  double ret;
  VectorType tmp(x.size());
  A.mv(x, tmp);          //            Ax
  ret = 0.5 * (tmp * x); // ret = 1/2 <Ax,x>

  ret -= b * x;  // ret = 1/2 <Ax,x> - <b,x>
  ret += phi(x); // ret = 1/2 <Ax,x> - <b,x> + phi(x)
  return ret;
}

/** \brief Base class for problems where each block can be solved with a
 * modified gradient method */
template <class MyConvexProblemTypeTEMPLATE> class MyBlockProblem {
public:
  typedef MyConvexProblemTypeTEMPLATE MyConvexProblemType;
  typedef typename MyConvexProblemType::VectorType VectorType;
  typedef typename MyConvexProblemType::MatrixType MatrixType;
  typedef typename MyConvexProblemType::LocalVectorType LocalVectorType;
  typedef typename MyConvexProblemType::LocalMatrixType LocalMatrixType;

  int static const block_size = MyConvexProblemType::block_size;
  int static const coarse_block_size = block_size;

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

  struct Linearization {
    static const int block_size = coarse_block_size;

    typedef typename MyBlockProblem<MyConvexProblemType>::LocalMatrixType
    LocalMatrixType;
    typedef Dune::BCRSMatrix<typename Linearization::LocalMatrixType>
    MatrixType;
    typedef Dune::BlockVector<
        Dune::FieldVector<double, Linearization::block_size>> VectorType;
    typedef Dune::BitSetVector<Linearization::block_size> BitVectorType;
    typename Linearization::MatrixType A;
    typename Linearization::VectorType b;
    typename Linearization::BitVectorType ignore;

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

  MyBlockProblem(Dune::ParameterTree const &parset,
                 const MyConvexProblemType &problem)
      : parset(parset), problem(problem) {
    bisection = Bisection(
        0.0, // acceptError: Stop if the search interval has
             // become smaller than this number
        parset.get<double>("bisection.acceptFactor"),
        parset.get<double>("bisection.requiredResidual"), true, // fastQuadratic
        0); // safety: acceptance factor for inexact minimization
  }

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

  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 (int 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) // This can (and needs to be able to) be very small
      return 0.0;

    v /= vnorm; // Rescale for numerical stability

    VectorType tmp = problem.f;
    problem.A.mmv(u, tmp);
    double const localb = tmp * v;

    problem.A.mv(v, tmp);
    double const localA = tmp * v;

    /*
      1/2 <A(u + hv),u + hv> - <b, u + hv>
      = 1/2 <Av,v> h^2 - <b - Au, v> h + const.

      localA = <Av,v>
      localb = <b - Au, v>
    */

    MyDirectionalConvexFunction<
        Dune::GlobalNonlinearity<VectorType, MatrixType>> const psi(localA,
                                                                    localb,
                                                                    problem.phi,
                                                                    u, v);

    Interval<double> D;
    psi.subDiff(0, D);
    if (D[1] > 0) {
      // NOTE: Numerical instability can actually get us here
      assert(std::abs(D[1]) < 1e-15);
      return 0;
    }

    int bisectionsteps = 0;
    Bisection bisection(0.0, 1.0, 1e-12, true, 0);                      // TODO
    return bisection.minimize(psi, vnorm, 0.0, bisectionsteps) / vnorm; // TODO
  }

  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) {
        linearization.truncation[i] = true;
        continue;
      }

      for (int 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) {
      typename MatrixType::row_type::ConstIterator it = problem.A[i].begin();
      typename MatrixType::row_type::ConstIterator end = problem.A[i].end();
      for (; it != end; ++it)
        Arithmetic::addProduct(linearization.A[i][it.index()], 1.0, *it);
    }

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

    // compute quadratic part of gradient
    linearization.b.resize(u.size());
    linearization.b = 0.0;
    problem.A.template umv<VectorType, VectorType>(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;

    // b should be a descent direction
    {
      VectorType const direction = linearization.b;
      VectorType tmp = linearization.b;      //  b
      linearization.A.mmv(u, tmp);           //  b-Au
      double const localA = tmp * direction; // <b-Au,v>

      linearization.A.mv(direction, tmp);    //  Av
      double const localb = tmp * direction; // <Av,v>

      MyDirectionalConvexFunction<
          Dune::GlobalNonlinearity<VectorType, MatrixType>> const
      psi(localA, localb, problem.phi, u, direction);

      Interval<double> D;
      psi.subDiff(0, D);
      if (!isnan(D[1]))
        assert(D[1] <= 0);
    }

    // apply truncation to system
    for (size_t row = 0; row < linearization.A.N(); ++row) {
      auto const end = linearization.A[row].end();
      for (auto it = linearization.A[row].begin(); it != end; ++it) {
        int const col = it.index();
        for (size_t i = 0; i < it->N(); ++i) {
          auto const blockEnd = (*it)[i].end();
          for (auto blockIt = (*it)[i].begin(); blockIt != blockEnd; ++blockIt)
            if (linearization.truncation[row][i] or linearization
                    .truncation[col][blockIt.index()])
              *blockIt = 0.0;
        }
      }

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

    for (int 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, bisection, problem);
  }

private:
  // problem data
  MyConvexProblemType const &problem;

  // commonly used minimization stuff
  Bisection bisection;

  Dune::ParameterTree const &parset;

  mutable std::ostringstream outStream;
};

/** \brief Solves one local system using a scalar Gauss-Seidel method */
template <class MyConvexProblemTypeTEMPLATE>
class MyBlockProblem<MyConvexProblemTypeTEMPLATE>::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,
                MyConvexProblemType const &problem)
      : parset(parset), problem(problem), bisection(bisection) {}

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(LocalVectorType const &ui, int 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(
      LocalVectorType &ui, int m,
      typename Dune::BitSetVector<block_size>::const_reference ignore) {
    {
      int ignore_component =
          block_size; // Special value that indicates nothing should be ignored
      switch (ignore.count()) {
        case 0: // Full problem
          break;
        case 1: // 1 Dimension is fixed
          assert(
              ignore[1]); // Only the second component is allowed to stay fixed
          ignore_component = 1;
          break;
        case block_size: // Ignore the whole node
          return;
        default:
          assert(false);
      }

      LocalMatrixType const *localA = nullptr;
      LocalVectorType localb(problem.f[m]);

      typename MatrixType::row_type::ConstIterator it;
      typename MatrixType::row_type::ConstIterator end = problem.A[m].end();
      for (it = problem.A[m].begin(); it != end; ++it) {
        int const j = it.index();
        if (j == m)
          localA = &(*it); // localA = A[m][m]
        else
          it->mmv(u[j], localb); // localb -= A[m][j] * u[j]
      }
      assert(localA != nullptr);

      auto const phi = problem.phi.restriction(m);
      Dune::SampleFunctional<block_size> localJ(*localA, localb, phi,
                                                ignore_component);

      LocalVectorType correction;
      Dune::minimise(localJ, ui, parset.get<size_t>("localsolver.steps"),
                     bisection);
    }
  }

private:
  Dune::ParameterTree const &parset;

  // problem data
  MyConvexProblemType const &problem;

  // commonly used minimization stuff
  Bisection bisection;

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

#endif