// Copyright (C) 2004,2005,2006,2007,2010 Michael Creel <michael.creel@uab.es>
//
// 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, see <http://www.gnu.org/licenses/>.

// the functions defined in this file are:
// __bfgsmin_obj: bulletproofed objective function that allows checking for availability of analytic gradient
// __numgradient: numeric gradient, used only if analytic not supplied
// __bisectionstep: fallback stepsize algorithm
// __newtonstep: default stepsize algorithm
// __bfgsmin: the DLD function that does the minimization, to be called from bfgsmin.m

#include <oct.h>
#include <octave/parse.h>
#include <octave/Cell.h>
#include <octave/lo-mappers.h>
#include <float.h>
#include <climits>

#include "error-helpers.h"

int __bfgsmin_obj(double &obj, const std::string f, const octave_value_list f_args, const ColumnVector theta, const int minarg)
{
  octave_value_list f_return, f_args_new;
  int success = 1;
  f_args_new = f_args;
  f_args_new(minarg - 1) = theta;
  bool err;
  SET_ERR (f_return = OCTAVE__FEVAL (f, f_args_new), err);
  if (! err)
    {
      SET_ERR (obj = f_return(0).double_value(), err);
    }
  // bullet-proof the objective function
  if (err)
    {
      warning("__bfgsmin_obj: objective function could not be evaluated - setting to DBL_MAX");
      obj = DBL_MAX;
      success = 0;
    }
  return success;
}


// __numgradient: numeric central difference gradient for bfgs.
// This is the same as numgradient, except the derivative is known to be a vector, it's defined as a column,
// and the finite difference delta is incorporated directly rather than called from a function
int __numgradient(ColumnVector &derivative, const std::string f, const octave_value_list f_args, const int minarg)
{
  double SQRT_EPS, diff, delta, obj_left, obj_right, p;
  int j, test, success = 0;
  ColumnVector parameter = f_args(minarg - 1).column_vector_value();
  int k = parameter.rows();
  ColumnVector g(k);
  SQRT_EPS = sqrt(DBL_EPSILON);
  diff = exp(log(DBL_EPSILON)/3.0);
  // get 1st derivative by central difference
  for (j=0; j<k; j++)
    {
      p = parameter(j);
      // determine delta for finite differencing
      test = (fabs(p) + SQRT_EPS) * SQRT_EPS > diff;
      if (test) delta = (fabs(p) + SQRT_EPS) * SQRT_EPS;
      else delta = diff;
      // right side
      parameter(j) = p + delta;
      success = __bfgsmin_obj(obj_right, f, f_args, parameter, minarg);
      if (! success)
        _p_error ("__numgradient: objective function failed, can't compute numeric gradient");
      // left size
      parameter(j) = p - delta;
      success = __bfgsmin_obj(obj_left, f, f_args, parameter, minarg);
      if (! success)
        _p_error ("__numgradient: objective function failed, can't compute numeric gradient");
      parameter(j) = p;  // restore original parameter for next round
      g(j) = (obj_right - obj_left) / (2.0*delta);
    }
  derivative = g;
  return success;
}

int __bfgsmin_gradient(ColumnVector &derivative, const std::string f, octave_value_list f_args, const ColumnVector theta, const int minarg, int have_analytic_gradient) {
  octave_value_list f_return;
  int k = theta.rows();
  int success;
  ColumnVector g(k);
  Matrix check_gradient(k,1);
  f_args(minarg - 1) = theta;
  if (have_analytic_gradient)
    {
      f_return = OCTAVE__FEVAL (f, f_args, 2);
      g = f_return(1).column_vector_value();
    }
  else
    __numgradient(g, f, f_args, minarg);
  // check that gradient is ok
  check_gradient.column(0) = g;
  if (check_gradient.any_element_is_inf_or_nan ())
    {
      _p_error("__bfgsmin_gradient: gradient contains NaNs or Inf");
      success = 0;
    }
  else success = 1;
  derivative = g;
  return success;
}


// this is the lbfgs direction, used if control has 5 elements
ColumnVector lbfgs_recursion(const int memory, const Matrix sigmas, const Matrix gammas, ColumnVector d)
{
  if (memory == 0)
    {
      const int n = sigmas.columns();
      ColumnVector sig = sigmas.column(n-1);
      ColumnVector gam = gammas.column(n-1);
      // do conditioning if there is any memory
      double cond = gam.transpose()*gam;
      if (cond > 0)
        {
          cond = (sig.transpose()*gam) / cond;
          d = cond*d;
        }
      return d;
    }
  else
    {
      ColumnVector sig = sigmas.column(memory-1);
      ColumnVector gam = gammas.column(memory-1);
      double rho;
      rho = 1.0 / (gam.transpose() * sig);
      double alpha;
      alpha = rho * (sig.transpose() * d);
      d = d - alpha*gam;
      d = lbfgs_recursion(memory - 1, sigmas, gammas, d);
      d = d + (alpha - rho * gam.transpose() * d) * sig;
    }
  return d;
}

// __bisectionstep: fallback stepsize method if __newtonstep fails
int __bisectionstep(double &step, double &obj, const std::string f, const octave_value_list f_args, const ColumnVector x, const ColumnVector dx, const int minarg, const int verbose)
{
  double best_obj, improvement, improvement_0;
  ColumnVector trial;
  // initial values
  best_obj = obj;
  improvement_0 = 0.0;
  step = 1.0;
  trial = x + step*dx;
  __bfgsmin_obj(obj, f, f_args, trial, minarg);
  if (verbose) printf("bisectionstep: trial step: %g  obj value: %g\n", step, obj);
  // this first loop goes until an improvement is found
  while (obj >= best_obj)
    {
      if (step < 2.0*DBL_EPSILON)
        {
          if (verbose) warning("bisectionstep: unable to find improvement, setting step to zero");
          step = 0.0;
          return 0;
        }
      step = 0.5*step;
      trial = x + step*dx;
      __bfgsmin_obj(obj, f, f_args, trial, minarg);
      if (verbose)
        printf("bisectionstep: trial step: %g  obj value: %g  best_value: %g\n", step, obj, best_obj);
    }
  // now keep going until rate of improvement is too low, or reach max trials
  best_obj = obj;
  while (step > 2.0*DBL_EPSILON)
    {
      step = 0.5*step;
      trial = x + step*dx;
      __bfgsmin_obj(obj, f, f_args, trial, minarg);
      if (verbose) printf("bisectionstep: trial step: %g  obj value: %g\n", step, obj);
      // if improved, record new best and try another step
      if (obj < best_obj)
        {
          improvement = best_obj - obj;
          best_obj = obj;
          if (improvement > 0.5*improvement_0) {
            improvement_0 = improvement;
          }
          else break;
        }
      else
        {
          step = 2.0*step; // put it back to best found
          obj = best_obj;
          break;
        }
    }
  return 1;
}

// __newtonstep: default stepsize algorithm
int __newtonstep(double &step, double &obj, const std::string f, const octave_value_list f_args, const ColumnVector x, const ColumnVector dx, const int minarg, const int verbose)
{
  double obj_0, obj_left, obj_right, delta, inv_delta_sq, gradient, hessian;
  int found_improvement = 0;
  obj_0 = obj;
  delta = 0.001; // experimentation shows that this is a good choice
  inv_delta_sq = 1.0 / (delta*delta);
  ColumnVector x_right = x + delta*dx;
  ColumnVector x_left = x  - delta*dx;
  // right
  __bfgsmin_obj(obj_right, f, f_args, x_right, minarg);
  // left
  __bfgsmin_obj(obj_left, f, f_args, x_left, minarg);
  gradient = (obj_right - obj_left) / (2.0*delta);  // take central difference
  hessian =  inv_delta_sq*(obj_right - 2.0*obj_0 + obj_left);
  hessian = fabs(hessian); // ensures we're going in a decreasing direction
  if (hessian < 2.0*DBL_EPSILON) hessian = 1.0; // avoid div by zero
  step = - gradient / hessian;  // hessian inverse gradient: the Newton step
  //      step = (step < 1.0)*step + 1.0*(step >= 1.0); // maximum stepsize is 1.0 - conservative
  // ensure that this is improvement, and if not, fall back to bisection
  __bfgsmin_obj(obj, f, f_args, x + step*dx, minarg);
  if (verbose) printf("newtonstep: trial step: %g  obj value: %g\n", step, obj);
  if (obj > obj_0)
    {
      obj = obj_0;
      if (verbose) warning("__stepsize: no improvement with Newton step, falling back to bisection");
      found_improvement = __bisectionstep(step, obj, f, f_args, x, dx, minarg, verbose);
    }
  else found_improvement = 1;
  if (OCTAVE__MATH__ISNAN (obj)) {
    obj = obj_0;
    if (verbose) warning("__stepsize: objective function crash in Newton step, falling back to bisection");
    found_improvement = __bisectionstep(step, obj, f, f_args, x, dx, minarg, verbose);
  }
  else found_improvement = 1;
  return found_improvement;
}

DEFUN_DLD(__bfgsmin, args, ,"__bfgsmin: backend for bfgs minimization\n\
Users should not use this directly. Use bfgsmin.m instead")
{
  std::string f (args(0).string_value());
  Cell f_args_cell (args(1).cell_value());
  octave_value_list f_args, f_return; // holder for return items

  int max_iters, verbosity, criterion, minarg, convergence, iter, memory,
    gradient_ok, i, j, k, conv_fun = 0, conv_param = 0, conv_grad = 0,
    have_gradient, warnings;
  double func_tol, param_tol, gradient_tol, stepsize, obj_value, obj_in,
    last_obj_value, denominator, test;
  Matrix H, H1, H2;
  ColumnVector thetain, d, g, g_new, p, q, sig, gam;

  // controls
  Cell control (args(2).cell_value());
  max_iters = control(0).int_value();
  if (max_iters == -1) max_iters = INT_MAX;
  verbosity = control(1).int_value();
  criterion = control(2).int_value();
  minarg = control(3).int_value();
  memory = control(4).int_value();
  func_tol = control(5).double_value();
  param_tol = control(6).double_value();
  gradient_tol = control(7).double_value();
  have_gradient = control(8).int_value();

  // want to see warnings?
  warnings = 0;
  if (verbosity == 3) warnings = 1;

  // copy cell contents over to octave_value_list to use feval()
  k = f_args_cell.numel ();
  f_args(k-1); // resize only once
  for (i = 0; i<k; i++) f_args(i) = f_args_cell(i);

  // get the minimization argument
  ColumnVector theta  = f_args(minarg - 1).column_vector_value();
  k = theta.rows();

  // containers for items in limited memory version
  Matrix sigmas(k, memory);
  Matrix gammas(k, memory);
  sigmas.fill(0.0);
  gammas.fill(0.0);

  // initialize things
  convergence = -1; // if this doesn't change, it means that maxiters were exceeded
  thetain = theta;
  H = OCTAVE__IDENTITY_MATRIX(k,k);

  // Initial obj_value
  __bfgsmin_obj(obj_in, f, f_args, theta, minarg);
  if (warnings) printf("initial obj_value %g\n", obj_in);

  __bfgsmin_gradient(g, f, f_args, theta, minarg, have_gradient);
  if (have_gradient)
    {
      // check equality if analytic available
      if (warnings) printf("function claims to provide analytic gradient\n");
      have_gradient = 0;				// force numeric
      __bfgsmin_gradient(g_new, f, f_args, theta, minarg, have_gradient);
      p = g - g_new;
      have_gradient = sqrt(p.transpose() * p) < gradient_tol;
      if (have_gradient && warnings) printf("function claims to provide analytic gradient, and it agrees with numeric - using analytic\n");
      if (!have_gradient && warnings) printf("function claims to provide analytic gradient, but it does not agree with numeric - using numeric\n");
    }

  last_obj_value = obj_in; // initialize, is updated after each iteration
  // MAIN LOOP STARTS HERE
  for (iter = 0; iter < max_iters; iter++)
    {
      if(memory > 0)
        {  // lbfgs
          if (iter < memory) d = lbfgs_recursion(iter, sigmas, gammas, g);
          else d = lbfgs_recursion(memory, sigmas, gammas, g);
          d = -d;
        }
      else
        d = -H*g; // ordinary bfgs
      // convergence tests

      // function convergence
      p = theta+d;
      __bfgsmin_obj(obj_value, f, f_args, p, minarg);
      if (fabs(last_obj_value) > 1.0)	conv_fun=(fabs((obj_value/last_obj_value-1)))<func_tol;
      else conv_fun = fabs(obj_value - last_obj_value) < func_tol;
      // parameter change convergence
      test = sqrt(theta.transpose() * theta);
      if (test > 1.0) conv_param = sqrt(d.transpose() * d) / test < param_tol;
      else conv_param = sqrt(d.transpose() * d) < param_tol;		// Want intermediate results?
      // gradient convergence
      conv_grad = sqrt(g.transpose() * g) < gradient_tol;
      // Are we done?
      if (criterion == 1)
        {
          if (conv_fun && conv_param && conv_grad)
            {
              convergence = 1;
              break;
            }
        }
      else if (conv_fun)
        {
          convergence = 1;
          break;
        }
      // if not done, then take a step
      // stepsize: try (l)bfgs direction, then steepest descent if it fails
      f_args(minarg - 1) = theta;
      obj_value = last_obj_value;
      __newtonstep(stepsize, obj_value, f, f_args, theta, d, minarg, warnings);
      if (stepsize == 0.0)
        {
          // fall back to steepest descent
          if (warnings) warning("bfgsmin: BFGS direction fails, switch to steepest descent");
          d = -g; // try steepest descent
          H = OCTAVE__IDENTITY_MATRIX(k,k); // accompany with Hessian reset, for good measure
          obj_value = last_obj_value;
          __newtonstep(stepsize, obj_value, f, f_args, theta, d, minarg, warnings);
          if (stepsize == 0.0)
            {  // if true, exit, we can't find a direction of descent
              warning("bfgsmin: failure, exiting. Try different start values?");
              f_return(0) = theta;
              f_return(1) = obj_value;
              f_return(2) = -1;
              f_return(3) = iter;
              return f_return;
            }
        }
      p = stepsize*d;
      // Want intermediate results?
      if (verbosity > 1)
        {
          printf("------------------------------------------------\n");
          printf("bfgsmin iteration %d  convergence (f g p): %d %d %d\n", iter, conv_fun, conv_grad, conv_param);
          if (warnings)
            {
              if (memory > 0) printf("Using LBFGS, memory is last %d iterations\n",memory);
            }
          printf("\nfunction value: %g  stepsize: %g  \n\n", last_obj_value, stepsize);
          if (have_gradient) printf("used analytic gradient\n");
          else printf("used numeric gradient\n");
          for (j = 0; j<k; j++) printf("%15.5f %15.5f %15.5f\n",theta(j), g(j), p(j));
        }
      //--------------------------------------------------
      // // Are we done?
      // if (criterion == 1) {
      // 	if (conv_fun && conv_param && conv_grad) {
      // 		convergence = 1;
      // 		break;
      // 	}
      // }
      // else if (conv_fun) {
      // 	convergence = 1;
      // 	break;
      // }
      //-------------------------------------------------- 
      last_obj_value = obj_value;
      theta = theta + p;
      // new gradient
      gradient_ok = __bfgsmin_gradient(g_new, f, f_args, theta, minarg, have_gradient);
      if (memory == 0)
        {  //bfgs?
          // Hessian update if gradient ok
          if (gradient_ok)
            {
              q = g_new-g;
              g = g_new;
              denominator = q.transpose()*p;
              if ((fabs(denominator) < DBL_EPSILON))
                {  // reset Hessian if necessary
                  if (verbosity == 1) printf("bfgsmin: Hessian reset\n");
                  H = OCTAVE__IDENTITY_MATRIX(k,k);
                }
              else
                {
                  H1 = (1.0+(q.transpose() * H * q) / denominator) / denominator
                    * (p * p.transpose());
                  H2 = (p * q.transpose() * H + H*q*p.transpose());
                  H2 = H2 / denominator;
                  H = H + H1 - H2;
                }
            }
          else
            H = OCTAVE__IDENTITY_MATRIX(k,k); // reset hessian if gradient fails
          // then try to start again with steepest descent
        }
      else
        {  // otherwise lbfgs
          // save components for Hessian if gradient ok
          if (gradient_ok)
            {
              sig = p; // change in parameter
              gam = g_new - g; // change in gradient
              g = g_new;
              // shift remembered vectors to the right (forget last)
              for(j = memory - 1; j > 0; j--)
                {
                  for(i = 0; i < k; i++)
                    {
                      sigmas(i,j) = sigmas(i,j-1);
                      gammas(i,j) = gammas(i,j-1);
                    }
                }
              // insert new vectors in left-most column
              for(i = 0; i < k; i++)
                {
                  sigmas(i, 0) = sig(i);
                  gammas(i, 0) = gam(i);
                }
            }
          else
            {
              // failed gradient - loose memory and use previous theta
              sigmas.fill(0.0);
              gammas.fill(0.0);
              theta = theta - p;
            }
        }
    }
  // Want last iteration results?
  if (verbosity > 0)
    {
      printf("------------------------------------------------\n");
      printf("bfgsmin final results: %d iterations\n", iter);
      if (warnings)
        {
          if (memory > 0) printf("Used LBFGS, memory is last %d iterations\n",memory);
        }
      printf("\nfunction value: %g\n\n", obj_value);
      if (convergence == -1)                      printf("NO CONVERGENCE: max iters exceeded\n");
      if ((convergence == 1) & (criterion == 1))  printf("STRONG CONVERGENCE\n");
      if ((convergence == 1) & !(criterion == 1)) printf("WEAK CONVERGENCE\n");
      if (convergence == 2)                       printf("NO CONVERGENCE: algorithm failed\n");
      printf("Function conv %d  Param conv %d  Gradient conv %d\n\n", conv_fun, conv_param, conv_grad);
      if (have_gradient) printf("used analytic gradient\n");
      else printf("used numeric gradient\n");
      printf("          param    gradient (n)          change\n");
      for (j = 0; j<k; j++) printf("%15.5f %15.5f %15.5f\n",theta(j), g(j), d(j));
    }
  f_return(3) = iter;
  f_return(2) = convergence;
  f_return(1) = obj_value;
  f_return(0) = theta;
  return f_return;
}
