##  Copyright (C) 2015-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{s}, @var{log}] =} so_replay_descent (@var{file}, @var{init})
## @deftypefnx {Function File} {[@var{s}, @var{log}] =} so_replay_descent (@var{file}, @var{init}, @var{nSteps})
##
## Replay a descent file saved by @code{so_save_descent}.  @var{file} should
## be either an already open file descriptor or a file name to be opened
## in @qcode{"rb"} mode.
##
## @var{init} must be a function that can be called with the ``header''
## saved in the descent log.  It should return the @var{data} structure
## to use for the descent run.  The handlers defined in
## @code{@var{data}.handler} will be called in the same way as they
## would by @code{so_run_descent}.  Also similar log chatter will be printed
## if @code{@var{data}.p.verbose} is @code{true}.
##
## If the state structs were compressed when saving with @code{so_save_descent},
## a routine should be provided to ``uncompress'' them as necessary.  This
## function should be defined in @code{@var{data}.compress.load.state}
## and should return the uncompressed state struct when called
## with the arguments @code{(@var{s}, @var{data})} where @var{s}
## is the compressed struct.
##
## If @var{nSteps} is not given, the file is read until EOF is encountered.
## If @var{nSteps} is present, the descent will be run until either
## this number is reached or @code{@var{data}.cb.check_stop} returns true.
## If the descent log does not contain enough data for this condition to be
## fulfilled, more steps will be performed according to @code{so_run_descent}.
## In this case, the callbacks in @code{@var{data}.cb} must be set in the
## same way as for @code{so_run_descent}.
##
## Returned are the final state and the produced log structure,
## in the same way as by @code{so_run_descent}.
##
## This routine requires @code{fload} from the @code{parallel} package
## to be available.
##
## @seealso{so_explore_descent, so_run_descent, so_save_descent, fload}
## @end deftypefn

function [s, descentLog] = so_replay_descent (file, init, nSteps)
  if (nargin () < 2 || nargin () > 3)
    print_usage ();
  endif

  if (nargin () == 3)
    forceSteps = true;
    if (!isscalar (nSteps) || nSteps != round (nSteps) || nSteps < 1)
      print_usage ();
    endif
  else
    forceSteps = false;
  endif

  openFile = ischar (file);
  if (!openFile && (!isnumeric (file) || !isscalar (file)))
    error ("FILE must be a string or file ID");
  endif

  hasFload = exist ("fload");
  if (hasFload != 2 && hasFload != 3)
    error ("'fload' is not available");
  endif

  % Open file and read header.
  if (openFile)
    fd = fopen (file, "rb");
    if (fd == -1)
      error ("failed to open file '%s' for reading the descent log", file);
    endif
  else
    fd = file;
  endif
  header = fload (fd);

  % Call init method and load initial state.
  data = init (header);
  data.s = so_load_compressed (fd, "state", data);

  % Add some (possibly) missing stuff with defaults.
  if (!isfield (data, "log"))
    data.log = struct ();
  endif
  if (!isfield (data, "handler"))
    data.handler = struct ();
  endif
  if (!isfield (data, "cb"))
    data.cb = struct ();
  endif
  if (!isfield (data.cb, "check_stop"))
    data.cb.check_stop = @() false;
  endif

  % Call "initialised" handler and save initial state struct.
  data.log.s0 = data.s;
  if (isfield (data.handler, "initialised"))
    data.log = data.handler.initialised (data);
  endif

  % Do the iteration.
  costs = [data.s.cost];
  k = 0;
  while (true)
    if (ferror (fd))
      error ("error reading descent log");
    endif

    % Honour stopping condition if nSteps is given.
    if (forceSteps && (k >= nSteps || data.cb.check_stop (data.s)))
      break;
    endif

    % Read next data and check for EOF.
    try
      f = fload (fd);
    catch
      if (feof (fd))
        break;
      endif
      error ("error reading descent log");
    end_try_catch
    dJ = fload (fd);
    step = fload (fd);
    newS = so_load_compressed (fd, "state", data);
    if (ferror (fd))
      error ("error reading descent log");
    endif

    % Call handlers.
    ++k;
    if (data.p.verbose)
      printf ("\nDescent iteration %d...\n", k);
      printf ("Starting cost: %.6f\n", data.s.cost);
      printf ("Directional derivative: %.6f\n", dJ);
      printf ("Armijo step %.6f: cost = %.6f\n", step, newS.cost);
    endif
    if (isfield (data.handler, "before_step"))
      data.log = data.handler.before_step (k, data);
    endif
    if (isfield (data.handler, "direction"))
      data.log = data.handler.direction (k, f, dJ, data);
    endif
    if (isfield (data.handler, "after_step"))
      data.log = data.handler.after_step (k, step, newS, data);
    endif

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

  % Make sure to close the file now.  Record EOF condition
  % before we do so.
  haveEOF = feof (fd);
  if (openFile)
    fclose (fd);
  endif

  % If we have an EOF condition, see if we need more steps.
  if (haveEOF && forceSteps)
    assert (k < nSteps && !data.cb.check_stop (data.s));

    if (k == 0)
      error ("no data in descent log");
    endif

    % Call through to so_run_descent, in a special "continuation"
    % mode (defined internally for this purpose).
    data._descentContinuation = struct ("k", k, "step", step, "costs", costs);
    [s, descentLog] = so_run_descent (nSteps, [], data);
    return;
  endif

  % Finish everything.
  data.log.steps = k;
  assert (length (costs) == k + 1);
  data.log.costs = costs;
  if (isfield (data.handler, "finished"))
    data.log = data.handler.finished (data);
  endif

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

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

% Test for errors.
%!error <Invalid call to>
%!  so_replay_descent (1);
%!error <Invalid call to>
%!  so_replay_descent (1, 2, 3, 4);
%!error <Invalid call to>
%!  so_replay_descent (1, 2, 0);
%!error <Invalid call to>
%!  so_replay_descent (1, 2, 1.5);
%!error <Invalid call to>
%!  so_replay_descent (1, 2, [2, 3]);
%!error <FILE must be a string or file ID>
%!  so_replay_descent (struct (), 2);
%!error <FILE must be a string or file ID>
%!  so_replay_descent ([2, 3], 2);
%!error <'fload' is not available>
%!  pkg unload parallel;
%!  so_replay_descent ("foo", 2);

% Define compression / uncompression functions.  They do not really
% "compress", but instead manipulate the data in ways such that it
% is later possible to check that they have been called correctly.
%
%!function x = compressionCheckStruct (s, data)
%!  x = struct ("weight", data.p.weight, "cost", s.cost);
%!endfunction
%
%!function checkCompression (s, data, field)
%!  assert (isfield (s, field));
%!  assert (getfield (s, field), compressionCheckStruct (s, data));
%!endfunction
%
%!function s = compressState (s, data)
%!  s.compressed = compressionCheckStruct (s, data);
%!endfunction
%
%!function s = uncompressState (s, data)
%!  checkCompression (s, data, "compressed");
%!  s.uncompressed = compressionCheckStruct (s, data);
%!endfunction

% Sample handlers that just build up "some" array from all the data
% they are passed.  This is used to compare replayed and computed runs.
%
%!function log = initialised (data)
%!  log = data.log;
%!  log.mashUp = [data.s.a, data.s.b];
%!endfunction
%
%!function log = beforeStep (k, data)
%!  log = data.log;
%!  log.mashUp = [log.mashUp, k, data.s.vol];
%!
%!  if (data.checkUncompression)
%!    checkCompression (data.s, data, "uncompressed");
%!  endif
%!endfunction
%
%!function log = direction (k, f, dJ, data)
%!  log = data.log;
%!  log.mashUp = [log.mashUp, k, mean(f), sum(f), dJ, data.s.integ];
%!endfunction
%
%!function log = afterStep (k, t, s, data)
%!  log = data.log;
%!  log.mashUp = [log.mashUp, k, t, s.a, s.b, s.cost, data.s.cost];
%!
%!  if (data.checkUncompression)
%!    checkCompression (s, data, "uncompressed");
%!  endif
%!endfunction
%
%!function log = finished (data)
%!  log = data.log;
%!  log.mashUp = [log.mashUp, data.s.cost, log.steps, log.costs];
%!endfunction

% Define a function that initialises the problem data.
%!function data = getData ()
%!  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 ("initialised", @initialised, ...
%!                         "before_step", @beforeStep, ...
%!                         "direction", @direction, ...
%!                         "after_step", @afterStep, ...
%!                         "finished", @finished);
%!
%!  % Unless explicitly set, do not expect compressed states.
%!  data.checkUncompression = false;
%!
%!  data.phi0 = ls_genbasic (data.g.x, "box", -3, 5);
%!endfunction

% Get data and enable compression.
%!function data = getDataForCompression (uncompress)
%!  data = getData ();
%!  data.compress = struct ();
%!  data.compress.save = struct ("state", @compressState);
%!  data.compress.load = struct ("state", @uncompressState);
%!  data.checkUncompression = uncompress;
%!endfunction

% Compare results (s, log) two runs and check if the mash-ups match.
% We cannot fully compare the log structs, though, since they contain
% private data used for saving.
%!function compareResults (s1, l1, s2, l2)
%!  assert (s2, s1);
%!  assert (l2.mashUp, l1.mashUp);
%!endfunction

% Create a basic log and replay it without forcing steps.  This tests
% both saving to a file name and file descriptor.
%!test
%!  pkg load parallel;
%!
%!  nSteps = 10;
%!  baseData = getData ();
%!
%!  f = tempname ();
%!  data = so_save_descent (f, struct (), baseData);
%!  [s1, l1] = so_run_descent (nSteps, data.phi0, data);
%!  [s2, l2] = so_replay_descent (f, @getData);
%!  unlink (f);
%!  compareResults (s1, l1, s2, l2);
%!
%!  f = tmpfile ();
%!  data = so_save_descent (f, struct (), baseData);
%!  [s1, l1] = so_run_descent (nSteps, data.phi0, data);
%!  frewind (f);
%!  [s2, l2] = so_replay_descent (f, @getData);
%!  fclose (f);
%!  compareResults (s1, l1, s2, l2);
%!
%!  [pr, pw] = pipe ();
%!  data = so_save_descent (pw, struct (), baseData);
%!  [s1, l1] = so_run_descent (nSteps, data.phi0, data);
%!  fclose (pw);
%!  [s2, l2] = so_replay_descent (pr, @getData);
%!  fclose (pr);
%!  compareResults (s1, l1, s2, l2);

% Check for "no data" error if we need to compute more steps
% but *none* are saved.  Also check that it *is* ok to have
% zero steps if the stopping criterion stops immediately.
%
%!test
%!  pkg load parallel;
%!
%!  baseData = getData ();
%!  [pr, pw] = pipe ();
%!  data = so_save_descent (pw, struct (), baseData);
%!  data.cb.check_stop = @() true;
%!  [s1, l1] = so_run_descent (1, data.phi0, data);
%!  fclose (pw);
%!  assert (l1.steps, 0);
%!  [s2, l2] = so_replay_descent (pr, @getData);
%!  fclose (pr);
%!  compareResults (s1, l1, s2, l2);
%!
%!  baseData = getData ();
%!  [pr, pw] = pipe ();
%!  data = so_save_descent (pw, struct (), baseData);
%!  [s1, l1] = so_run_descent (1, data.phi0, data);
%!  fclose (pw);
%!  [s2, l2] = so_replay_descent (pr, @getData, 1);
%!  fclose (pr);
%!  compareResults (s1, l1, s2, l2);
%
%!error <no data in descent log>
%!  pkg load parallel;
%!
%!  baseData = getData ();
%!  unwind_protect
%!    [pr, pw] = pipe ();
%!    data = so_save_descent (pw, struct (), baseData);
%!    data.cb.check_stop = @() true;
%!    [s1, l1] = so_run_descent (1, data.phi0, data);
%!    fclose (pw);
%!    [s2, l2] = so_replay_descent (pr, @getData, 1);
%!  unwind_protect_cleanup
%!    fclose (pr);
%!  end_unwind_protect
%!  compareResults (s1, l1, s2, l2);

% Check forcing the steps to be less or to require computation.
%!test
%!  pkg load parallel;
%!
%!  nStepsShort = 5;
%!  nStepsLong = 10;
%!
%!  baseData = getData ();
%!  [prl, pwl] = pipe ();
%!  data = so_save_descent (pwl, struct (), baseData);
%!  [sl, ll] = so_run_descent (nStepsLong, data.phi0, data);
%!  fclose (pwl);
%!
%!  baseData = getData ();
%!  [prs, pws] = pipe ();
%!  data = so_save_descent (pws, struct (), baseData);
%!  [ss, ls] = so_run_descent (nStepsShort, data.phi0, data);
%!  fclose (pws);
%!
%!  [s, l] = so_replay_descent (prl, @getData, nStepsShort);
%!  fclose (prl);
%!  compareResults (s, l, ss, ls);
%!
%!  [s, l] = so_replay_descent (prs, @getData, nStepsLong);
%!  fclose (prs);
%!  compareResults (s, l, sl, ll);

% Check compression / uncompression feature.
%!test
%!  pkg load parallel;
%!
%!  nSteps = 10;
%!  [pr, pw] = pipe ();
%!
%!  data = getDataForCompression (false);
%!  data = so_save_descent (pw, struct (), data);
%!  [s1, l1] = so_run_descent (nSteps, data.phi0, data);
%!  fclose (pw);
%!
%!  init = @(header) getDataForCompression (true);
%!  [s2, l2] = so_replay_descent (pr, init);
%!  fclose (pr);
%!
%!  % No result comparison here, since the compress / uncompress
%!  % functions introduced changes to the state structs.

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

% This is the "same" demo as for so_run_descent.  The difference is just
% that the descent is first saved and only later replayed to produce
% the actual plots.
%
%!demo
%!  pkg load parallel;
%!
%!  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);
%!  phi0 = ls_genbasic (data.g.x, "box", -3, 7);
%!
%!  printf ("Computing descent...\n");
%!  [pr, pw] = pipe ();
%!  d = data;
%!  d.handler = struct ();
%!  d = so_save_descent (pw, struct (), d);
%!  s = so_run_descent (5, phi0, d);
%!  fclose (pw);
%!  printf ("Final cost: %.6d\n", s.cost);
%!
%!  printf ("\nNow replaying...\n");
%!  d = data;
%!  d.p.verbose = true;
%!  init = @() d;
%!  s = so_replay_descent (pr, init);
%!  fclose (pr);
%!  printf ("Final cost: %.6d\n", s.cost);
