##  Copyright (C) 2013-2019  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{u} =} fastmarching (@var{u0}, @var{f})
## @deftypefnx {Function File} {[@var{u}, @var{g}] =} fastmarching (@var{u0}, @var{g0}, @var{f})
## 
## Solve the Eikonal equation using the Fast-Marching Method.  In particular,
## the equation
## @tex
## \begin{equation*}
##   \left| \nabla u \right| = f
## \end{equation*}
## @end tex
## @ifnottex
##
##
## @example
## | grad u | = f
## @end example
##
## @end ifnottex
## is solved on a multi-dimensional domain.  @var{u0} and @var{f} must be
## multi-dimensional arrays of the same dimension and size, and the returned
## solution @var{u} will be also of this size.  @var{f} should be positive,
## and we assume the grid spacing is unity (or, equivalently, has been
## absorbed into f already).
## 
## @var{u0} defines the initial setting and the domain.  It should contain the
## initial distances for points that are used as initially alive points,
## be @code{Inf} for points not part of the domain, and @code{NA} at all
## other points.  The latter are considered as far-away at the beginning
## and their distances will be calculated.  @var{f} is ignored at all but
## the far-away points.
## 
## If @var{g0} is also given, it should contain function values on initially
## alive points that will be extended onto the domain along the
## ``normal directions''.  This means that
## @tex
## \begin{equation*}
##  \nabla g \cdot \nabla u = 0
## \end{equation*}
## @end tex
## @ifnottex
##
##
## @example
## grad g * grad u = 0
## @end example
##
## @end ifnottex
## will be fulfilled for the solution @var{u} of the Eikonal equation.
## @var{g} will be @code{NA} at points not in the domain or not reached by the
## fast marching.
## 
## At return, @var{u} will contain the calculated distances.
## @code{Inf} values will be retained, and it may still contain @code{NA}
## at points to which no suitable path can be found.
##
## @seealso{ls_signed_distance, ls_solve_stationary}
## @end deftypefn

function [u, g] = fastmarching (u0, arg2, arg3)
  if (nargin () == 2)
    f = arg2;
    g0 = NA (size (u0));
    if (nargout () > 1)
      print_usage ();
    endif
  elseif (nargin () ~= 3)
    print_usage ();
  else
    assert (nargin () == 3)
    f = arg3;
    g0 = arg2;
  endif

  where = isna (u0);
  if (~all (f(where) > 0))
    error ("F must be positive");
  endif

  sz = size (u0);
  if (~all (sz == size (f)))
    error ("U0 and F must be of the same size");
  endif
  if (~all (sz == size (g0)))
    error ("U0 and G0 must be of the same size");
  endif

  domain = ~isinf (u0);
  alive = ~isna (u0);
  if (~all (alive(domain) == (isfinite (u0(domain)))))
    error ("U0 shold only contain Inf, NA and finite values");
  endif
  
  [u, g, reached] = __levelset_internal_fastmarching (domain, u0, g0, alive, f);
  u(~reached) = NA;
  u(~domain) = Inf;
  g(~reached | ~domain) = NA;
endfunction

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

% Test for error conditions.
%!error <Invalid call to>
%!  fastmarching ()
%!error <Invalid call to>
%!  fastmarching (1, 2, 3, 4)
%!error <Invalid call to>
%!  [a, b] = fastmarching (1, 2)
%!error <U0 and F must be of the same size>
%!  fastmarching ([1, 2], [1; 2])
%!error <U0 and G0 must be of the same size>
%!  fastmarching ([1, 2], [1; 2], [1, 2])
%!error <F must be positive>
%!  fastmarching ([NA, 1], [0, 5]);

% Basic test for obstacle with a square for which propagation is only
% allowed along the outer edges.
%!test
%!  n = 600;
%!  x = linspace (0, 1, n);
%!  h = x(2) - x(1);
%!  [XX, YY] = meshgrid (x, x);
%!  u0 = NA (n, n);
%!  u0(2 : end - 1, 2 : end - 1) = Inf;
%!  u0(1, 1) = 0;
%!  f = 1 + YY;
%!  u = fastmarching (u0, h * f);
%!
%!  % Length of vertical side: integrate (1 + y, y, 0, x)
%!  vertLength = @(x) x/2 .* (x + 2);
%!
%!  % Speed difference at top and bottom, gives where the
%!  % paths will meet on the bottom side.
%!  %   solve (x * fBot = fTop + (1 - x) * fBot, x)
%!  fTop = mean (f(1, :));
%!  fBot = mean (f(end, :));
%!  xMeet = (fTop + fBot) / 2 / fBot;
%!  botRow = vertLength (1) + fBot * x;
%!  where = (x > xMeet);
%!  botRow(where) = vertLength (1) + 1 + fBot * (1 - x(where));
%!
%!  assert (u(1, :), x, sqrt (eps));
%!  assert (u(:, 1), vertLength (x'), 1e-3);
%!  assert (u(:, end), 1 + vertLength (x'), 1e-3);
%!  assert (u(end, :), botRow, 1e-3);

% Check 1D situation.  Also verifies that it works with non-zero initial values.
%!test
%!  n = 1100;
%!  x = linspace (0, 1, n);
%!  h = x(2) - x(1);
%!
%!  obst1 = round (0.5 / h);
%!  obst2 = round (0.75 / h);
%!  s1 = round (0.2 / h);
%!  s1val = 1;
%!  s2 = round (0.9 / h);
%!  s2val = -1;
%!
%!  u0 = NA (n, 1);
%!  u0([obst1, obst2]) = Inf;
%!  u0([s1, s2]) = [s1val, s2val];
%!  f = ones (size (u0));
%!  g0 = NA (size (u0));
%!  g0([s1, s2]) = [s1val, s2val];
%!
%!  [u, g] = fastmarching (u0, g0, f * h);
%!
%!  gDesired = NA (size (u0));
%!  gDesired(1 : obst1 - 1) = s1val;
%!  gDesired(obst2 + 1 : end) = s2val;
%!  assert (g, gDesired, sqrt (eps));
%!
%!  uDesired = NA (size (u0));
%!  uDesired([obst1, obst2]) = Inf;
%!  uDesired(1 : obst1 - 1) = s1val + abs (s1 * h - x(1 : obst1 - 1));
%!  uDesired(obst2 + 1 : end) = s2val + abs (s2 * h - x(obst2 + 1 : end));
%!  assert (u, uDesired, 1e-3);

% Check 0D situation (which is trivial but nevertheless).
%!test
%!  assert (fastmarching ([], []), []);

% Extend from a circular initial region outwards, also including a function
% value which is the angle.  For this configuration, the expected result
% is easy to calculate.  Test in 3D.
%!test
%!  n = 42;
%!  r = 1;
%!  R = 2;
%!  x = linspace (-R, R, n);
%!  h = x(2) - x(1);
%!
%!  [XX, YY, ZZ] = ndgrid (x);
%!  RR = sqrt (XX.^2 + YY.^2 + ZZ.^2);
%!  PPhi = NA (size (XX));
%!  for i = 1 : numel (PPhi)
%!    pt = [XX(i), YY(i), ZZ(i)];
%!    PPhi(i) = acos (dot (pt, [0, 0, 1]) / norm (pt));
%!  endfor
%!
%!  initial = (RR < r);
%!  u0 = NA (size (XX));
%!  u0(initial) = 0;
%!  g0 = NA (size (XX));
%!  g0(initial) = PPhi(initial);
%!  f = ones (size (u0));
%!  [u, g] = fastmarching (u0, g0, h * f);
%!
%!  assert (u(~initial), RR(~initial) - r, 1e-1);
%!  assert (g(~initial), PPhi(~initial), 2e-1);

% Function for calculating the approximate gradient of a given function
% using central differences.  This is used to verify grad g * grad u = 0.
%!function [dx, dy] = gradient (fcn, h)
%!  ext = NA (size (fcn) + 2);
%!  ext(2 : end - 1, 2 : end - 1) = fcn;
%!  ext(:, 1) = 2 * ext(:, 2) - ext(:, 3);
%!  ext(:, end) = 2 * ext(:, end - 1) - ext(:, end - 2);
%!  ext(1, :) = 2 * ext(2, :) - ext(3, :);
%!  ext(end, :) = 2 * ext(end - 1, :) - ext(end - 2, :);
%!  dx = (ext(2 : end - 1, 3 : end) - ext(2 : end - 1, 1 : end - 2)) / (2 * h);
%!  dy = (ext(3 : end, 2 : end - 1) - ext(1 : end - 2, 2 : end - 1)) / (2 * h);
%!endfunction

% Test extending of function values in a more general setting.  We verify
% that the "defining equation"
%
%   grad g * grad u = 0
%
% is indeed valid.
%!test
%!  n = 1000;
%!  x = linspace (-2, 2, n);
%!  h = x(2) - x(1);
%!  [XX, YY] = meshgrid (x, x);
%!
%!  phi0 = 2 * XX.^2 + 0.4 * YY.^2 - 1;
%!  initial = (phi0 < 0);
%!  %contour (XX, YY, phi0, [0, 0], "r");
%!
%!  f = ones (size (XX));
%!  u0 = NA (size (XX));
%!  u0(initial) = 0;
%!  g0 = NA (size (XX));
%!  g0(initial) = sin (atan2 (YY(initial), XX(initial)));
%!
%!  [u, g] = fastmarching (u0, g0, h * f);
%!
%!  [dxu, dyu] = gradient (u, h);
%!  [dxg, dyg] = gradient (g, h);
%!
%!  % We only verify the mean error here, since there tend
%!  % to be lower-dimensional lines on which the error is
%!  % large while it is very low otherwise.
%!  err = abs (dxu .* dxg + dyu .* dyg);
%!  assert (mean (err(initial)), 0, 1e-3);

% Particular test situation where during the fast marching
% procedure the distance of a point is later calculated to
% be further from the initial set.
%!test
%!  u0 = [1.257221, NA, NA; Inf, 0.275128, 0.433001];
%!  f = [NA, 0.85714, 0.66667; NA, NA, NA];
%!  u = fastmarching (u0, f);

% Test the case of far-away alive point.
%!test
%!  u0 = [0, NA; Inf, 2];
%!  f = ones (size (u0));
%!  u = fastmarching (u0, f);
%!  assert (u(1, 2), 1, sqrt (eps));

% Test for the case with multiple alive neighbours along a single dimension.
%!test
%!  u0 = [0, NA, NA, NA, 0.1; ...
%!        Inf, Inf, Inf, Inf, Inf; ...
%!        0.1, NA, NA, NA, 0];
%!  f = ones (size (u0));
%!  u = fastmarching (u0, f);
%!
%!  uExpected = [0, 1, 2, 1.1, 0.1; ...
%!               Inf, Inf, Inf, Inf, Inf; ...
%!               0.1, 1.1, 2, 1, 0];
%!  assert (u, uExpected, sqrt (eps));

% Estimate the runtime complexity.
%!test
%!  ns = [101, 201, 501, 1001, 2001];
%!  realNs = NA (size (ns));
%!  times = NA (size (ns));
%!  for i = 1 : length (ns)
%!    nMid = (ns(i) + 1) / 2;
%!    assert (nMid, round (nMid));
%!    u0 = NA (ns(i), ns(i));
%!    u0(nMid, nMid) = 0;
%!    f = ones (size (u0));
%!    realNs(i) = numel (u0);
%!
%!    old = cputime ();
%!    u = fastmarching (u0, f);
%!    new = cputime ();
%!    times(i) = new - old;
%!
%!    printf ("n = %4d: t = %6.3f s\n", ns(i), times(i));
%!  endfor
%!
%!  nLogN = realNs .* log (realNs);
%!  p = polyfit (log (realNs), log (times), 1);
%!  pLogN = polyfit (log(nLogN), log (times), 1);
%!  printf ("Exponents:\n        O(n): %4.2f\n  O(n log n): %4.2f\n", ...
%!          p(1), pLogN(1));
%!  assert (p(1), 1, 1e-1);
%!  assert (pLogN(1), 1, 1e-1);

% Compare the solution of the level-set equation using finite differences
% and a viscosity term to our formula.  If the initial profile is constructed
% by fastmarching to satisfy F | grad phi0 | = 1, then it should just
% be shifted down in time.
%!test
%!  T = 1.6;
%!  timeSteps = 10000;
%!  L = 2;
%!  gridSteps = 10000;
%!  eps = 1e-4;
%!  floorValue = -1/2;
%!  tol = 7e-2;
%!  
%!  t = linspace (0, T, timeSteps);
%!  dt = t(2) - t(1);
%!  x = linspace (-L/2, L/2, gridSteps);
%!  h = x(2) - x(1);
%!  
%!  F = sqrt (abs (x));
%!  phi0 = NA (size (x));
%!  phi0(1 : end/4) = floorValue;
%!  phi0(3*end/4 : end) = floorValue;
%!  phi0 = fastmarching (phi0, h ./ F);
%!  
%!  % Plot the speed field and initial data.
%!  %{
%!  figure ();
%!  plot (x, F, "r", x, phi0, "k");
%!  legend ("F", "\\phi_0");
%!  title ("Speed Field and Initial Data");
%!  ax = [-L/2, L/2, -1, 1];
%!  axis (ax);
%!  %}
%!  
%!  % Create the Laplacian as matrix.  We assume homogeneous Neumann boundary
%!  % conditions so that the initial data can be chosen as we like.
%!  Delta = spdiags (repmat ([1, -2, 1], gridSteps, 1), [-1, 0, 1], ...
%!                   gridSteps, gridSteps);
%!  Delta(1, 1) += Delta(1, 2);
%!  Delta(end, end) += Delta(end - 1, end);
%!  Delta /= h^2;
%!  
%!  % Perform the calculation.  The Laplacian is done implicitly, the
%!  % non-linear term explicitly:
%!  %
%!  %   phi+ - phi0 = dt (eps Delta phi+ - F |grad phi0|)
%!  %   => (I - dt eps Delta) phi+ = phi0 - F |grad phi0|
%!  
%!  %figure ();
%!  phi = phi0';
%!  M = eye (gridSteps) - dt * eps * Delta;
%!  for tInd = 2 : timeSteps
%!  
%!    % Calculate derivatives as symmetric differences, and zero on the
%!    % boundary.  Since the data is flat there, this is correct.
%!    d = NA (size (phi));
%!    low = 1 : gridSteps - 2;
%!    mid = 2 : gridSteps - 1;
%!    high = 3 : gridSteps;
%!    d(mid) = (phi(high) - phi(low)) / (2 * h);
%!    d([1, end]) = 0;
%!    gradNorm = abs (d);
%!  
%!    % Do the solve.
%!    phi = M \ (phi - dt * F' .* gradNorm);
%!    phiEx = max (phi0 - t(tInd), floorValue);
%!    assert (phi, phiEx', tol);
%!  
%!    % Show solution as animation in real time.
%!    %{
%!    if (mod (tInd, 100) == 0)
%!      plot (x, phi, "r", x, phiEx, "k");
%!      legend ("Approximate", "Exact");
%!      axis (ax);
%!      title (sprintf ("t = %.2f", t(tInd)));
%!      drawnow ();
%!      %pause (0.1);
%!    endif
%!    %}
%!  endfor

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Demos.

% Plot how the distance and an extended function "flow" around an obstacle.
%!demo
%!  n = 100;
%!  x = linspace (-1, 1, n);
%!  h = x(2) - x(1);
%!  [XX, YY] = meshgrid (x, x);
%!
%!  u0 = NA (size (XX));
%!  u0(sqrt (XX.^2 + YY.^2) < 0.5) = Inf;
%!  u0(end, :) = 0;
%!  f = 3 + XX;
%!  g0 = NA (size (XX));
%!  g0(end, :) = x;
%!
%!  [u, g] = fastmarching (u0, g0, h * f);
%!
%!  figure ();
%!  subplot (2, 2, 1);
%!  imagesc (f);
%!  colorbar ();
%!  title ("Speed Field");
%!  subplot (2, 2, 2);
%!  imagesc (u0);
%!  colorbar ();
%!  title ("Initial Distances");
%!  subplot (2, 2, 3);
%!  imagesc (u);
%!  colorbar ();
%!  title ("Distance u");
%!  subplot (2, 2, 4);
%!  imagesc (g);
%!  colorbar ();
%!  title ("Extended Function g");

% Demo solving a "maze".
%!demo
%!  pkgs = pkg ("list");
%!  file = false;
%!  for i = 1 : length (pkgs)
%!    if (strcmp (pkgs{i}.name, "level-set"))
%!      file = fullfile (pkgs{i}.dir, "maze.png");
%!    endif
%!  endfor
%!  if (!file)
%!    error ("could not locate example image file");
%!  endif
%!  img = double (imread (file)) / 255;
%!
%!  % Red is the start area, white (not black) the accessible domain.
%!  start = (img(:, :, 1) > img(:, :, 2) + img(:, :, 3));
%!  domain = (img(:, :, 2) + img(:, :, 3) > 1.0) | start;
%!
%!  u0 = NA (size (domain));
%!  u0(start) = 0;
%!  u0(~domain) = Inf;
%!  f = ones (size (u0));
%!
%!  u = fastmarching (u0, f);
%!  u = u / max (u(isfinite (u)));
%!
%!  % Format result as image.
%!  walls = isinf (u);
%!  notReached = isna (u);
%!  ur = u;
%!  ug = u;
%!  ub = u;
%!  ur(notReached) = 1;
%!  ug(notReached) = 0;
%!  ub(notReached) = 0;
%!  ur(walls) = 0;
%!  ug(walls) = 0;
%!  ub(walls) = 1;
%!  ur(start) = 1;
%!  ug(start) = 1;
%!  ub(start) = 0;
%!  u = NA (size (u, 1), size (u, 2), 3);
%!  u(:, :, 1) = ur;
%!  u(:, :, 2) = ug;
%!  u(:, :, 3) = ub;
%!  figure ();
%!  imshow (u);
%!  title ("Solved Maze.  Blue: Walls, Red: Unreachable, Yellow: Start.");
