%# Copyright (C) 2015-2016, Jacopo Corno <jacopo.corno@gmail.com>
%# Copyright (C) 2013, Roberto Porcu' <roberto.porcu@polimi.it>
%# Copyright (C) 2006-2012, Thomas Treichl <treichl@users.sourceforge.net>
%# OdePkg - A package for solving ordinary differential equations and more
%#
%# 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 2 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{t}, @var{y}] =} odeRATTLE (@var{fun}, @var{constr}, @var{trange}, @var{init})
%# @deftypefnx {Function File} {[@var{t}, @var{y}] =} odeRATTLE (@var{fun}, @var{constr}, @var{trange}, @var{init}, @var{ode_opt})
%# @deftypefnx {Function File} {[@var{t}, @var{y}] =} odeRATTLE (@dots{}, @var{par1}, @var{par2}, @dots{})
%# @deftypefnx {Function File} {[@var{t}, @var{y}, @var{te}, @var{ye}, @var{ie}] =} odeRATTLE (@dots{})
%# @deftypefnx {Function File} {@var{solution} =} odeRATTLE (@dots{})
%#
%# This function is a geometric integrator for constrained Hamilton's equations
%# of motion that uses the RATTLE method of order 2 to find the new solution,
%# step by step.
%#
%# @var{fun} is a function_handle or an inline function that
%# defines the system
%# @ifhtml
%# @example
%# @math{q' =  dH/dp (t,[q;p])}
%# @math{p' = -dH/dq (t,[q;p])}
%# @end example
%# @end ifhtml
%# @ifnothtml
%# @math{q' =  dH/dp (t,[q;p])}
%# @math{p' = -dH/dq (t,[q;p])},
%# @end ifnothtml
%# where @math{H(t,[q;p])} is the Hamiltonian of the problem, @math{q} are
%# generalized coordinates and @math{p} are generalized momenta. This function
%# must take two input arguments: the time and the unknowns. It must return a
%# vector or a cell (if the Hamiltonian Hessian is passed to speedup
%# computations) containing as first element the Hamilton's equations of motion
%# (the previous system of equations) and as second element the Hessian of the
%# Hamiltonian.
%#
%# @var{constr} is a function_handle or an inline function
%# that defines the constraints of the problem @math{g(q)=0}. This
%# function must take one input argument and must return a vector or a cell (if
%# the constraint gradient and constraint Hessian are passed) containing as
%# first element the equations defining the holonomic constraints and as second
%# and third elements respectively the gradient and the Hessian of this
%# constraints function (this is not mandatory but it can speedup the
%# computation and improve accuracy).
%#
%# @var{trange} specifies the time interval over which the ODE will be
%# evaluated.  Typically, it is a two-element vector specifying the initial and
%# final times (@code{[tinit, tfinal]}).  If there are more than two elements
%# then the solution will also be evaluated at these intermediate time
%# instances using linear interpolation.
%#
%# By default, @code{odeRATTLE} uses an adaptive timestep with the
%# @code{integrate_adaptive} algorithm.  The tolerance for the timestep
%# computation may be changed by using the options @code{"RelTol"}
%# and @code{"AbsTol"}.
%#
%# @var{init} contains the initial value for the unknowns.  It is a row or
%# column vector containing the generalized coordinates in its first half and
%# the generalized momenta in the second half.  Each momentum correspond to the
%# coordinate with the same relative position inside the vector of initial
%# conditions.
%#
%# The optional fourth argument @var{ode_opt} specifies non-default options to
%# the ODE solver. It is a structure generated by @code{odeset}.  @code{ode12}
%# will ignore the following options: "BDF", "InitialSlope", "Jacobian",
%# "JPattern", "MassSingular", "MaxOrder", "MvPattern", "Vectorized".  The user
%# can set extra options by adding fields to the structure.  @code{odeRATTLE}
%# allows for:
%#
%# @table @asis
%# @item  NewtonTol
%# the tolerance for the Newton algorithm [Default = 1e-7].
%#
%# @item  MaxNewtonIterations
%# maximum number of Newton iterations [Default = 100].
%#
%# @item  InexactSolver
%# the Newton solver. It can be set to either ["newton_raphson"],
%# "inexact_newton" or "fsolve".
%#
%# @item  ConstraintFcn
%# function_handle representing the constraints of the system.
%#
%# @item  ConstraintGradFcn
%# function_handle representing the gradient of the constraints of the system.
%#
%# @item  ConstraintHessFcn
%# function_handle representing the Hessian of the Hamiltonian of the system,
%#
%# @item  ConstraintsNb
%# positive integer equal to the number of constraints of the system.
%#
%# @item  HamiltonianHessFcn
%# function_handle representing the Hessian of the Hamiltonian of the system.
%# @end table
%#
%# If the "inexact_newton" solver is set, a further set of options can be
%# provided:
%#
%# @table @asis
%# @item  UseJacobian
%# It can be set to either "yes" or ["no].  In the first case, the Hamiltonian
%# Hessian, the constraint gradient and the constraint Hessian will be used.
%# If they are not provided, they will be approximated by the utility
%# functions @command{approx_Hamilt_Hess}, @command{approx_Constr_grad}
%# and @command{approx_Constr_Hess}.
%#
%# @item  Eta
%# initial forcing term (must be in the interval [0,1)).  The dafault value is
%# set to 0.5.  For details see [1].
%#
%# @item  Choice
%# select the forcing term. It can be 1 (Default) or 2.  For details see [1].
%#
%# @item  Algorithm
%# iterative method to solve the linearized system.  It can be ["gmres"], "pcg"
%# or "bicgstab".
%#
%# @item  Restart
%# restart parameter for the GMRES solver (ignored for other solvers).
%# [Default = 20].
%# @end table
%#
%# The function typically returns two outputs.  Variable @var{t} is a
%# column vector and contains the times where the solution was found.  The
%# output @var{y} is a matrix in which each column refers to a different
%# unknown of the problem and each row corresponds to a time in @var{t}.
%#
%# The output can also be returned as a structure @var{solution} which
%# has field @var{x} containing the time where the solution was evaluated and
%# field @var{y} containing the solution matrix for the times in @var{x}.
%# Use @code{fieldnames (@var{solution})} to see the other fields and additional
%# information returned.
%#
%# If using the @code{"Events"} option then three additional outputs may
%# be returned.  @var{te} holds the time when an Event function returned a
%# zero.  @var{ye} holds the value of the solution at time @var{te}.  @var{ie}
%# contains an index indicating which Event function was triggered in the case
%# of multiple Event functions.
%#
%# References:
%#
%# [1] S.C. Eisenstat and H.F. Walker, "Choosing the Forcing Terms in an Inexact
%# Newton Method." SIAM Journal on Scientific Computing, 17(1), pp. 16-16, 1996.
%# @seealso{odeset, odeSE, odeSPVI, odeSV, odeVV}
%# @end deftypefn

function varargout = odeRATTLE (fun, constraint, trange, init, varargin)

  if (nargin < 4)
    print_usage ();
  end

  order = 2;
  solver = 'odeRATTLE';

  if (nargin >= 5)
    if (! isstruct (varargin{1}))
      %# varargin{1:len} are parameters for fun
      odeopts = odeset ();
      funarguments = varargin;
    elseif (length (varargin) > 1)
      %# varargin{1} is an ODE options structure opt
      odeopts = varargin{1};
      funarguments = {varargin{2:length(varargin)}};
    else  # if (isstruct (varargin{1}))
      odeopts = varargin{1};
      funarguments = {};
    end
  else  # nargin == 4
    odeopts = odeset ();
    funarguments = {};
  end

  if (! isnumeric (trange) || ! isvector (trange))
    error ("Octave:invalid-input-arg",
           "odeRATTLE: TRANGE must be a numeric vector");
  end

  if (length (trange) < 2)
    error ("Octave:invalid-input-arg",
           "odeRATTLE: TRANGE must contain at least 2 elements");
  elseif (trange(1) == trange(2))
    error ("Octave:invalid-input-arg",
           "odeRATTLE: invalid time span, TRANGE(1) == TRANGE(2)");
  else
    direction = sign (trange(2) - trange(1));
  end
  trange = trange(:);

  if (! isnumeric (init) || ! isvector (init))
    error ("Octave:invalid-input-arg",
           "odeRATTLE: INIT must be a numeric vector");
  end
  init = init(:);

  if (ischar (fun))
    try
      fun = str2func (fun);
    catch
      warning (lasterr);
    end
  end
  if (! isa (fun, "function_handle"))
    error ("Octave:invalid-input-arg",
           "odeRATTLE: FUN must be a valid function handle");
  end

  if ! (isa (constraint, "function_handle") || isa (constraint, "inline"))
    error ("Octave:invalid-input-arg",
           "odeRATTLE: CONSTR must be a valid function handle");
  end

  %# Start preprocessing, have a look which options are set in odeopts,
  %# check if an invalid or unused option is set
  [defaults, classes, attributes] = odedefaults (numel (init),
                                                 trange(1), trange(end));

  persistent odeRATTLE_ignore_options = ...
    {"BDF", "InitialSlope", "Jacobian", "JPattern",
     "MassSingular", "MaxOrder", "MvPattern", "Vectorized"};

  defaults   = rmfield (defaults, odeRATTLE_ignore_options);
  classes    = rmfield (classes, odeRATTLE_ignore_options);
  attributes = rmfield (attributes, odeRATTLE_ignore_options);

  % Specific options for odeRATTLE
  defaults.NewtonTol   = 1e-7;
  classes.NewtonTol    = {"float"};
  attributes.NewtonTol = {"scalar", "positive"};
  defaults.MaxNewtonIterations   = 100;
  classes.MaxNewtonIterations    = {"float"};
  attributes.MaxNewtonIterations = {"scalar", "integer", "positive"};
  defaults.InexactSolver   = "newton_raphson";
  classes.InexactSolver    = {"char"};
  attributes.InexactSolver = {"newton_raphson", "inexact_newton"};
  # Specific options for Inexact Newton solver (see inexact_newton.m)
  defaults.UseJacobian   = "no";
  classes.UseJacobian    = {"char"};
  attributes.UseJacobian = {"yes", "no"};
  defaults.Eta   = 0.5;
  classes.Eta    = {"float"};
  attributes.Eta = {"scalar", ">=", 0, "<", 1};
  defaults.Choice   = 1;
  classes.Choice    = {"float"};
  attributes.Choice = {">=", 1, "<=", 2, "integer"};
  defaults.Algorithm   = "gmres";
  classes.Algorithm    = {"char"};
  attributes.Algorithm = {"gmres", "pcg", "bicgstab"};
  defaults.Restart   = 20;
  classes.Restart    = {"float"};
  attributes.Restart = {"integer"};
  % Other options (see rattle.m)
  defaults.ConstraintFcn   = [];
  classes.ConstraintFcn    = {"function_handle"};
  attributes.ConstraintFcn = {};
  defaults.ConstraintGradFcn   = [];
  classes.ConstraintGradFcn    = {"function_handle"};
  attributes.ConstraintGradFcn = {};
  defaults.ConstraintHessFcn   = [];
  classes.ConstraintHessFcn    = {"function_handle"};
  attributes.ConstraintHessFcn = {};
  defaults.ConstraintsNb   = [];
  classes.ConstraintsNb    = {"function_handle"};
  attributes.ConstraintsNb = {};
  defaults.HamiltonianHessFcn   = [];
  classes.HamiltonianHessFcn    = {"function_handle"};
  attributes.HamiltonianHessFcn = {};

  odeopts = odemergeopts ("odeRATTLE", odeopts, defaults, classes, attributes);

  odeopts.funarguments = funarguments;
  odeopts.direction    = direction;

  if (! isempty (odeopts.NonNegative))
    if (isempty (odeopts.Mass))
      odeopts.havenonnegative = true;
    else
      odeopts.havenonnegative = false;
      warning ("Octave:invalid-input-arg",
               ["odeRATTLE: option 'NonNegative' is ignored", ...
                " when mass matrix is set\n"]);
    end
  else
    odeopts.havenonnegative = false;
  end

  if (isempty (odeopts.OutputFcn) && nargout == 0)
    odeopts.OutputFcn = @odeplot;
    odeopts.haveoutputfunction = true;
  else
    odeopts.haveoutputfunction = ! isempty (odeopts.OutputFcn);
  end

  if (isempty (odeopts.InitialStep))
    odeopts.InitialStep = odeopts.direction * ...
                          starting_stepsize (order, fun, trange(1),
                                             init, odeopts.AbsTol,
                                             odeopts.RelTol,
                                             strcmp (odeopts.NormControl,
                                             "on"), odeopts.funarguments);
  end

  if (! isempty (odeopts.Mass) && isnumeric (odeopts.Mass))
    havemasshandle = false;
    mass = odeopts.Mass;  # constant mass
  elseif (isa (odeopts.Mass, "function_handle"))
    havemasshandle = true;    # mass defined by a function handle
  else  # no mass matrix - creating a diag-matrix of ones for mass
    havemasshandle = false;   # mass = diag (ones (length (init), 1), 0);
  end

  if (havemasshandle)   # Handle only the dynamic mass matrix,
    if (! strcmp (odeopts.MStateDependence, "none")) # constant mass matrices have already
      mass = @(t,x) odeopts.Mass (t, x, odeopts.funarguments{:});
      fun = @(t,x) mass (t, x, odeopts.funarguments{:}) ...
             \ fun (t, x, odeopts.funarguments{:});
    else                 # if ((! strcmp (odeopts.MStateDependence, "none")) == false)
      mass = @(t) odeopts.Mass (t, odeopts.funarguments{:});
      fun = @(t,x) mass (t, odeopts.funarguments{:}) ...
             \ fun (t, x, odeopts.funarguments{:});
    end
  end

  %# Starting the initialisation of the core solver odeRATTLE
  dim = length (init) / 2;

  flag = 0;
  try
    fun (trange(1), init){1};
  catch
    flag = 1;
  end

  if (flag == 0)
    fun_new = @(t,x) fun (t, x){1};
  else
    fun_new = fun;
  end

  flag = 0;
  try
    constraint (init(1:dim)){1};
  catch
    flag = 1;
  end

  if (flag == 0)
    constraint_new = @(q) constraint (q){1};
  else
    constraint_new = constraint;
  end

  odeopts.ConstraintFcn = constraint_new;
  odeopts.ConstraintsNb = length (constraint_new (init(1:dim)));

  flag = 0;
  try
    constraint (init(1:dim)){2};
  catch
    flag = 1;
    warning ("OdePkg:InvalidArgument", ...
             "gradient of constraint will be approximated");
  end

  if (flag == 1)
    GradConstrainFcn = @(q) approx_Constr_grad (constraint, q);
  else
    GradConstrainFcn = @(q) constraint (q){2};
  end

  odeopts.ConstraintGradFcn = GradConstrainFcn;

  use_Jacobian = odeget (odeopts, "UseJacobian", "no", "fast_not_empty");
  use_Jacobian = lower (deblank (use_Jacobian));
  if (strcmp (use_Jacobian, "yes"))
    flag = 0;
    try
      fun (trange(1), init){2};
    catch
      flag = 1;
      warning ("OdePkg:InvalidArgument", ...
               "consider returning the Hessian of Hamiltonian in function ", ...
               "''%s'', this may speedup the program", func2str (fun));
    end

    if (flag == 1)
      Hessian = @(t,x) approx_Hamilt_Hess (fun_new, t, x);
    else
      Hessian = @(t,x) fun(t, x){2};
    end

    odeopts.HamiltonianHessFcn = Hessian;

    flag = 0;
    try
      constraint (init(1:dim)){3};
    catch
      flag = 1;
      warning ("OdePkg:InvalidArgument", ...
               "consider returning the Hessian of constraint in function ", ...
               "''%s'', this may speedup the program", func2str (constraint));
    end

    if (flag == 1)
      ConstraintHess = @(q) approx_Constr_Hess (GradConstrainFcn, q);
    else
      ConstraintHess = @(q) constraint (q){3};
    end

    odeopts.ConstraintHessFcn = ConstraintHess;
  end

  if (strcmp (odeopts.InexactSolver, "inexact_newton"))
    odeopts.solver = @inexact_newton;
  else
    odeopts.solver = @fsolve;
  end
  
  #SobOpts.Tau = odeget (odeopts, "Tau", 1.e-6, "fast_not_empty");
  solution = integrate_adaptive (@rattle, order, fun_new, trange,
                                 init, odeopts);

  %# Postprocessing, do whatever when terminating integration algorithm
  if (odeopts.haveoutputfunction)  # Cleanup plotter
    feval (odeopts.OutputFcn, [], [], "done", odeopts.funarguments{:});
  end
  if (! isempty (odeopts.Events))   # Cleanup event function handling
    ode_event_handler (odeopts.Events, solution.t(end),
                       solution.x(end,:).', "done", odeopts.funarguments{:});
  end

  %# Print additional information if option Stats is set
  if (strcmp (odeopts.Stats, "on"))
    havestats = true;
    nsteps    = solution.vcntloop-2;              # vcntloop from 2..end
    nfailed   = (solution.vcntcycles-1)-nsteps+1; # vcntcycl from 1..end
    ndecomps  = 0;  # number of LU decompositions
    npds      = 0;  # number of partial derivatives
    nlinsols  = 0;  # no. of solutions of linear systems
    %# Print cost statistics if no output argument is given
    if (nargout == 0)
      printf ("Number of successful steps: %d\n", nsteps);
      printf ("Number of failed attempts:  %d\n", nfailed);
    end
  else
    havestats = false;
  end

  if (nargout == 2)
    varargout{1} = solution.t;      # Time stamps are first output argument
    varargout{2} = solution.x;      # Results are second output argument
  elseif (nargout == 1)
    varargout{1}.x = solution.t.';    # Time stamps are saved in field x
    varargout{1}.y = solution.x.';    # Results are saved in field y
    varargout{1}.solver = solver;   # Solver name is saved in field solver
    if (! isempty (odeopts.Events))
      varargout{1}.ie = solution.event{2};  # Index info which event occurred
      varargout{1}.xe = solution.event{3};  # Time info when an event occurred
      varargout{1}.ye = solution.event{4};  # Results when an event occurred
    end
    if (havestats)
      varargout{1}.stats = struct ();
      varargout{1}.stats.nsteps   = nsteps;
      varargout{1}.stats.nfailed  = nfailed;
      varargout{1}.stats.nfevals  = nfevals;
      varargout{1}.stats.npds     = npds;
      varargout{1}.stats.ndecomps = ndecomps;
      varargout{1}.stats.nlinsols = nlinsols;
    end
  elseif (nargout == 5)
    varargout = cell (1,5);
    varargout{1} = solution.t;
    varargout{2} = solution.x;
    if (! isempty (odeopts.Events))
      varargout{3} = solution.event{3};  # Time info when an event occurred
      varargout{4} = solution.event{4};  # Results when an event occurred
      varargout{5} = solution.event{2};  # Index info which event occurred
    end
  end

end

%! # We are using the "armonic oscillator" implementation for all tests that
%! # are done for this function.
%!
%!function [ydot] = pendulum (vt, vy)
%!  ydot = [0.5*vy(length(vy)/2+1:end); -(2*5*9.806)*[zeros(length(vy)/4);ones(length(vy)/4)]];
%!function [ydot] = pendulum_with_H (vt, vy)
%!  Hamilton = [0.5*vy(length(vy)/2+1:end); -(2*5*9.806)*[zeros(length(vy)/4);ones(length(vy)/4)]];
%!  Hess = [zeros(length(vy)/2,length(vy));zeros(length(vy)/2),0.5*eye(length(vy)/2)];
%!  ydot = {Hamilton,Hess};
%!function [g] = constraint (vy)
%!  constr = (vy(:)')*vy(:) - 25.0;
%!  grad = 2.0*vy(:)';
%!  g = {constr,grad};
%!
%!shared m, L, g
%!  m = 2.0; 
%!  L = 5.0; 
%!  g = 9.806;
%!
%! %# Turn off output of warning messages for all tests, turn them on
%! %# again if the last test is called
%!error %# output argument
%!  warning ('off', 'OdePkg:InvalidArgument');
%!  B = odeRATTLE (1, @constraint, [0 1], [L*cos(-pi/4) L*sin(-pi/4) 0 0]);
%!error %# input argument number one
%!  [vt, vy] = odeRATTLE (1, @constraint, [0 1], [L*cos(-pi/4) L*sin(-pi/4) 0 0]);
%!error %# input argument number two
%!  [vt, vy] = odeRATTLE (@pendulum, 2, [0 1], [L*cos(-pi/4) L*sin(-pi/4) 0 0]);
%!error %# input argument number three
%!  [vt, vy] = odeRATTLE (@pendulum, @constraint, 0, [L*cos(-pi/4) L*sin(-pi/4) 0 0]);
%!error %# input argument number four
%!  [vt, vy] = odeRATTLE (@pendulum, @constraint,  [0 1], [L*cos(-pi/4) L*sin(-pi/4) 0 0 0]);
%!test %# two output arguments
%!  [vt, vy] = odeRATTLE (@pendulum, @constraint, [0 0.4], [L*cos(-pi/4) L*sin(-pi/4) 0 0]);
%!  assert ([vt(end), sqrt(vy(end,1)^2+vy(end,2)^2)], [0.4, L], 1e-3);
%!test %# extra input arguments passed through
%!  [vt, vy] = odeRATTLE (@pendulum, @constraint, [0 0.4], [L*cos(-pi/4) L*sin(-pi/4) 0 0], 12, 13, 'KL');
%!  assert ([vt(end), sqrt(vy(end,1)^2+vy(end,2)^2)], [0.4, L], 1e-2);
%!test %# empty OdePkg structure *but* extra input arguments
%!  opt = odeset();
%!  [vt, vy] = odeRATTLE (@pendulum, @constraint, [0 0.4], [L*cos(-pi/4) L*sin(-pi/4) 0 0], opt, 12, 13, 'KL');
%!  assert ([vt(end), sqrt(vy(end,1)^2+vy(end,2)^2)], [0.4, L], 1e-2);
%!test  # Solve with intermidiate step
%!  [vt, vy] = odeRATTLE (@pendulum, @constraint, [0 0.2 0.4], [L*cos(-pi/4) L*sin(-pi/4) 0 0]);
%!  assert ([vt(end), sqrt(vy(end,1)^2+vy(end,2)^2)], [0.4, L], 1e-2);
%!test %# use_jacobian option
%!  opt = odeset();
%!  opt.UseJacobian = 'yes';
%!  [vt, vy] = odeRATTLE (@pendulum_with_H, @constraint, [0 0.4], [L*cos(-pi/4) L*sin(-pi/4) 0 0], opt);
%!  assert ([vt(end), sqrt(vy(end,1)^2+vy(end,2)^2)], [0.4, L], 1e-2);
%!test %# use approximated jacobian
%!  opt = odeset();
%!  opt.UseJacobian = 'yes';
%!  [vt, vy] = odeRATTLE (@pendulum, @constraint, [0 0.4], [L*cos(-pi/4) L*sin(-pi/4) 0 0], opt);
%!  assert ([vt(end), sqrt(vy(end,1)^2+vy(end,2)^2)], [0.4, L], 1e-2);
%!
%!  warning ('on', 'OdePkg:InvalidArgument');

%# Local Variables: ***
%# mode: octave ***
%# End: ***
