function y = dlode45(odefun,tspan,y0,theta,nvArgs)
%DLODE45  Solve ordinary differential equation (ODE) for neural ODE.
%
%    Y = DLODE45(NET,TSPAN,Y0) integrates the system of ODEs given by the
%    dlnetwork object NET. 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 must be a numeric or unformatted dlarray row vector. Y0 must be
%    a formatted dlarray.
%
%    Y = DLODE45(ODEFUN,TSPAN,Y0,THETA) integrates the system of ODEs given
%    by ODEFUN on the time interval defined by the first and last elements
%    of TSPAN, with initial conditions Y0 and parameters THETA.
%
%    ODEFUN is a function handle that accepts three input arguments T, Y,
%    and P, where T is a scalar, Y is a dlarray, and P is a set of
%    parameters, and returns a dlarray with the same size and format as Y.
%    THETA must be a dlarray, dlnetwork, cell array, struct, or table.
%    The software can compute gradients with respect to dlarray and dlnetwork
%    entries in THETA.
%    For table inputs, THETA must be a table with variables Layer,
%    Parameter and Value, where the Value variable must contain dlarray
%    objects.
%
%    Y is a formatted dlarray containing the numerical solution of the ODE
%    at the times given by TSPAN(2:end). If TSPAN contains exactly two
%    elements, then Y has the same format as Y0. If TSPAN contains more
%    than two elements, then Y has the same format as Y0 with an additional
%    appended T dimension.
%
%    The dlode45 function best suits neural ODE and custom training loop
%    workflows. To solve ODEs for other workflows, use ode45.
%
%    Y = DLODE45(_,Name=Value) specifies optional parameter name/value
%    pairs:
%
%      DataFormat                - Specify the data format for unformatted
%                                  initial conditions Y0. The format must
%                                  contain "S", "C", and "B" dimension
%                                  labels only. When using a function
%                                  handle ODEFUN it must accept unformatted
%                                  dlarray objects.
%
%      GradientMode              - Method to compute gradients with respect
%                                  to the initial conditions and parameters
%                                  when using DLGRADIENT, specified as one
%                                  of the following:
%                                  - 'direct'           - Compute gradients
%                                                         by backpropagating                                                         
%                                                         through the
%                                                         operations
%                                                         undertaken by the
%                                                         numerical solver.
%                                  - 'adjoint'          - Compute gradients
%                                                         by solving the
%                                                         associated
%                                                         adjoint ODE
%                                                         system.
%                                  - 'adjoint-seminorm' - Compute gradients
%                                                         by solving the
%                                                         associated
%                                                         adjoint ODE
%                                                         system and using
%                                                         a seminorm for
%                                                         error checks as
%                                                         described in [1].
%
%                                  The default is 'direct'.
%
%      MaxStepSize              - Maximum step size, specified as a
%                                 positive scalar or []. If MaxStepSize is
%                                 [], then the function uses a tenth of the
%                                 interval of integration size. The default
%                                 is [].
%
%      InitialStepSize          - Initial step size, specified as a
%                                 positive scalar or []. If InitialStepSize
%                                 is [], then the function automatically
%                                 determines the initial step size based on
%                                 the interval of integration and the
%                                 output of the ODE function corresponding
%                                 to the initial conditions. The default is
%                                 [].
%
%      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 apples to all components of the
%                                 solution. The default is 1e-6.
%
% Example 1:
%     % Create an initial condition, a vector of output times, and a
%     % dlnetwork to represent the ODE model.
%     y0 = dlarray(rand(4,3),"CB");
%     tspan = [0 1];
%     layers = [
%         featureInputLayer(4)
%         fullyConnectedLayer(10)
%         tanhLayer
%         fullyConnectedLayer(4)];
%     net = dlnetwork(layers);
%
%     % Solve the ODE defined by the dlnetwork net.
%     y = dlode45(net,tspan,y0);
%
% Example 2:
%     % Create an initial condition, a vector of output times, and a
%     % dlnetwork to represent the ODE model.
%     y0 = dlarray(rand(4,3),"CB");
%     tspan = [0 1];
%     layers = [
%         featureInputLayer(4)
%         fullyConnectedLayer(10)
%         tanhLayer
%         fullyConnectedLayer(4)];
%     net = dlnetwork(layers);
%     odeModel = @(t,y,net) forward(net,y);
%
%     % Solve the ODE defined by the function odeModel.
%     y = dlode45(odeModel,tspan,y0,net);
%
% Example 3:
%     % Create an initial condition, a vector of output times, a set of
%     % parameters as a dlarray, and an ODE model.
%     y0 = dlarray(rand(4,3));
%     tspan = [0 1];
%     theta = dlarray(rand(4));
%     odeModel = @(t,y,theta) tanh(theta*y);
%
%     % Solve the ODE defined by odeModel.
%     y = dlode45(odeModel,tspan,y0,theta,DataFormat="CB");
%
% Example 4:
%    % Create an initial condition, a vector of output times, a set of
%    % parameters as a structure of dlarray, and an ODE model.
%    y0 = dlarray(rand(11,11,3,10),"SSCB");
%    tspan = [0:0.05:1];
%    params = struct(Weights=dlarray(rand(5,5,3,3)), ...
%         Scale=dlarray(rand(3,1)), ...
%         Offset=dlarray(rand(3,1)));
%
%    % Solve the ODE defined by imgOdeModel.
%    y = dlode45(@imgOdeModel,tspan,y0,params);
%
%    function y = imgOdeModel(~,y,p)
%        y = dlconv(y,p.Weights,0,Padding="same");
%        y = instancenorm(y,p.Scale,p.Offset);
%        y = tanh(y);
%    end
%
% References
% [1] "Hey, that's not an ODE": Faster ODE Adjoints via Seminorms,
%     P. Kidger, R.T.Q. Chen, T. Lyons,
%     https://arxiv.org/abs/2009.09457

%   Copyright 2022-2023 The MathWorks, Inc.
arguments
    odefun {mustBeA(odefun,["dlnetwork","function_handle"])}
    tspan (1,:) {mustBeNumeric,mustBeReal,iValidateTspan}
    y0     {mustBeA(y0,'dlarray'),mustBeReal}
    theta  {iValidateTheta} = {}
    nvArgs.DataFormat char {deep.internal.dlarray.mustBeDataFormatType} = ''
    nvArgs.MaxStepSize double {mustBeScalarOrEmpty,mustBePositive} = []
    nvArgs.InitialStepSize double {mustBeScalarOrEmpty,mustBePositive} = []
    nvArgs.RelativeTolerance (1,1) double {mustBePositive} = 1.e-3
    nvArgs.AbsoluteTolerance (1,1) double {mustBePositive} = 1.e-6
    nvArgs.GradientMode (1,1) string = "direct"
end

if isa(odefun,"function_handle") && nargin<4
    error(message("deep:dlarray:DlOdeFunctionHandleNoTheta"));
end

nvArgs.GradientMode = iValidateGradientMode(nvArgs.GradientMode);

wasDlnetworkWithNonEmptyDataFormat = false;
if isa(odefun,"dlnetwork")    
    % For dlnetwork as first input argument we format the y0 according to
    % nvArgs.DataFormat.
    % To restore the original unformatted data permutation we store the
    % permutation here.
    [y0,permForDlnetInput] = deep.internal.dlarray.validateDataFormatArg(y0, nvArgs.DataFormat);
    permFormat = [];
    wasDlnetworkWithNonEmptyDataFormat = ~isempty(nvArgs.DataFormat);

    % Set DataFormat to empty so the remaining dlode45 code works like a
    % formatted case.
    nvArgs.DataFormat = '';
    if ~isempty(theta)
        error(message("nnet_cnn:dlode45:NoThetaForDlnetworkSyntax"));
    end
    theta = odefun;
    odefun = iPrepareDlnetworkInput(odefun);
end

[nDimsY0,originalSize,originalFormat] = iOriginalSizeAndFormat(y0,nvArgs.DataFormat);
if ~wasDlnetworkWithNonEmptyDataFormat
    % Validate data format.
    [y0,permFormat] = deep.internal.dlarray.validateDataFormatArg(y0, nvArgs.DataFormat);
end

% Store the format of the data to compute where the batch dim is in
% internal_dlodeBackward method.

y0Format = dims(y0);
iValidateY0Format(y0Format);

if ~isempty(nvArgs.DataFormat)
    y0 = stripdims(y0);
    y0 = ipermute(y0, permFormat);
    y0Format = nvArgs.DataFormat;
end

if isdlarray(tspan)
    t0 = parenReference(tspan,1);
    t1 = parenReference(tspan,2);
    tfinal = parenReference(tspan,numel(tspan));
else
    t0 = tspan(1);
    t1 = tspan(2);
    tfinal = tspan(end);
end

try
    % Test odefun does not error and returns a dlarray with the same size
    % and format of y0.
    if nargin(odefun) ~= 3
        error(message("deep:dlarray:DlOdeOdefunThreeInputs"));
    end

    testy = odefun(t0, y0, theta);

    testySize = size(testy,1:nDimsY0);
    testyDims = dims(testy);
    if ~isequal(testySize,originalSize)
        error(message("deep:dlarray:DlOdeOutputSizeSameAsInput",mat2str(originalSize),mat2str(testySize)));
    elseif ~isequal(testyDims,originalFormat)
        error(message("deep:dlarray:DlOdeOutputFormatSameAsInput",originalFormat,testyDims))
    end
catch caughtException
    % odefun errored internally: throw general error plus add the specific
    % cause.
    ex = MException(message("deep:dlarray:DlOdeUnableToApplyOdefun"));
    ex = addCause(ex,caughtException);
    throw(ex)
end

% Standardize type.
underlyingT0type = underlyingType(t0);
underlyingY0type = underlyingType(y0);
underlyingF0type = underlyingType(testy);

if isequal(underlyingT0type,'single') ...
        || isequal(underlyingY0type,'single') ...
        || isequal(underlyingF0type,'single')
    tspan = cast(tspan,'single');
    y0 = cast(y0,'single');
    t0 = cast(t0,'single');
    t1 = cast(t1,'single');
    tfinal = cast(tfinal,'single');
end

% Construct option object for internal ODE solver.
options = iGetOptionsForSolver(nvArgs);
isTspanDlarray  = isdlarray(tspan);

% Get a list of all of dlarray inputs in theta.
collector = deep.internal.recording.Collection;
deep.internal.basicContainerFun(@(x)iCollectFun(x,collector),theta);
collectedTheta = collector.Parameters;

useSeminorm = strcmp(nvArgs.GradientMode,"adjoint-seminorm");

if strcmp( nvArgs.GradientMode, 'adjoint' ) || useSeminorm

    % Error out if there is any non-empty state in any dlnetwork.
    iErrorIfThetaContainsStatefulNetwork(theta);

    % Create traced versions of odefun and the corresponding adjoint
    % system. We flatten the system into a vector.
    y0Size = size(y0);
    y0Flat = stripdims(y0);
    y0Flat = reshape(y0Flat,[],1);

    % To use the provided odefun we need to compose a new
    % flatOdeFun that unflattens the flattened inputs back to their
    % original shape and format for odefun, then flattens the odefun outputs.
    if isempty(dims(y0))
        unflattenFun = @(y) reshape(y,y0Size);
        flattenFun = @(y) reshape(y,[],1);
    else
        unflattenFun = @(y) dlarray(reshape(y,y0Size),y0Format);
        flattenFun = @(y) reshape(stripdims(y),[],1);
    end
    flatOdeFun = @(t,y,p) flattenFun(odefun(t,unflattenFun(y),p));

    % Trace the flatOdeFun, and the adjoint system
    tracedOdeFun = deep.AcceleratedFunction.createStandaloneInstance(flatOdeFun);
    tracedAdjointFun = deep.AcceleratedFunction.createStandaloneInstance(@deep.internal.dlode.adjointODE);

    % Perform dummy calls to cache the trace.
    initialTime = dlarray(t0);
    dummyOutput = tracedOdeFun(initialTime,y0Flat,theta); %#ok<NASGU>


    % Compute meta-data for the augmented adjoint system
    stateSize = numel(y0Flat);

    % Count learnables
    numLearnablesCollector = deep.internal.recording.Collection;
    deep.internal.networkContainerFun(@(th) iCollectNumelIfDlarray(numLearnablesCollector,th), theta);
    numLearnables = sum([numLearnablesCollector.Parameters{:}]);

    % Compute augmented state size and create dummyState.
    augStateSize = 2*stateSize + numLearnables;
    dummyState = dlarray(randn(augStateSize,1));

    % Call augDynamicsRightHandSide - this will in turn call tracedAdjoint
    % and cache its trace.
    try
        deep.internal.dlode.augDynamicsRightHandSide(...
            initialTime,dummyState,theta,flatOdeFun,stateSize,...
            numLearnables,tracedAdjointFun);
    catch e
        if strcmp(e.identifier,"deep:ad:CannotUseRecordingArraysBetweenTapes")
            error(message("nnet_cnn:dlode45:MultipleTapesDetected"));
        else
            rethrow(e);
        end
    end

    % Compute parameters to be traced.
    visitor = deep.internal.accelerate.ExtractSignatureAndParametersVisitor;
    deep.internal.container.accept({theta}, visitor);
    parameters = visitor.Parameters;

    [y0,fd] = deep.internal.dlarray.toInternal(y0Flat);
    % Extract for initial condition data.
    [fd, y0data] = extractData(fd);

    if isTspanDlarray
        [~,fd_tspan] = deep.internal.dlarray.toInternal(tspan);
        [~,tspan_data] = extractData(fd_tspan);
    else
        tspan_data = tspan;
    end

    % Call internal ode function.    
    [~,ydata] = internal_dlode(y0data,tspan_data,tracedOdeFun,tracedAdjointFun,options,...
        @nnet.internal.cnnhost.ode45,useSeminorm,...
        parameters{:});

    ydata = reshape(ydata,[originalSize,numel(tspan_data)-1]);

    targetFormat = iComputeTargetFormat(originalFormat,originalSize,size(ydata));

    fd.Layout = deep.internal.format.DataLayout(targetFormat,[]);
    fd = setDataAndLayout(fd, matlab.lang.internal.move(ydata), fd.Layout);

    y = y0;
    y = deep.internal.dlarray.fromInternal(fd,y);
else
    % Direct backpropagation of gradients.

    isY0Recorded    = ~isempty(getRecorded(y0));
    isTspanRecorded = isTspanDlarray && ~isempty(getRecorded(tspan));
    isThetaRecorded = any(cellfun(@(x) ~isempty(getRecorded(x)), collectedTheta), 'all');

    if isY0Recorded || isTspanRecorded || isThetaRecorded
        [~,fd] = deep.internal.dlarray.toInternal(y0);
        % Extract for initial condition data.
        [fd, y0data] = extractData(fd);

        if isTspanDlarray
            [tspan,fd_tspan] = deep.internal.dlarray.toInternal(tspan);
            [~,tspan_data] = extractData(fd_tspan);
        else
            tspan_data = tspan;
        end

        theta_data = cellfun(@iExtractInternalDlarrayAndFormat, collectedTheta, 'UniformOutput', false);

        % Record direct identity operation.
        [tspan_recording,y0_recording] = internal_dlode_direct_identity(tspan_data,y0data,"dlode45",theta_data{:});

        if isY0Recorded
            fd.Layout = deep.internal.format.DataLayout(originalFormat,[]);
            fd = setDataAndLayout(fd, matlab.lang.internal.move(y0_recording), fd.Layout);
            y0 = deep.internal.dlarray.fromInternal(fd,y0);
        end

        if isTspanDlarray
            fd_tspan.Layout = deep.internal.format.DataLayout('',[]);
            fd_tspan = setDataAndLayout(fd_tspan, matlab.lang.internal.move(tspan_recording), fd_tspan.Layout);
            tspan = deep.internal.dlarray.fromInternal(fd_tspan,tspan);
        end
    end

    [~,ydata] = internal_dlode_direct(odefun,tspan,stripdims(y0),options,...
        originalSize,originalFormat,theta,@nnet.internal.cnnhost.ode45,...
        t0,t1,tfinal);

    targetFormat = iComputeTargetFormat(originalFormat,originalSize,size(ydata));

    y = dlarray(ydata,targetFormat);
end

if wasDlnetworkWithNonEmptyDataFormat
    % In the case that odefun was originally a dlnetwork, and DataFormat
    % was originally non-empty we return the solution to the user's
    % original permutation, potentially with an appended T dimension.
    y = stripdims(y);
    % Permute leading dimensions to user format, leave time dimension last.
    perm = [permForDlnetInput,numel(permForDlnetInput)+1];
    y = ipermute(y,perm);
end
end

function [nDimsY0,originalSize,originalFormat] = iOriginalSizeAndFormat(y0,dataFormat)
% Original size and format to be restored when calling the provided
% function handle in the ODE solver.
nDimsY0 =  max(ndims(stripdims(y0)), strlength(dataFormat) );
originalSize   = size(stripdims(y0), 1:nDimsY0);
originalFormat = dims(y0);
end

function odefun = iPrepareDlnetworkInput(net)
% Ensure network has exactly 1 or 2 inputs and is initialized.
inputNames = net.InputNames;
numIn = numel(inputNames);
if numIn>2
    error(message("nnet_cnn:dlode45:TooManyNetworkInputs"));
end
if ~net.Initialized
    error(message("nnet_cnn:dlode45:NetworkUnitialized"))
end

% Construct function_handle-s.
if numIn == 2
    validTInput = false;
    exampleInputs = net.getExampleInputs();
    if ~isempty(exampleInputs)
        firstInputSize = size(exampleInputs{1});
        firstInputDims = dims(exampleInputs{1});
        validTInput = firstInputSize(1)==1 && isequal(firstInputDims,'CB');
    else        
        firstInputLayerName = net.InputNames{1};
        firstInputLayer = net.Layers(strcmp(firstInputLayerName,{net.Layers.Name}));
        % Only feature and generic inputs can be used for t.
        isFeature = isa(firstInputLayer,'nnet.cnn.layer.FeatureInputLayer');
        isGeneric = isa(firstInputLayer,'nnet.cnn.layer.InputLayer');
        inputSize = firstInputLayer.InputSize;
        if isFeature     
            % Check size in C dim
            validTInput = isequal(inputSize(1),1);
        end
        if isGeneric
            firstInputDims = firstInputLayer.InputFormat;
            isCB = isequal(firstInputDims,'CB');
            isBC = isequal(firstInputDims,'BC');
            if isCB
                validTInput = isequal(inputSize(1),1);
            end
            if isBC
                validTInput = isequal(inputSize(2),1);
            end
        end
    end
    if ~validTInput
        error(message("nnet_cnn:dlode45:FirstInputNotScalar"));
    end
    odefun = @iExpandTAndPredict;
else
    odefun = @(t,y,net) net.predict(y);
end
end

function targetFormat = iComputeTargetFormat(originalFormat,originalSize,ydataSize)
% Append 'T' label in the case of formatted output and tspan has more than
% two elements and hence ydata has an additional dimension with respect to
% originalSize.
targetFormat = originalFormat;
if ~isempty(originalFormat) && ~isequal(ydataSize, originalSize)
    targetFormat = [targetFormat, 'T'];
end
end

function [data, format] = iExtractInternalDlarrayAndFormat(theta)
[~,fd] = deep.internal.dlarray.toInternal(theta);
[~,data] = extractData(fd);
format = fd.DimensionMetadata;
end

function iCollectFun(x,collector)
if isdlarray(x)
    collect(collector,x);
end
end

function iValidateY0Format(fmt)
if any(fmt == 'U') || any(fmt == 'T')
    error(message("deep:dlarray:DlOdeY0Fmt"));
end
end

function iValidateTspan(tspan)
isDlarray = isdlarray(tspan);
if isDlarray && ~isempty(dims(tspan))
    error(message("deep:dlarray:DlOdeTspanFormatted"));
elseif isDlarray
    tspan = extractdata(tspan);
end
if ~(numel(tspan) >= 2 && issorted(tspan,'strictmonotonic'))
    error(message("deep:dlarray:DlOdeInvalidTspan"));
end
end

function options = iGetOptionsForSolver(nvArgs)
options = odeset('RelTol',nvArgs.RelativeTolerance,'AbsTol',nvArgs.AbsoluteTolerance);
options = odeset(options,'MaxStep',nvArgs.MaxStepSize);
options = odeset(options,'InitialStep',nvArgs.InitialStepSize);
options = odeset(options,"Refine",1);
end

function gradientMode = iValidateGradientMode(gradientMode)
gradientMode = validatestring(gradientMode, ["direct", "adjoint","adjoint-seminorm"]);
end

function y = iExpandTAndPredict(t,y,theta)

% Expand t over batch dimension and cast to formatted dlarray.
bdim = finddim(y,'B');
if isempty(bdim)
    batchSize = 1;
else
    batchSize = size(y,bdim);
end
t = repmat(t,1,batchSize);
t = dlarray(t,"CB");

% Call net.predict
y = theta.predict(t,y);
end

function iErrorIfThetaContainsStatefulNetwork(theta)
try
    deep.internal.basicContainerFun(@iCheckIfThetaElementIsStatefulDlnetwork,theta);
catch ex
    throw(ex);
end
end

function iCheckIfThetaElementIsStatefulDlnetwork(th)
if isa(th,"dlnetwork") && size(th.State,1)>0
    error(message("deep:dlarray:DlOdeInvalidStatefulNet"));
end
end

function iCollectNumelIfDlarray(collector,x)
if isdlarray(x)
    collector.collect(numel(x));
end
end

function iValidateTheta(theta)
if ...
  ~(isdlarray(theta) || ...
  isa(theta,"dlnetwork") || ...
  iscell(theta) || ...
  isstruct(theta) || ...
  istable(theta))
    error(message('deep:dlarray:DlOdeInvalidParameter'));
end
end