##  Copyright (C) 2013-2015  Daniel Kraft <d@domob.eu>
##  GNU Octave level-set package.
##
##  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/>.

## -*- texinfo -*-
## @deftypefn  {Function File} {[@var{s}, @var{t}] =} so_step_armijo (@var{t0}, @var{d}, @var{f}, @var{dJ}, @var{data})
##
## Perform a line search according to the Armijo rule.  A backtracking
## strategy is employed, such that the returned step length @var{t}
## satisfies
## @tex
## \begin{equation*}
##  J(t) \le J(0) + t \tau \cdot dJ.
## \end{equation*}
## @end tex
## @ifnottex
##
## @example
## J(t) <= J(0) + t * tau * dJ.
## @end example
##
## @end ifnottex
##
## @var{t0} defines the initial step length (see below).  It can be set to
## some suitable constant or to the last successful step length.
## @var{d} should be the result of @code{ls_solve_stationary} for
## the current geometry, and @var{f} the matching speed field on the grid.
## @var{dJ} must be the directional shape derivative of the cost functional
## in direction @var{f}.  This value must be negative.
##
## Data defining the problem, the current state and parameters
## should be passed in @var{data}.  See @code{so_run_descent} for a
## description of this struct.
##
## Returned is the final state struct in @var{s} and the corresponding
## step length taken in @var{t}.
##
## The following parameters must be set in @code{@var{data}.p.lineSearch}
## to determine how the Armijo line search is done:
##
## @table @code
## @item relaxation
## The parameter tau in the Armijo rule.
##
## @item backtrack
## Backtracking factor.  Should be less than one.
##
## @item initial
## Factor for increasing @var{t0} to get the initial guess.  This should
## be larger than one in order to allow increasing of the step
## length when this is possible.
##
## @item minStep
## Minimum step length.  If backtracking goes below this value, use
## @code{minStep} and disregard the condition.
## @end table
##
## Furthermore, the following general parameters in @code{@var{data}.p}
## also influence the behaviour of the line search:
##
## @table @code
## @item verbose
## Whether or not log chatter should be printed.
##
## @item nProc
## Number of parallel threads to use for the line search.  If this
## is larger than one, multiple trial steps are computed in parallel to
## speed up the computation.  For this to work, @code{pararrayfun}
## of the @code{parallel} package must be available.
## @end table
##
## The routine @code{so_init_params} can be used to initialise the parameter
## structure with default values.
## 
## @seealso{so_init_params, so_run_descent, ls_solve_stationary}
## @end deftypefn

function [s, t] = so_step_armijo (t0, d, f, dJ, data)
  if (nargin () != 5)
    print_usage ();
  endif

  if (!isscalar (t0) || !isscalar (dJ) || !isstruct (data))
    print_usage ();
  endif

  if (dJ >= 0)
    error ("DJ should be negative");
  endif

  p = data.p.lineSearch;
  if (p.relaxation <= 0 || p.relaxation >= 1)
    error ("invalid Armijo relaxation parameter");
  endif
  if (p.backtrack <= 0 || p.backtrack >= 1)
    error ("invalid backtracking factor");
  endif
  if (p.initial <= 1)
    error ("invalid initial increase factor");
  endif
  if (p.minStep < 0)
    error ("invalid minimum step length");
  endif

  sz = size (d);
  if (!all (sz == size (f)))
    error ("sizes mismatch between D and F");
  endif
  if (!all (sz == size (data.s.phi)))
    error ("sizes mismatch between D and PHI");
  endif

  t = t0 * p.initial;
  while (true)

    % Use parcellfun to check a couple of steps in parallel.
    parTs = t * p.backtrack.^(0 : data.p.nProc - 1);
    errorHandler = @(err, curT) deal (err, NA);
    handler = @(curT) getCostForStep (curT, d, f, data);

    if (length (parTs) == 1)
      [costps, sps] = handler (parTs);
      costps = {costps};
      sps = {sps};
    else
      [costps, sps] = pararrayfun (data.p.nProc, handler, parTs, ...
                                   "UniformOutput", false, ...
                                   "ErrorHandler", errorHandler);
    endif

    % Process the results looking for a step that is ok.
    for i = 1 : length (parTs)
      if (isstruct (costps{i}))
        error ("Subprocess error: %s", costps{i}.message);
      endif
      if (data.p.verbose)
        printf ("Armijo step %.6f: cost = %.6f\n", parTs(i), costps{i});
      endif

      % Check for minimum step length.  Note that it may happen
      % that even this step is "too long", eliminating the domain.
      % In that case, error out.
      if (parTs(i) < p.minStep)
        warning ("level-set:shape-optimisation:no-armijo-step", ...
                 "no step found in line search until minimum %.6f", p.minStep);

        t = p.minStep;
        [cost, s] = getCostForStep (p.minStep, d, f, data);
        if (isinf (cost))
          error ("minimum step length results in invalid state");
        endif

        return;
      endif

      if (costps{i} <= data.s.cost + parTs(i) * p.relaxation * dJ)
        t = parTs(i);
        s = sps{i};
        return;
      endif
    endfor

    % Decrease step length.
    t *= p.backtrack .^ data.p.nProc;
  endwhile
endfunction

% Helper routine to calculate costs given a step length.
function [costp, sp] = getCostForStep (t, d, f, data)
  phip = ls_extract_solution (t, d, data.s.phi, f);

  % Do not allow steps that eliminate the domain!
  if (ls_isempty (phip))
    costp = Inf;
    sp = NA;
    return;
  endif

  sp = data.cb.update_state (phip, data);
  costp = sp.cost;
endfunction

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Tests.

% Cost function for our example problem.
%!function s = updateState (phi, data)
%!  inside = ls_inside (phi);
%!  vol = data.g.h * length (find (inside));
%!  cost = 1 / vol;
%!
%!  % Throw if the volume is too large.  This is used
%!  % to check the error handling code.
%!  if (vol > 15)
%!    error ("volume too large");
%!  endif
%!
%!  s = struct ("phi", phi, "cost", cost);
%!endfunction

% Define default parameters and phi used in the test later on.
%!shared n, phi, baseData, d, f, dJ
%!  n = 100;
%!  x = linspace (-10, 10, n);
%!  h = x(2) - x(1);
%!  a = 0;
%!  b = 5;
%!  phi = ls_genbasic (x, "box", a, b);
%!
%!  baseData = struct ();
%!  baseData.p = so_init_params (false);
%!  baseData.g = struct ("x", x, "h", h);
%!  baseData.cb = struct ("update_state", @updateState);
%!  baseData.s = updateState (phi, baseData);
%!  baseData.s.phi = phi;
%!
%!  f = ones (size (phi));
%!  d = ls_solve_stationary (phi, f, h);
%!
%!  curVol = b - a;
%!  dJ = -2 / curVol^2;

% Test for error conditions.
%!error <Invalid call to>
%!  so_step_armijo (1, 2, 3, 4)
%!error <Invalid call to>
%!  so_step_armijo (1, 2, 3, 4, 5, 6)
%!error <Invalid call to>
%!  so_step_armijo ([1, 1], 2, 3, 4, struct ())
%!error <Invalid call to>
%!  so_step_armijo (1, 2, 3, [4, 4], struct ())
%!error <Invalid call to>
%!  so_step_armijo (1, 2, 3, 4, 5)
%!error <DJ should be negative>
%!  so_step_armijo (1, d, f, 4, baseData)
%!error <sizes mismatch between D and F>
%!  so_step_armijo (1, 2, f, dJ, baseData)
%!error <sizes mismatch between D and PHI>
%!  so_step_armijo (1, 2, 3, dJ, baseData)

% Test parameter validation.
%!error <invalid Armijo relaxation parameter>
%!  data = baseData;
%!  data.p.lineSearch.relaxation = 0;
%!  so_step_armijo (1, d, f, dJ, data);
%!error <invalid Armijo relaxation parameter>
%!  data = baseData;
%!  data.p.lineSearch.relaxation = 1;
%!  so_step_armijo (1, d, f, dJ, data);
%!error <invalid backtracking factor>
%!  data = baseData;
%!  data.p.lineSearch.backtrack = 0;
%!  so_step_armijo (1, d, f, dJ, data);
%!error <invalid backtracking factor>
%!  data = baseData;
%!  data.p.lineSearch.backtrack = 1;
%!  so_step_armijo (1, d, f, dJ, data);
%!error <invalid initial increase factor>
%!  data = baseData;
%!  data.p.lineSearch.initial = 1;
%!  so_step_armijo (1, d, f, dJ, data);
%!error <invalid minimum step length>
%!  data = baseData;
%!  data.p.lineSearch.minStep = -1;
%!  so_step_armijo (1, d, f, dJ, data);
%!test
%!  data = baseData;
%!  data.p.lineSearch.minStep = 0;
%!  so_step_armijo (1, d, f, dJ, data);

% Test that backtracking works.  Set the parameters such that
% backtracking is actually required.
%!test
%!  data = baseData;
%!  data.p.verbose = true;
%!  data.p.lineSearch.relaxation = 0.9;
%!  [~, t] = so_step_armijo (1, d, f, dJ, data);
%!  assert (t < 0.5);

% Test error handling.  With a large step, we trigger
% the "volume too large" error above.  Verify that the
% error is handled both with and without multi-threading.
%
% For evolution time t, the resulting volume will be 5 + 4 t.
% Thus the volume constraint is hit at t = 2.5.
%
% Note that we don't get 'Subprocess error' prefixed.  This is
% presumably due to how Octave handles the error checks in
% combination with the forks induced by multi-threading.
%
%!test
%!  so_step_armijo (2, d, f, dJ, baseData);
%!error <volume too large>
%!  so_step_armijo (3, d, f, dJ, baseData);
%!error <volume too large>
%!  pkg load parallel;
%!  data = baseData;
%!  data.p.nProc = 2;
%!  so_step_armijo (3, d, f, dJ, data);
