/* multilarge_nlinear/dogleg.c
 * 
 * Copyright (C) 2016 Patrick Alken
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or (at
 * your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

#include <config.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_multilarge_nlinear.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_blas.h>

/*
 * This module contains an implementation of the Powell dogleg
 * algorithm for nonlinear optimization problems. This implementation
 * closely follows the following works:
 *
 * [1] H. B. Nielsen, K. Madsen, Introduction to Optimization and
 *     Data Fitting, Informatics and Mathematical Modeling,
 *     Technical University of Denmark (DTU), 2010.
 *
 * [2] J. E. Dennis and H. H. W. Mei, Two new unconstrained optimization
 *     algorithms which use function and gradient values, J. Opt. Theory and
 *     Appl., 28(4), 1979.
 */

typedef struct
{
  size_t n;                  /* number of observations */
  size_t p;                  /* number of parameters */
  gsl_vector *dx_gn;         /* Gauss-Newton step, size p */
  gsl_vector *dx_sd;         /* steepest descent step, size p */
  double norm_Dgn;           /* || D dx_gn || */
  double norm_Dsd;           /* || D dx_sd || */
  double norm_Dinvg;         /* || D^{-1} g || */
  double norm_JDinv2g;       /* || J D^{-2} g || */
  gsl_vector *workp1;        /* workspace, length p */
  gsl_vector *workp2;        /* workspace, length p */
  gsl_vector *workn;         /* workspace, length n */

  /* tunable parameters */
  gsl_multilarge_nlinear_parameters params;
} dogleg_state_t;

#include "common.c"

static void * dogleg_alloc (const void * params, const size_t n, const size_t p);
static void dogleg_free(void *vstate);
static int dogleg_init(const void *vtrust_state, void *vstate);
static int dogleg_preloop(const void * vtrust_state, void * vstate);
static int dogleg_step(const void * vtrust_state, const double delta,
                       gsl_vector * dx, void * vstate);
static int dogleg_double_step(const void * vtrust_state, const double delta,
                              gsl_vector * dx, void * vstate);
static int dogleg_preduction(const void * vtrust_state, const gsl_vector * dx,
                             double * pred, void * vstate);
static int dogleg_calc_gn(const gsl_multilarge_nlinear_trust_state * trust_state, gsl_vector * dx);
static double dogleg_beta(const double t, const double delta,
                          const gsl_vector * diag, dogleg_state_t * state);

static void *
dogleg_alloc (const void * params, const size_t n, const size_t p)
{
  dogleg_state_t *state;
  
  state = calloc(1, sizeof(dogleg_state_t));
  if (state == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate dogleg state", GSL_ENOMEM);
    }

  state->dx_gn = gsl_vector_alloc(p);
  if (state->dx_gn == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate space for dx_gn", GSL_ENOMEM);
    }

  state->dx_sd = gsl_vector_alloc(p);
  if (state->dx_sd == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate space for dx_sd", GSL_ENOMEM);
    }

  state->workp1 = gsl_vector_alloc(p);
  if (state->workp1 == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate space for workp1", GSL_ENOMEM);
    }

  state->workp2 = gsl_vector_alloc(p);
  if (state->workp2 == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate space for workp2", GSL_ENOMEM);
    }

  state->workn = gsl_vector_alloc(n);
  if (state->workn == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate space for workn", GSL_ENOMEM);
    }

  state->n = n;
  state->p = p;
  state->params = *(const gsl_multilarge_nlinear_parameters *) params;

  return state;
}

static void
dogleg_free(void *vstate)
{
  dogleg_state_t *state = (dogleg_state_t *) vstate;

  if (state->dx_gn)
    gsl_vector_free(state->dx_gn);

  if (state->dx_sd)
    gsl_vector_free(state->dx_sd);

  if (state->workp1)
    gsl_vector_free(state->workp1);

  if (state->workp2)
    gsl_vector_free(state->workp2);

  if (state->workn)
    gsl_vector_free(state->workn);

  free(state);
}

/*
dogleg_init()
  Initialize dogleg solver

Inputs: vtrust_state - trust state
        vstate       - workspace

Return: success/error
*/

static int
dogleg_init(const void *vtrust_state, void *vstate)
{
  (void)vtrust_state;
  (void)vstate;

  return GSL_SUCCESS;
}

/*
dogleg_preloop()
  Initialize dogleg method prior to iteration loop.
This involves computing the steepest descent step. The
Gauss-Newton step is computed later in the _step() functions
if required.

Notes: on output,
1) state->dx_sd contains steepest descent step
2) state->norm_Dinvg contains || D^{-1} g ||
3) state->norm_JDinv2g contains || J D^{-2} g ||
*/

static int
dogleg_preloop(const void * vtrust_state, void * vstate)
{
  const gsl_multilarge_nlinear_trust_state *trust_state =
    (const gsl_multilarge_nlinear_trust_state *) vtrust_state;
  dogleg_state_t *state = (dogleg_state_t *) vstate;
  double u;
  double alpha; /* ||g||^2 / ||Jg||^2 */

  /* calculate the steepest descent step */

  /* compute workp1 = D^{-1} g and its norm */
  gsl_vector_memcpy(state->workp1, trust_state->g);
  gsl_vector_div(state->workp1, trust_state->diag);
  state->norm_Dinvg = gsl_blas_dnrm2(state->workp1);

  /* compute workp1 = D^{-2} g */
  gsl_vector_div(state->workp1, trust_state->diag);

  /* compute workp2 = J^T J D^{-2} g */
  gsl_blas_dsymv(CblasLower, 1.0, trust_state->JTJ, state->workp1, 0.0, state->workp2);

  /* compute norm_JDinv2g = || J D^{-2} g || */
  gsl_blas_ddot(state->workp1, state->workp2, &u);
  state->norm_JDinv2g = sqrt(u);

  u = state->norm_Dinvg / state->norm_JDinv2g;
  alpha = u * u;

  /* dx_sd = -alpha D^{-2} g */
  gsl_vector_memcpy(state->dx_sd, state->workp1);
  gsl_vector_scale(state->dx_sd, -alpha);

  state->norm_Dsd = scaled_enorm(trust_state->diag, state->dx_sd);
  state->norm_Dgn = -1.0; /* computed later if needed */

  return GSL_SUCCESS;
}

/*
dogleg_step()
  Calculate a new step vector
*/

static int
dogleg_step(const void * vtrust_state, const double delta,
            gsl_vector * dx, void * vstate)
{
  const gsl_multilarge_nlinear_trust_state *trust_state =
    (const gsl_multilarge_nlinear_trust_state *) vtrust_state;
  dogleg_state_t *state = (dogleg_state_t *) vstate;

  if (state->norm_Dsd >= delta)
    {
      /* steepest descent step is outside trust region;
       * truncate steepest descent step to trust region boundary */
      gsl_vector_memcpy(dx, state->dx_sd);
      gsl_vector_scale(dx, delta / state->norm_Dsd);
    }
  else
    {
      /* compute Gauss-Newton step if needed */
      if (state->norm_Dgn < 0.0)
        {
          int status = dogleg_calc_gn(trust_state, state->dx_gn);

          if (status)
            return status;

          /* compute || D dx_gn || */
          state->norm_Dgn = scaled_enorm(trust_state->diag, state->dx_gn);
        }

      if (state->norm_Dgn <= delta)
        {
          /* Gauss-Newton step is inside trust region, use it as final step
           * since it is the global minimizer of the quadratic model function */
          gsl_vector_memcpy(dx, state->dx_gn);
        }
      else
        {
          /* Gauss-Newton step is outside trust region, but steepest
           * descent is inside; use dogleg step */

          double beta = dogleg_beta(1.0, delta, trust_state->diag, state);

          /* compute: workp1 = dx_gn - dx_sd */
          scaled_addition(1.0, state->dx_gn, -1.0, state->dx_sd, state->workp1);

          /* dx = dx_sd + beta*(dx_gn - dx_sd) */
          scaled_addition(beta, state->workp1, 1.0, state->dx_sd, dx);
        }
    }

  return GSL_SUCCESS;
}

/*
dogleg_double_step()
  Calculate a new step with double dogleg method. Based on
section 3 of [2]
*/

static int
dogleg_double_step(const void * vtrust_state, const double delta,
                   gsl_vector * dx, void * vstate)
{
  const double alpha_fac = 0.8; /* recommended value from Dennis and Mei */
  const gsl_multilarge_nlinear_trust_state *trust_state =
    (const gsl_multilarge_nlinear_trust_state *) vtrust_state;
  dogleg_state_t *state = (dogleg_state_t *) vstate;

  if (state->norm_Dsd >= delta)
    {
      /* steepest descent step is outside trust region;
       * truncate steepest descent step to trust region boundary */
      gsl_vector_memcpy(dx, state->dx_sd);
      gsl_vector_scale(dx, delta / state->norm_Dsd);
    }
  else
    {
      /* compute Gauss-Newton step if needed */
      if (state->norm_Dgn < 0.0)
        {
          int status = dogleg_calc_gn(trust_state, state->dx_gn);

          if (status)
            return status;

          /* compute || D dx_gn || */
          state->norm_Dgn = scaled_enorm(trust_state->diag, state->dx_gn);
        }

      if (state->norm_Dgn <= delta)
        {
          /* Gauss-Newton step is inside trust region, use it as final step
           * since it is the global minimizer of the quadratic model function */
          gsl_vector_memcpy(dx, state->dx_gn);
        }
      else
        {
          double t, u, v, c;

          /* compute: u = ||D^{-1} g||^2 / ||J D^{-2} g||^2 */
          v = state->norm_Dinvg / state->norm_JDinv2g;
          u = v * v;

          /* compute: v = g^T dx_gn */
          gsl_blas_ddot(trust_state->g, state->dx_gn, &v);

          /* compute: c = ||D^{-1} g||^4 / (||J D^{-2} g||^2 * |g^T dx_gn|) */
          c = u * (state->norm_Dinvg / fabs(v)) * state->norm_Dinvg;

          /* compute: t = 1 - alpha_fac*(1-c) */
          t = 1.0 - alpha_fac*(1.0 - c);

          if (t * state->norm_Dgn <= delta)
            {
              /* set dx = (delta / ||D dx_gn||) dx_gn */
              gsl_vector_memcpy(dx, state->dx_gn);
              gsl_vector_scale(dx, delta / state->norm_Dgn);
            }
          else
            {
              /* Cauchy point is inside, Gauss-Newton is outside trust region;
               * use double dogleg step */

              double beta = dogleg_beta(t, delta, trust_state->diag, state);

              /* compute: workp1 = t*dx_gn - dx_sd */
              scaled_addition(t, state->dx_gn, -1.0, state->dx_sd, state->workp1);

              /* dx = dx_sd + beta*(t*dx_gn - dx_sd) */
              scaled_addition(beta, state->workp1, 1.0, state->dx_sd, dx);
            }
        }
    }

  return GSL_SUCCESS;
}

static int
dogleg_preduction(const void * vtrust_state, const gsl_vector * dx,
                  double * pred, void * vstate)
{
  const gsl_multilarge_nlinear_trust_state *trust_state =
    (const gsl_multilarge_nlinear_trust_state *) vtrust_state;
  dogleg_state_t *state = (dogleg_state_t *) vstate;

  *pred = quadratic_preduction(trust_state, dx, state->workn);

  return GSL_SUCCESS;
}

/*
dogleg_calc_gn()
  Calculate Gauss-Newton step by solving

J^T J dx_gn = -J^T f

Inputs: trust_state - trust state variables
        dx          - (output) Gauss-Newton step vector

Return: success/error
*/

static int
dogleg_calc_gn(const gsl_multilarge_nlinear_trust_state * trust_state, gsl_vector * dx)
{
  int status;
  const gsl_multilarge_nlinear_parameters *params = trust_state->params;

  /* initialize linear least squares solver */
  status = (params->solver->init)(trust_state, trust_state->solver_state);
  if (status)
    return status;

  /* prepare the linear solver to compute Gauss-Newton step */
  status = (params->solver->presolve)(0.0, trust_state, trust_state->solver_state);
  if (status)
    return status;

  /* solve: J dx_gn = -f for Gauss-Newton step */
  status = (params->solver->solve)(trust_state->g,
                                   dx,
                                   trust_state,
                                   trust_state->solver_state);
  if (status)
    return status;

  return GSL_SUCCESS;
}

/*
dogleg_beta()
  This function finds beta in [0,1] such that the step

dx = dx_sd + beta*(t*dx_gn - dx_sd)

has norm

||D dx|| = delta

beta is the positive root of the quadratic:

a beta^2 + b beta + c = 0

with

a = ||D(t*dx_gn - dx_sd)||^2
b = 2 dx_sd^T D^T D (t*dx_gn - dx_sd)
c = ||D dx_sd||^2 - delta^2 

Inputs: t     - amount of Gauss-Newton step to use for dogleg
                (= 1 for classical dogleg, <= 1 for double dogleg)
        delta - trust region radius
        diag  - diag(D) scaling matrix
        state - workspace
*/

static double
dogleg_beta(const double t, const double delta,
            const gsl_vector * diag, dogleg_state_t * state)
{
  double beta;
  double a, b, c;

  /* compute: workp1 = t*dx_gn - dx_sd */
  scaled_addition(t, state->dx_gn, -1.0, state->dx_sd, state->workp1);

  /* a = || D (t*dx_gn - dx_sd) ||^2 */
  a = scaled_enorm(diag, state->workp1);
  a *= a;

  /* workp1 = D^T D (t*dx_gn - dx_sd) */
  gsl_vector_mul(state->workp1, diag);
  gsl_vector_mul(state->workp1, diag);

  /* b = 2 dx_sd^T D^T D (t*dx_gn - dx-sd) */
  gsl_blas_ddot(state->dx_sd, state->workp1, &b);
  b *= 2.0;

  /* c = || D dx_sd ||^2 - delta^2 = (||D dx_sd|| + delta) (||D dx_sd|| - delta) */
  c = (state->norm_Dsd + delta) * (state->norm_Dsd - delta);

  if (b > 0.0)
    {
      beta = (-2.0 * c) / (b + sqrt(b*b - 4.0*a*c));
    }
  else
    {
      beta = (-b + sqrt(b*b - 4.0*a*c)) / (2.0 * a);
    }

  return beta;
}

static const gsl_multilarge_nlinear_trs dogleg_type =
{
  "dogleg",
  dogleg_alloc,
  dogleg_init,
  dogleg_preloop,
  dogleg_step,
  dogleg_preduction,
  dogleg_free
};

const gsl_multilarge_nlinear_trs *gsl_multilarge_nlinear_trs_dogleg = &dogleg_type;

static const gsl_multilarge_nlinear_trs ddogleg_type =
{
  "double-dogleg",
  dogleg_alloc,
  dogleg_init,
  dogleg_preloop,
  dogleg_double_step,
  dogleg_preduction,
  dogleg_free
};

const gsl_multilarge_nlinear_trs *gsl_multilarge_nlinear_trs_ddogleg = &ddogleg_type;