##  Copyright (C) 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{data} =} so_example_problem (@var{data})
##
## Construct an example problem for shape optimisation.  The example problem
## is one-dimensional.  Roughly speaking, the cost is chosen such that
## the optimal solution should be symmetric around the origin and
## have a determined volume.
##
## The parameters in @code{@var{data}.p} as well as the basic grid in
## @code{@var{data}.g} should already be set.  This routine fills in
## the necessary callbacks to compute the cost and to find a descent direction.
##
## These parameters in @code{@var{data}.p} are used to refine the
## example problem:
##
## @table @code
## @code vol
## The desired volume.  Deviations from this volume will be penalised in the
## cost functional.
##
## @code weight
## Weight of the volume penalisation term.  This is relative to the term
## that penalises non-symmetry and should be relatively large.
## @end table
##
## Note that these routines assume (for simplicity) that the level-set domain
## will always be a single interval.  Thus, the initial domain should
## fulfil this assumption.  (It will then automatically stay that way
## even when the shape is evolved.)
##
## In addition to the callbacks, also utility routines will be returned
## that can be used in handlers for plotting.  They will be set
## in @code{@var{data}.util}:
##
## @table @code
## @item [@var{a}, @var{b}] = bounds (@var{phi}, @var{data})
## Find the left and right bounds of the (assumed) interval that the
## level-set function @var{phi} describes.  This makes use of the
## grid in @code{@var{data}.g}.  The bounds will also be filled into
## the state struct as @code{@var{s}.a} and @code{@var{s}.b}.
## @end table
## 
## @seealso{so_run_descent}
## @end deftypefn

function data = so_example_problem (data)
  if (nargin () != 1)
    print_usage ();
  endif
  if (!isstruct (data))
    error ("DATA should be a struct");
  endif

  if (!isfield (data, "figs"))
    data.figs = struct ();
  endif

  data.util = struct ("bounds", @bounds);
  data.cb = struct ("update_state", @updateState, ...
                    "get_direction", @getDirection);
  data.handler = struct ("direction", @plotSpeed);

  % Set also an "onFirst" handler for so_explore_descent.
  % This handler plots all (already known) costs as we go.
  if (isfield (data.figs, "exploreCosts"))
    data.onFirst = struct ("before_step", @plotCosts);
  endif

  holdall = true (size (data.g.x));
  holdall([1, end]) = false;
  data.g.constraints = struct ("holdall", holdall);
endfunction

% Cost function of the example problem.
function s = updateState (phi, data)
  s = struct ("cost", 0);

  % Disallow empty domain.
  if (ls_isempty (phi))
    s.cost = Inf;
    return;
  endif

  % Penalise deviations from the desired volume.
  [s.a, s.b] = bounds (phi, data);
  s.vol = s.b - s.a;
  assert (s.vol > 0);
  s.cost += data.p.weight * (s.vol - data.p.vol)^2;

  % Penalise non-symmetry.  This is done by integrating over x.
  where = ls_inside (phi);
  s.integ = data.g.h * sum (data.g.x(where));
  s.cost += s.integ^2;
endfunction

% Compute shape derivative and gradient.
function [f, dJ] = getDirection (data)

  % Find derivatives at the end points (assuming F = 1).
  bdry = [data.s.a, data.s.b];
  deriv = zeros (size (bdry));
  for i = 1 : length (bdry)
    deriv(i) += 2 * data.p.weight * (data.s.vol - data.p.vol);
    deriv(i) += 2 * data.s.integ * bdry(i);
  endfor

  % Use derivatives at the end points to compute dJ.  Furthermore,
  % build the descent direction as a linear interpolation of the
  % end-point derivatives.

  f = NA (size (data.g.x));

  left = (data.g.x < data.s.a);
  fact = (data.g.x(left) - data.g.x(1)) ./ (data.s.a - data.g.x(1));
  f(left) = -deriv(1) * fact;

  mid = (data.g.x >= data.s.a & data.g.x <= data.s.b);
  fact = (data.g.x(mid) - data.s.a) ./ (data.s.b - data.s.a);
  f(mid) = -deriv(2) * fact - (1 - fact) * deriv(1);

  right = (data.g.x > data.s.b);
  fact = (data.g.x(end) - data.g.x(right)) ./ (data.g.x(end) - data.s.b);
  f(right) = -deriv(2) * fact;

  assert (all (isfinite (f)));
  dJ = -dot (deriv, deriv);
endfunction

% Handler to plot speed and the current situation.
function log = plotSpeed (k, f, dJ, data)
  fmin = min (f);
  fmax = max (f);

  if (!isfield (data.figs, "speed"))
    figure ();
  else
    figure (data.figs.speed);
    clf ();
  endif
  hold ("on");
  plot (data.g.x, f, "b");
  plot ([-5, -5], [fmin, fmax], "k");
  plot ([data.s.a, data.s.a], [fmin, fmax], "r");
  plot ([5, 5], [fmin, fmax], "k");
  plot ([data.s.b, data.s.b], [fmin, fmax], "r");
  hold ("off");
  axis ([data.g.x(1), data.g.x(end), fmin, fmax]);
  legend ("Speed", "Target", "Current");
  title (sprintf ("Step %d", k));

  log = data.log;
endfunction

% Plot the currently known costs.  This is used for so_explore_descent.
function log = plotCosts (k, data)
  assert (k <= data.log.steps);

  figure (data.figs.exploreCosts);
  clf ();
  costs = postpad (data.log.costs, data.p.nSteps + 1, NA);
  semilogy (0 : data.p.nSteps, costs, "o");
  axis ([0, data.p.nSteps]);
  title (sprintf ("Costs up to step %d", data.log.steps));

  log = data.log;
endfunction

% Bounds utility function.
function [a, b] = bounds (phi, data)
  if (ls_isempty (phi))
    error ("PHI describes empty domain");
  endif
  if (!all (size (phi) == size (data.g.x)))
    error ("PHI and DATA.g.x size mismatch");
  endif

  % Hack around infinite phi values.  This could be done
  % nicer by checking for infinite values directly (as it
  % is done in ls_find_geometry, for instance), but just
  % do it like this for simplicity.
  whereInf = !isfinite (phi);
  phi(whereInf) = sign (phi(whereInf)) * max (abs (phi(!whereInf)));

  insideInd = find (ls_inside (phi));
  minInd = min (insideInd);
  maxInd = max (insideInd);

  convexZero = @(phiA, phiB, xA, xB) ...
                  xA + (xB - xA) * abs (phiA) / (abs (phiA) + abs (phiB));

  if (minInd == 1)
    a = data.g.x(1);
  else
    a = convexZero (phi(minInd - 1), phi(minInd), ...
                    data.g.x(minInd - 1), data.g.x(minInd));
  endif

  if (maxInd == length (phi))
    b = data.g.x(end);
  else
    b = convexZero (phi(maxInd), phi(maxInd + 1), ...
                    data.g.x(maxInd), data.g.x(maxInd + 1));
  endif
endfunction

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

% Build basic data with example grid.
%!shared basicData
%!  basicData = struct ();
%!  basicData.p = so_init_params (true);
%!  basicData.p.vol = 4;
%!  basicData.p.weight = 10;
%!
%!  n = 50;
%!  x = linspace (-4, 5, n);
%!  h = x(2) - x(1);
%!  basicData.g = struct ("x", x, "h", h);

% Test for errors.
%!error <Invalid call to>
%!  so_example_problem ();
%!error <Invalid call to>
%!  so_example_problem (1, 2);
%!error <DATA should be a struct>
%!  so_example_problem (1);

% Test the bounds routine created.
%!test
%!  data = so_example_problem (basicData);
%!  phi = ls_genbasic (data.g.x, "box", -2.5, 1.3);
%!
%!  [a, b] = data.util.bounds (phi, data);
%!  assert (a, -2.5, sqrt (eps));
%!  assert (b, 1.3, sqrt (eps));

% Test that the cost works qualitatively as it should.
%!test
%!  data = so_example_problem (basicData);
%!
%!  phi = ls_genbasic (data.g.x, "box", -2, 2); 
%!  s0 = data.cb.update_state (phi, data);
%!
%!  phi = ls_genbasic (data.g.x, "box", -3, 3); 
%!  s1 = data.cb.update_state (phi, data);
%!
%!  phi = ls_genbasic (data.g.x, "box", -1, 1); 
%!  s2 = data.cb.update_state (phi, data);
%!
%!  phi = ls_genbasic (data.g.x, "box", -1, 3); 
%!  s3 = data.cb.update_state (phi, data);
%!
%!  phi = ls_genbasic (data.g.x, "box", -3, 1); 
%!  s4 = data.cb.update_state (phi, data);
%!
%!  assert (all ([s0.cost, s1.cost, s2.cost, s3.cost, s4.cost] >= 0));
%!  assert (s0.cost, 0, 1e-1);
%!  assert (s1.cost > s3.cost);
%!  assert (s1.cost > s4.cost);
%!  assert (s2.cost > s3.cost);
%!  assert (s2.cost > s4.cost);
%!  assert (s3.cost > s0.cost);
%!  assert (s4.cost > s0.cost);

% Do a rough finite-difference test for the derivative and ensure
% that we have a descent direction.
%
%!function testDeriv (a, b, data)
%!  phi = ls_genbasic (data.g.x, "box", a, b);
%!  data.s = data.cb.update_state (phi, data);
%!  [f, dJ] = data.cb.get_direction (data);
%!  assert (dJ < 0);
%!
%!  dt = 5e-3;
%!  dists = ls_solve_stationary (phi, f, data.g.h);
%!  phip = ls_extract_solution (dt, dists, phi, f);
%!  sp = data.cb.update_state (phip, data);
%!
%!  assert (sp.cost < data.s.cost);
%!  relErr = abs (sp.cost - (data.s.cost + dt * dJ)) / data.s.cost;
%!  assert (relErr < 0.1);
%!endfunction
%
%!test
%!  data = so_example_problem (basicData);
%!
%!  testDeriv (-1, 1, data);
%!  testDeriv (-3, 3, data);
%!  testDeriv (-1, 3, data);
%!  testDeriv (-3, 1, data);

% This routine is also "tested" by the tests and demos in so_run_descent.
