function layer = neuralODELayer(net,tspan,args)
% neuralODELayer   Neural ODE Layer
%
%   layer = neuralODELayer(net,tspan) creates a neural ODE layer. net is a
%   dlnetwork that defines the neural ODE system. If net has 1 input, then
%   predict(net,y) defines the ODE system. If net has 2 inputs, then
%   predict(net,t,y) defines the ODE system, where t is a timestep repeated
%   over the batch dimension. tspan is the integration time interval.
%
%   layer = neuralODELayer(net,tspan,Name=Value) specifies the additional
%   options using one or more name-value arguments:
%
%      GradientMode              - Mode for computing gradients, specified 
%                                  as one of these values:
%                                    - "direct"  - Compute gradients by
%                                                  backpropagating through
%                                                  the operations undertaken
%                                                  by the numerical solver.
%                                    - "adjoint" - Compute gradients by 
%                                                  solving the associated 
%                                                  adjoint ODE system.
%                                 The default is "direct".
%
%      RelativeTolerance        - Relative error tolerance, specified as a
%                                 positive scalar. The relative tolerance
%                                 applies to all components of the
%                                 solution. The default is 1e-3.
%
%      AbsoluteTolerance        - Absolute error tolerance, specified as a
%                                 positive scalar. The absolute error
%                                 tolerance applies to all components of the
%                                 solution. The default is 1e-6.
%
%      Name                     - Name for the layer.
%                                 The default is "".
%
%   Example:
%      % Create a neural ODE layer to solve a one dimensional ODE defined
%      % by a two layer feedforward network with hidden size 10.
%      % Define the ODE with a dlnetwork.
%      inputSize = 1;
%      hiddenSize = 10;
%      odenet = [
%        featureInputLayer(inputSize)
%        fullyConnectedLayer(hiddenSize)
%        tanhLayer
%        fullyConnectedLayer(inputSize)];
%      odenet = dlnetwork(odenet);
%      % Create a neuralODELayer to solve the ODE on the time interval
%      % [0,1].
%      tspan = [0,1];
%      odeLayer = neuralODELayer(odenet,tspan);

% Copyright 2022-2023 The MathWorks, Inc.
arguments
    net (1,1) {mustBeA(net,"dlnetwork")}
    tspan (1,:) {mustBeNumeric,iMustBeTwoOrMoreElements}
    args.Name {iAssertValidLayerName} = ''
    args.GradientMode {mustBeTextScalar} = 'direct'
    args.RelativeTolerance (1,1) {mustBeNumeric,mustBePositive} = 1.e-3
    args.AbsoluteTolerance (1,1) {mustBeNumeric,mustBePositive} = 1.e-6
end

tspan = iConvertToDouble(tspan);
iMustBeMonotone(tspan);
args = iValidateAndConvertToCanonical(args);
args = namedargs2cell(args);
layer = nnet.cnn.layer.NeuralODELayer(net,tspan,args{:});
end

function iAssertValidLayerName(name)
nnet.internal.cnn.layer.paramvalidation.validateLayerName(name)
end

function gmode = iValidateGradientMode(gmode)
gmode = validatestring(gmode,{'direct','adjoint'});
end

function args = iValidateAndConvertToCanonical(args)
args = nnet.internal.cnn.layer.util.gatherParametersToCPU(args);
args.Name = char(args.Name);
args.GradientMode = iValidateGradientMode(args.GradientMode);
args.RelativeTolerance = iConvertToDouble(args.RelativeTolerance);
args.AbsoluteTolerance = iConvertToDouble(args.AbsoluteTolerance);
end

function iMustBeTwoOrMoreElements(val)
% Verify val contains at least 2 elements.
if numel(val)<2
    id = "deep:dlarray:DlOdeInvalidTspan";
    e = MException(id,message(id));
    throwAsCaller(e);
end
end

function iMustBeMonotone(val)
% Verify that a vector val is strictly monotonic
d = diff(val);
increasing = all(d>0);
decreasing = all(d<0);
if ~(increasing || decreasing)
    id = "deep:dlarray:DlOdeInvalidTspan";
    e = MException(id,message(id));
    throwAsCaller(e);
end
end

function value = iConvertToDouble(value)
value = nnet.internal.cnn.layer.util.convertToDouble(value);
end