// ---------------------------------------------------------------------
//
// Copyright (C) 1998 - 2017 by the deal.II authors
//
// This file is part of the deal.II library.
//
// The deal.II library is free software; you can use it, redistribute
// it, and/or modify it under the terms of the GNU Lesser General
// Public License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
// The full text of the license can be found in the file LICENSE at
// the top level of the deal.II distribution.
//
// ---------------------------------------------------------------------

#ifndef dealii_solver_fire_h
#define dealii_solver_fire_h


#include <deal.II/base/config.h>
#include <deal.II/base/logstream.h>
#include <deal.II/lac/diagonal_matrix.h>
#include <deal.II/lac/solver.h>

#include <functional>


DEAL_II_NAMESPACE_OPEN


/*!@addtogroup Solvers */
/*@{*/

/**
 * FIRE (Fast Inertial Relaxation Engine) for minimization of (potentially
 * non-linear) objective function $E(\mathbf x)$, $\mathbf x$ is a vector of
 * $n$ variables ($n$ is the number of variables of the objective function).
 * Like all other solver classes, it can work on any kind of vector and matrix
 * as long as they satisfy certain requirements (for the requirements on
 * matrices and vectors in order to work with this class, see the documentation
 * of the Solver base class). The type of the solution vector must be passed as
 * template argument, and defaults to dealii::Vector<double>.
 *
 * FIRE is a damped dynamics method described in
 * <a href="https://doi.org/10.1103/PhysRevLett.97.170201">Structural
 * Relaxation Made Simple</a> by Bitzek et al. 2006, typically used to find
 * stable equilibrium configurations of atomistic systems in computational
 * material science. Starting from a given initial configuration of the
 * atomistic system, the algorithm relies on inertia to obtain (nearest)
 * configuration with least potential energy.
 *
 * Notation:
 *  - The global vector of unknown variables: $\mathbf x$.
 *  - Objective function:                     $E(\mathbf x)$.
 *  - Rate of change of unknowns:             $\mathbf v$.
 *  - Gradient of the objective
 *    function w.r.t unknowns:                $\mathbf g = \nabla E(\mathbf x)$.
 *  - Mass matrix:                            $\mathbf M$.
 *  - Initial guess of unknowns:              $\mathbf x_0$.
 *  - Time step:                              $\Delta t$.
 *
 * Given initial values for $\Delta t$, $\alpha = \alpha_0$, $\epsilon$,
 * $\mathbf x = \mathbf x_0$ and $\mathbf v= \mathbf 0$ along with a given mass
 * matrix $\mathbf M$, FIRE algorithm is as follows,
 * 1. Calculate $\mathbf g = \nabla E(\mathbf x)$ and check for convergence
 *    ($\mathbf g \cdot \mathbf g < \epsilon^2 $).
 * 2. Update $\mathbf x$ and $V$ using simple (forward) Euler integration step,
 *    <BR>
 *        $\mathbf x = \mathbf x + \Delta t \mathbf v$,                 <BR>
 *        $\mathbf v = \mathbf v + \Delta t \mathbf M^{-1} \cdot \mathbf g$.
 * 3. Calculate $p = \mathbf g \cdot \mathbf v$.
 * 4. Set $\mathbf v = (1-\alpha) \mathbf v
 *                   + \alpha \frac{|\mathbf v|}{|\mathbf g|} \mathbf g$.
 * 5. If $p<0$ and number of steps since $p$ was last negative is larger
 *    than certain value, then increase time step $\Delta t$ and decrease
 *    $\alpha$.
 * 6. If $p>0$, then decrease the time step, freeze the system i.e.,
 *    $\mathbf v = \mathbf 0$ and reset $\alpha = \alpha_0$.
 * 7. Return to 1.
 *
 * Also see
 * <a href="http://onlinelibrary.wiley.com/doi/10.1002/pamm.201110246/full">
 * Energy-Minimization in Atomic-to-Continuum Scale-Bridging Methods </a> by
 * Eidel et al. 2011.
 *
 * @author Vishal Boddu, Denis Davydov, 2017
 */
template <typename VectorType = Vector<double> >
class SolverFIRE : public Solver<VectorType>
{

public:

  /**
   * Standardized data struct to pipe additional data to the solver.
   */
  struct AdditionalData
  {
    /**
     * Constructor. By default, set the initial time step for the (forward)
     * Euler integration step to 0.1, the maximum time step to 1 and the
     * maximum change allowed in any variable (per iteration) to 1.
     */
    explicit
    AdditionalData (const double  initial_timestep    = 0.1,
                    const double  maximum_timestep    = 1,
                    const double  maximum_linfty_norm = 1);

    /**
     * Initial time step for the (forward) Euler integration step.
     */
    const double initial_timestep;

    /**
     * Maximum time step for the (forward) Euler integration step.
     */
    const double maximum_timestep;

    /**
     * Maximum change allowed in any variable of the objective function.
     */
    const double maximum_linfty_norm;

  };

  /**
   * Constructor.
   */
  SolverFIRE (SolverControl            &solver_control,
              VectorMemory<VectorType> &vector_memory,
              const AdditionalData     &data = AdditionalData());

  /**
   * Constructor. Use an object of type GrowingVectorMemory as a default to
   * allocate memory.
   */
  SolverFIRE (SolverControl         &solver_control,
              const AdditionalData  &data          );

  /**
   * Virtual destructor.
   */
  virtual ~SolverFIRE();

  /**
   * Obtain a set of variables @p x that minimize an objective function
   * described by the polymorphic function wrapper @p compute, with a given
   * preconditioner @p inverse_mass_matrix and initial @p x values.
   * The function @p compute returns the objective function's value and updates
   * the objective function's gradient (with respect to the variables) when
   * passed in as first argument based on the second argument-- the state of
   * variables.
   */
  template <typename PreconditionerType = DiagonalMatrix<VectorType>>
  void solve
  (const std::function<double(VectorType &, const VectorType &)> &compute,
   VectorType                                                    &x,
   const PreconditionerType                                      &inverse_mass_matrix);

  /**
   * Solve for x that minimizes $E(\mathbf x)$ for the <EM>special case</EM>
   * when $E(\mathbf x)
   * = \frac{1}{2} \mathbf x^{T} \mathbf A \mathbf x - \mathbf x^{T} \mathbf b$.
   */
  template <typename MatrixType, typename PreconditionerType>
  void solve (const MatrixType         &A,
              VectorType               &x,
              const VectorType         &b,
              const PreconditionerType &preconditioner);

protected:

  /**
   * Interface for derived class. This function gets the current iteration
   * @p x (variables), @p v (x's time derivative) and @p g (the gradient) in
   * each step.
   * It can be used for graphical output of the convergence history.
   */
  virtual void print_vectors (const unsigned int,
                              const VectorType &x,
                              const VectorType &v,
                              const VectorType &g) const;

  /**
   * Additional data to the solver.
   */
  const AdditionalData additional_data;

};

/*@}*/

/*------------------------- Implementation ----------------------------*/

#ifndef DOXYGEN

template <typename VectorType>
SolverFIRE<VectorType>::AdditionalData::
AdditionalData (const double  initial_timestep,
                const double  maximum_timestep,
                const double  maximum_linfty_norm)
  :
  initial_timestep(initial_timestep),
  maximum_timestep(maximum_timestep),
  maximum_linfty_norm(maximum_linfty_norm)
{
  AssertThrow (initial_timestep    > 0. &&
               maximum_timestep    > 0. &&
               maximum_linfty_norm > 0.,
               ExcMessage("Expected positive values for initial_timestep, "
                          "maximum_timestep and maximum_linfty_norm but one "
                          "or more of the these values are not positive."));
}



template <typename VectorType>
SolverFIRE<VectorType>::
SolverFIRE (SolverControl            &solver_control,
            VectorMemory<VectorType> &vector_memory,
            const AdditionalData     &data          )
  :
  Solver<VectorType>(solver_control, vector_memory),
  additional_data(data)
{}



template <typename VectorType>
SolverFIRE<VectorType>::
SolverFIRE (SolverControl         &solver_control,
            const AdditionalData  &data          )
  :
  Solver<VectorType>(solver_control),
  additional_data(data)
{}



template <typename VectorType>
SolverFIRE<VectorType>::~SolverFIRE()
{}



template <typename VectorType>
template <typename PreconditionerType>
void
SolverFIRE<VectorType>::solve
(const std::function<double(VectorType &, const VectorType &)> &compute,
 VectorType                                                    &x,
 const PreconditionerType                                      &inverse_mass_matrix)
{
  LogStream::Prefix prefix("FIRE");

  // FIRE algorithm constants
  const double DELAYSTEP       = 5;
  const double TIMESTEP_GROW   = 1.1;
  const double TIMESTEP_SHRINK = 0.5;
  const double ALPHA_0         = 0.1;
  const double ALPHA_SHRINK    = 0.99;

  using real_type = typename VectorType::real_type;

  typename VectorMemory<VectorType>::Pointer v(this->memory);
  typename VectorMemory<VectorType>::Pointer g(this->memory);

  // Set velocities to zero but not gradients
  // as we are going to compute them soon.
  v->reinit(x,false);
  g->reinit(x,true);

  // Refer to v and g with some readable names.
  VectorType &velocities = *v;
  VectorType &gradients  = *g;

  // Update gradients for the new x.
  compute(gradients, x);

  unsigned int iter = 0;

  SolverControl::State conv = SolverControl::iterate;
  conv = this->iteration_status (iter, gradients * gradients, x);
  if (conv != SolverControl::iterate)
    return;

  // Refer to additional data members with some readable names.
  const auto &maximum_timestep   = additional_data.maximum_timestep;
  double timestep                = additional_data.initial_timestep;

  // First scaling factor.
  double alpha = ALPHA_0;

  unsigned int previous_iter_with_positive_v_dot_g = 0;

  while (conv == SolverControl::iterate)
    {
      ++iter;
      // Euler integration step.
      x.add (timestep, velocities);                    // x += dt     * v
      inverse_mass_matrix.vmult(gradients, gradients); // g  = M^{-1} * g
      velocities.add (-timestep, gradients);           // v -= dt     * h

      // Compute gradients for the new x.
      compute(gradients, x);

      const real_type gradient_norm_squared = gradients * gradients;
      conv = this->iteration_status(iter, gradient_norm_squared, x);
      if (conv != SolverControl::iterate)
        break;

      // v_dot_g = V * G
      const real_type v_dot_g = velocities * gradients;

      if (v_dot_g < 0.)
        {
          const real_type velocities_norm_squared =
            velocities * velocities;

          // Check if we divide by zero in DEBUG mode.
          Assert (gradient_norm_squared > 0., ExcInternalError());

          // beta = - alpha |V|/|G|
          const real_type beta = -alpha *
                                 std::sqrt (velocities_norm_squared
                                            /
                                            gradient_norm_squared);

          // V = (1-alpha) V + beta G.
          velocities.sadd (1. - alpha, beta, gradients);

          if (iter - previous_iter_with_positive_v_dot_g > DELAYSTEP)
            {
              // Increase timestep and decrease alpha.
              timestep = std::min (timestep*TIMESTEP_GROW, maximum_timestep);
              alpha *= ALPHA_SHRINK;
            }
        }
      else
        {
          // Decrease timestep, reset alpha and set V = 0.
          previous_iter_with_positive_v_dot_g = iter;
          timestep *= TIMESTEP_SHRINK;
          alpha = ALPHA_0;
          velocities = 0.;
        }

      real_type vmax = velocities.linfty_norm();

      // Change timestep if any dof would move more than maximum_linfty_norm.
      if (vmax > 0.)
        {
          const double minimal_timestep = additional_data.maximum_linfty_norm
                                          /
                                          vmax;
          if (minimal_timestep < timestep)
            timestep = minimal_timestep;
        }

      print_vectors(iter, x, velocities, gradients);

    } // While we need to iterate.

  // In the case of failure: throw exception.
  if (conv != SolverControl::success)
    AssertThrow (false,
                 SolverControl::NoConvergence (iter, gradients * gradients));

}



template <typename VectorType>
template <typename MatrixType, typename PreconditionerType>
void SolverFIRE<VectorType>::solve (const MatrixType         &A,
                                    VectorType               &x,
                                    const VectorType         &b,
                                    const PreconditionerType &preconditioner)
{

  std::function<double(VectorType &,  const VectorType &)> compute_func =
    [&]               (VectorType &g, const VectorType &x) -> double
  {
    // Residual of the quadratic form $ \frac{1}{2} xAx - xb $.
    // G = b - Ax
    A.residual(g, x, b);

    // Gradient G = Ax -b.
    g *= -1.;

    // The quadratic form $\frac{1}{2} xAx - xb $.
    return 0.5*A.matrix_norm_square(x) - x*b;
  };

  this->solve (compute_func, x, preconditioner);
}



template <typename VectorType>
void
SolverFIRE<VectorType>::print_vectors (const unsigned int,
                                       const VectorType &,
                                       const VectorType &,
                                       const VectorType &) const
{}



#endif // DOXYGEN

DEAL_II_NAMESPACE_CLOSE

#endif
