##  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{log}] =} so_run_descent (@var{nSteps}, @var{phi0}, @var{data})
##
## Run a descent method for shape optimisation.  This routine is quite
## general, since the actual ``descent direction'' used must be computed
## by a callback (see below).  It is thus not just for steepest descent.
## At each step, a line search according to @code{so_step_armijo} is employed.
##
## @var{nSteps} defines the maximal number of descent steps to take.  After
## this number, the method returns no matter whether or not the stopping
## criterion is fulfilled.  @var{phi0} is the initial geometry to use.
## Returned are the final state and the descent log structs (see below).
##
## @var{data} contains all the information about the actual problem.
## This includes various parameters that influence the behaviour of this
## and other routines.  Furthermore, also some callback functions must
## be defined to evaluate the cost functional and to determine the
## next descent direction.  This struct is also updated with the current
## state at each optimisation step and passed to the callback functions.
##
## Most content of @var{data} is read-only, with @code{@var{data}.s} and
## @code{@var{data}.log} being the only exceptions.  They are updated
## accordingly from the results of the callback functions and the
## optimisation descent.  These two structs will be created and filled in
## by @code{so_run_descent}.  @code{@var{data}.log} is preserved if it is
## already present.  @code{@var{data}.s} will always be overwritten
## with the initial result of the @code{update_state} callback.
##
## The struct @code{@var{data}.p} contains parameters.  They can contain
## problem-dependent parameters that are handled by the callback routines
## as well as parameters for the line search with @code{so_step_armijo}.
## This routine is influenced by the following parameters directly:
##
## @table @code
## @item verbose
## Whether or not log chatter should be printed.
##
## @item descent.initialStep
## Initial trial step length for the line search in the first iteration.
##
## @item descent.projectSpeed
## If set, project the speed field (in L2, discontinuously) to enforce
## geometric constraints (see below).  If not set or missing, the
## geometry is projected instead.
## @end table
##
## The routine @code{so_init_params} can be used to initialise the parameter
## structure with some defaults.
##
## Information about the used grid and grid-dependent data should be in
## @code{@var{data}.g}.  Generally defined are the following fields:
##
## @table @code
## @item h
## The grid spacing.
##
## @item constraints
## Optional, can be used to specify the hold-all domain and contained
## regions as shape constraints.  If specified, the shape or the speed
## field will be ``projected'' (depending on
## @code{@var{data}.p.descent.projectSpeed}) to satisfy the constraints.
## The callback computing the descent direction is also a good place
## to take the constraints into account in a problem-specific way.
##
## The following fields (as logical arrays with the grid size) can
## be used to specify the constraints:
##
## @table @code
## @item holdall
## The hold-all domain.  The current shape must always be part of it.
## By excluding interior parts of the grid, forbidden regions can be defined.
## 
## @item contained
## A part of the hold-all domain that should always be part of the shape.
## @end table
## @end table
##
## The specific problem is handled by @emph{callback routines} to compute
## the cost and find descent directions.  They must be set in
## @code{@var{data}.cb}, with these interfaces:
##
## @table @code
## @item @var{s} = update_state (@var{phi}, @var{data})
## Compute the state (see below) for the geometry described by @var{phi}.
## The main work done by this callback is evaluating the cost functional.
## It does not need to update @code{@var{s}.phi} to @var{phi}, this will
## be done automatically after the call.
##
## @item [@var{f}, @var{dJ}] = get_direction (@var{data})
## Compute the speed field @var{f} and its associated directional
## shape derivative @var{dJ} to be used as the next descent direction.
## The current cost and geometry can be extracted from @code{@var{data}.s}.
##
## @item @var{exit} = check_stop (@var{data})
## Determine whether or not to stop the iteration.  Should return
## @code{true} if the state in @var{data} satisfies the problem-specific
## stopping criterion.  If not present, no stopping criterion
## is used and the iteration done for precisely @var{nSteps} iterations.
##
## @item @var{d} = solve_stationary (@var{f}, @var{data})
## Solve the stationary Eikonal equation for the speed field.
## If not present, @code{ls_solve_stationary} is used.  It can be
## set explicitly to give a more appropriate routine.  For instance,
## @code{ls_nb_from_geom} could be used if geometry information is computed.
## @end table
##
## Additional callback functions can be defined optionally
## in @code{@var{data}.handler}.  They can keep problem-specific
## logs of the descent, plot figures, print output or something else.
## They should always return a (possibly) updated log struct, which
## will be updated in @code{@var{data}.log} after the call.
## This struct can also be used to persist information between
## calls to the handlers.  These handlers can be defined:
##
## @table @code
## @item @var{log} = initialised (@var{data})
## Called before the iteration is started.  @code{@var{data}.s}
## contains the initial geometry and computed state data for it.
##
## @item @var{log} = before_step (@var{k}, @var{data})
## Called when starting a new step, before any action for the
## step is taken.  @var{k} is the number of the iteration and
## @code{@var{data}.s} contains the current state.
##
## @item @var{log} = direction (@var{k}, @var{f}, @var{dJ}, @var{data})
## Called when the descent direction has been determined.
## @var{f} and @var{dJ} are the results of @code{@var{data}.cb.get_direction}.
##
## @item @var{log} = after_step (@var{k}, @var{t}, @var{s}, @var{data})
## Called after a step has been found.  @var{t} is the step length
## found by the line search.  @var{s} is the state struct after the step
## has been taken.  The previous state can still be found in
## @code{@var{data}.s}.
##
## @item @var{log} = finished (@var{data})
## Called after the descent run.  @code{@var{data}.log} is already filled
## in with the standard data (number of steps and costs at each step).
## @end table
##
## Note that information about the speed field @var{f} is only
## passed to the @code{direction} handler and @emph{not again}
## to @code{after_step}.  If it is required by the latter,
## it should be saved somehow in the log struct.  This way, the handlers
## really only receive new information and no redundant arguments.
##
## Finally, we come to the state struct @code{@var{data}.s}.  It must
## contain at least the fields below, but can also contain problem-specific
## information that is, for instance, computed in the @code{update_state}
## handler and reused for computing the descent direction.
##
## @table @code
## @item phi
## The level-set function of the current geometry.  This is automatically
## updated after the @code{update_state} handler has been called.
##
## @item cost
## Value of the cost functional for the geometry in @code{@var{data}.s.phi}.
## @end table
##
## The log struct can contain mostly user-defined data, but
## @code{so_run_descent} will also fill in some data by itself:
##
## @table @code
## @item s0
## The state struct corresponding to the initial geometry defined
## by @var{phi0}.
##
## @item steps
## Total number of steps done.  This is useful if the stopping criterion
## was satisfied before @var{nSteps} iterations.
##
## @item costs
## Cost values for all steps.  This will be a vector of length
## @code{@var{data}.log.steps + 1}, since it contains both the initial
## and the final cost value in addition to all intermediate ones.
## @end table
##
## @seealso{so_init_params, so_save_descent, so_replay_descent,
## so_explore_descent, so_step_armijo}
## @end deftypefn

function [s, descentLog] = so_run_descent (nSteps, phi0, data)
  if (nargin () != 3)
    print_usage ();
  endif

  if (!isscalar (nSteps) || nSteps != round (nSteps) || nSteps < 1)
    print_usage ();
  endif

  if (data.p.descent.initialStep <= 0)
    error ("invalid initial step length");
  endif

  % Construct empty structs for optional fields.  This simplifies the
  % process later on.
  if (!isfield (data.g, "constraints"))
    data.g.constraints = struct ();
  endif
  if (!isfield (data, "handler"))
    data.handler = struct ();
  endif

  % Augment the update_state and get_direction callbacks.
  data._descent = struct ();
  data._descent.userUpdateState = data.cb.update_state;
  data._descent.userGetDirection = data.cb.get_direction;
  data.cb.update_state = @augmentedUpdateState;
  data.cb.get_direction = @augmentedGetDirection;

  % Create 'null' check_stop callback if none is present.
  if (!isfield (data.cb, "check_stop"))
    data.cb.check_stop = @() false;
  endif

  % Create default 'solve_stationary' callback.
  if (!isfield (data.cb, "solve_stationary"))
    data.cb.solve_stationary ...
      = @(f, data) ls_solve_stationary (data.s.phi, f, data.g.h);
  endif

  % Do some initialisations, but only if we're not in "continuation" mode.
  % In continuation mode, the "initialised" handler is already called
  % and the data to use is explicitly given by the caller.
  if (!isfield (data, "_descentContinuation"))
    data.s = data.cb.update_state (phi0, data);
    data.s.phi = phi0;
    if (!isfield (data, "log"))
      data.log = struct ();
    endif

    lastStep = data.p.descent.initialStep;
    costs = [data.s.cost];
    k = 0;

    data.log.s0 = data.s;
    if (isfield (data.handler, "initialised"))
      data.log = data.handler.initialised (data);
    endif
  else
    lastStep = data._descentContinuation.step;
    costs = data._descentContinuation.costs;
    k = data._descentContinuation.k;
  endif

  % Do the descent iteration.
  while (k < nSteps && !data.cb.check_stop (data.s))
    ++k;
    if (data.p.verbose)
      printf ("\nDescent iteration %d...\n", k);
      printf ("Starting cost: %.6f\n", data.s.cost);
      ticId = tic ();
    endif

    % Call "before_step" handler if present.
    if (isfield (data.handler, "before_step"))
      data.log = data.handler.before_step (k, data);
    endif

    % Check the constraints.  Since we enforce them, they
    % should always be satisfied.
    if (isfield (data.g.constraints, "holdall"))
      assert (ls_check (data.s.phi, "inside", data.g.constraints.holdall));
    endif
    if (isfield (data.g.constraints, "contained"))
      assert (ls_check (data.s.phi, "contain", data.g.constraints.contained));
    endif

    % Get the descent direction.
    [f, dJ] = data.cb.get_direction (data);
    if (dJ >= 0)
      error ("no descent direction found");
    endif

    % Print log and call "direction" handler if present.
    if (data.p.verbose)
      printf ("Directional derivative: %.6f\n", dJ);
    endif
    if (isfield (data.handler, "direction"))
      data.log = data.handler.direction (k, f, dJ, data);
    endif

    % Calculate level-set stationary distances and perform Armijo step.
    dists = data.cb.solve_stationary (f, data);
    [newS, lastStep] = so_step_armijo (lastStep, dists, f, dJ, data);

    % Call "after_step" handler.
    if (isfield (data.handler, "after_step"))
      data.log = data.handler.after_step (k, lastStep, newS, data);
    endif

    % Update cost tracking and switch over to new state.
    costs(k + 1) = newS.cost;
    data.s = newS;

    % Finish timing for verbose output.
    if (data.p.verbose)
      toc (ticId);
    endif
  endwhile

  % Fill in the final results.
  data.log.steps = k;
  assert (length (costs) == k + 1);
  data.log.costs = costs;

  % Call "finished" handler as the very last action.
  if (isfield (data.handler, "finished"))
    data.log = data.handler.finished (data);
  endif

  % Fill return values.
  s = data.s;
  descentLog = data.log;
endfunction

% Function that augments the 'update_state' handler with general stuff
% that it should do.  In particular, we enforce the constraints
% before doing anything (so that they are definitely not violated
% before calling the user code).  Also, update data.s.phi after the call.
function s = augmentedUpdateState (phi, data)
  phi = ls_normalise (phi, data.g.h);
  if (isfield (data.g.constraints, "holdall"))
    phi = ls_enforce (phi, "inside", data.g.constraints.holdall);
  endif
  if (isfield (data.g.constraints, "contained"))
    phi = ls_enforce (phi, "contain", data.g.constraints.contained);
  endif

  s = data._descent.userUpdateState (phi, data);
  s.phi = phi;
endfunction

% Function that augments the 'get_direction' handler.  It hooks the
% speed field projection onto it when requested by the parameters.
function [F, DJ] = augmentedGetDirection (data)
  [F, DJ] = data._descent.userGetDirection (data);

  % Note that the projection below does not take care of adapting DJ
  % accordingly!  Using this feature makes the descent method potentially
  % "inconsistent".  The same is true if the geometry is projected
  % after a step, though.

  if (isfield (data.p.descent, "projectSpeed") && data.p.descent.projectSpeed)
    if (isfield (data.g.constraints, "holdall"))
      F = ls_enforce_speed (F, "inside", data.g.constraints.holdall);
    endif
    if (isfield (data.g.constraints, "contained"))
      F = ls_enforce_speed (F, "contain", data.g.constraints.contained);
    endif
  endif
endfunction

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

% Construct basic example problem for testing.
%!shared data, phi0
%!  data = struct ();
%!  data.p = so_init_params (false);
%!  data.p.vol = 10;
%!  data.p.weight = 50;
%!
%!  n = 100;
%!  x = linspace (-10, 10, n);
%!  h = x(2) - x(1);
%!  data.g = struct ("x", x, "h", h);
%!
%!  data = so_example_problem (data);
%!  data.handler = struct ();
%!
%!  phi0 = ls_genbasic (data.g.x, "box", -3, 5);

% Test for errors.
%!error <Invalid call to>
%!  so_run_descent (1, 2);
%!error <Invalid call to>
%!  so_run_descent (1, 2, 3, 4);
%!error <Invalid call to>
%!  so_run_descent (0, phi0, data);
%!error <Invalid call to>
%!  so_run_descent (1.5, phi0, data);
%!error <Invalid call to>
%!  so_run_descent ([2, 3], phi0, data);

% Validate parameters.
%!error <invalid initial step length>
%!  d = data;
%!  d.p.descent.initialStep = 0;
%!  so_run_descent (1, phi0, d);

% Test optimisation of the example problem.
%!test
%!  nSteps = 12;
%!
%!  [s, descentLog] = so_run_descent (nSteps, phi0, data);
%!  assert (s.cost, 0, 1e-3);
%!  assert (length (descentLog.costs), nSteps + 1);
%!  assert (descentLog.costs(end), s.cost);
%!  %semilogy (descentLog.costs, "o");

% Test stopping criterion.
%!test
%!  tol = 5e-2;
%!  nSteps = 100;
%!  d = data;
%!  d.cb.check_stop = @(data) (data.cost < tol);
%!
%!  [s, descentLog] = so_run_descent (nSteps, phi0, d);
%!  assert (s.cost < tol);
%!  printf ("Steps needed: %d\n", descentLog.steps);
%!  assert (descentLog.steps < nSteps);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Demo.

% Use so_example_problem to define a demo.
%
%!demo
%!  data = struct ();
%!  data.p = so_init_params (true);
%!  data.p.vol = 10;
%!  data.p.weight = 50;
%!
%!  n = 100;
%!  x = linspace (-10, 10, n);
%!  h = x(2) - x(1);
%!  data.g = struct ("x", x, "h", h);
%!
%!  % Look at so_example_problem to see how the callbacks
%!  % and the plotting handler is defined.
%!  data = so_example_problem (data);
%!
%!  phi0 = ls_genbasic (data.g.x, "box", -3, 7);
%!  [s, descentLog] = so_run_descent (5, phi0, data);
%!
%!  figure ();
%!  semilogy (0 : descentLog.steps, descentLog.costs, "o");
%!  title ("Cost Descrease");
%!
%!  printf ("\nFinal interval: [%.6d, %.6d]\n", s.a, s.b);
%!  printf ("Final cost: %.6d\n", s.cost);