function varargout = dlgradient(y,varargin)
%DLGRADIENT Compute gradient of dlarray object
%   [dX1, dX2, ... dXn] = DLGRADIENT(Y, X1, X2, ..., Xn) computes gradients
%   of scalar dlarray object Y. Each dXi is the gradient of Y with respect
%   to Xi. Call DLGRADIENT from inside the function passed to DLFEVAL. The
%   X inputs can be traced dlarray objects, cell arrays, structures of
%   traced dlarray objects, or tables containing traced dlarrays in a Value
%   variable, recursively.
%
%   Some operations on dlarray result in untraced values. For example,
%   EXTRACTDATA and SIZE are untraced. Gradients may not be computed
%   correctly if intermediate values are untraced. The outputs of DLGRADIENT
%   are traced if 'EnableHigherDerivatives' name-value pair argument below
%   is true and untraced otherwise.
%
%    __ = DLGRADIENT(__,Name,Value) specifies additional options using one
%   or more name-value pair arguments.
%
%   Name-Value Pair Arguments:
%     'RetainData'       - Flag to retain intermediate values for reuse in subsequent DLGRADIENT
%                          calls, specified as true or false. If 'EnableHigherDerivatives'
%                          name-value pair argument is true then the function retains data
%                          automatically and this option has no effect. The default is false.
%
%     'EnableHigherDerivatives'
%                        - Flag to enable computing higher order derivatives, specified as true
%                          or false. If true, then the software traces the backward pass so that
%                          returned gradients are traced. When true, intermediate values are
%                          retained and 'RetainData' name-value pair argument has no effect.
%                          If evaluating a deep.AcceleratedFunction object, the default is true.
%                          Otherwise, the default is false.
%
%     'AllowComplex'
%                        - Flag to enable gradients with respect to complex
%                          Xi and for gradients dXi to be complex. If true,
%                          dXi can be complex even if input Xi is real. If
%                          false, Xi must be real and outputs dXi are all
%                          real. The default is true.
%
%   Examples:
% Example 1
%     % Evaluate gradient of a function
%     a = dlarray([0 0]); % point at which to evaluate gradient
%     [val,grad] = dlfeval(@objectiveAndGradient,a);
%     function [y,grad] = objectiveAndGradient(x)
%         y = x(1).^2 + x(2) - 2;
%         grad = dlgradient(y,x);
%     end
%
% Example 2
%     % Evaluate gradients of a non-scalar valued function
%     a = dlarray([1 2; 3 4]);
%     scaleFactors = dlarray([0 0; 0 1]);
%     [val,grad] = dlfeval(@objectiveAndGradient,a,scaleFactors);
%     % val is dlarray([1 4; 9 16])
%     % grad is dlarray([0 0; 0 8])
%     function [y,grad] = objectiveAndGradient(x,scaleFactors)
%         y = x.^2;
%         s = sum(scaleFactors.*y,'all');
%         grad = dlgradient(s,x);
%     end
%
% Example 3
%     % Evaluate gradients of complex functions
% 
%     % Evaluate gradient of complex function with complex and real inputs
%     x = dlarray([2+3i 5]);
%     grad = dlfeval(@complexGrad,x);
%     % grad is dlarray([2+3i, 2+3i])
%     
%     % Evaluate gradient allowing only real inputs and returning only real
%     % gradients. Intermediate values can still be complex.
%     x = dlarray([2 5]);
%     gradReal = dlfeval(@complexWithRealGradients,x);
%     % gradReal is dlarray([2, 2])
%     
%     function grad = complexGrad(x)
%        y = real((2+3i)*x);
%        y = sum(y);
%        grad = dlgradient(y,x);
%     end
%     function grad = complexWithRealGradients(x)
%        y = real((2+3i)*x);
%        y = sum(y);
%        grad = dlgradient(y,x,'AllowComplex',false);
%     end
%     
%   See also: DLFEVAL, DLACCELERATE, EXTRACTDATA

%   Copyright 2018-2021 The MathWorks, Inc.

persistent tm;
if isempty(tm)
    tm = deep.internal.recording.TapeManager();
end

narginchk(2,inf);
if ~isa(y, 'dlarray')
    error(message('deep:dlgradient:InvalidDlarrayDependent'));
end
if ~isscalar(y)
    error(message('deep:dlgradient:InvalidScalarDependent'));
end
if ~internal_isRecordingArray(getRecorded(y))
    error(message('deep:dlgradient:InvalidTracedDependent'));
end
if ~isreal(y)
    error(message('deep:dlgradient:InvalidRealDependent'));
end
n = numel(varargin);
while n > 2 && (ischar(varargin{n-1}) || isstring(varargin{n-1}))
    n = n - 2;
end
x = varargin(1:n);
deep.internal.recording.validateGradientInput(x,true);
pvpairs = varargin(n+1:end);
y = getRecorded(y);

% Gather a flattened list of all of the dlarray inputs
collector = deep.internal.recording.Collection;
    function collectfun(x)
        if isa(x, 'dlarray')
            collect(collector,x);
        end
    end
deep.internal.basicContainerFun(@collectfun,x);
xc = collector.Parameters;
xcDlarray = xc;

% Extract RecordingArray objects from flattened list and perform backwards
xc = cellfun(@getRecorded,xc, 'UniformOutput', false);
[grad,isTracedGrad] = backwardPass(y,xc,pvpairs{:});

% Create dlarrays to wrap each gradient, and put those gradients back into
% the collector that we originally used to get a flattened list.
stopTracing = ~isTracedGrad;
for i=1:length(grad)
    xDlarray = xcDlarray{i};
    fd = xDlarray.FormattedData;
    val = grad{i};
    fd = insertData(fd,val);
    xDlarray.FormattedData = fd;
    if stopTracing
        xDlarray = stop(xDlarray);
    end
    grad{i} = xDlarray;
end
collector.Parameters = grad;

% Place the gradients into the original structure by iterating down the
% original and pulling gradients back out of the collector in the
% order that we originally collected them.
newgrad = deep.internal.basicContainerFun(@(~)retrieve(collector),x);

varargout = newgrad;
end

