function [dX,dRHS] = internal_parenAssignBackward(dZ,Flags,ShapeX,ShapeRHS,varargin)
% This function is undocumented and will change in a future release.

% Backward method for ParenAssignOp.
%
% Z = parenAssign of X using subscripts in Index to RHS.
%
% Inputs:
%   dZ       = gradient of loss w.r.t Z
%   Flags    = logical vector indicating which outputs to compute
%   ShapeX   = size of X
%   ShapeRHS = size of RHS
%   Index    = subscripts used to index into X
%
% Outputs:
%   dX   = gradient of loss w.r.t X
%   dRHS = gradient of loss w.r.t RHS

%   Copyright 2019 The MathWorks, Inc.

Index = varargin;
Index = iRemoveObjectIndex(Index);

dXNeeded = Flags(1);
dRHSNeeded = Flags(2);

dX = [];
dRHS = [];

if dXNeeded && dRHSNeeded
    dRHS = iComputeRHSGradient(dZ,Index,ShapeRHS);
    dZ = iComputeXGradient(dZ,Index,ShapeX);
    dX = dZ;
elseif dXNeeded && ~dRHSNeeded
    dZ = iComputeXGradient(dZ,Index,ShapeX);
    dX = dZ;
elseif ~dXNeeded && dRHSNeeded
    dRHS = iComputeRHSGradient(dZ,Index,ShapeRHS);
end
end

function dZ = iComputeXGradient(dZ,Index,ShapeX)
SizedZ = size(dZ);
dZ(Index{:}) = 0;
if ~isequal(SizedZ, ShapeX)
    % Cut dX back to its original size
    % dZ has been grown compared to original size of X
    shapeXcell = cell(1, length(ShapeX));
    for ii=1:length(ShapeX)
        shapeXcell{ii} = 1:ShapeX(ii);
    end
    dZ = dZ(shapeXcell{:});
end
dZ = reshape(dZ, ShapeX); % Just for safety
end

function dRHS = iComputeRHSGradient(dZ,Index,ShapeRHS)
% Uniquify the indices: a([1 1 2]) = b is equivalent to
% a([1 2]) = b([2 3]) - only the last unique instance of each
% index is used. Must have dRHS(1) == 0.
uindLHS = Index;
uindRHS = repmat({':'}, size(Index));
needUniqueInd = false;
for ii=1:length(Index)
    indi = Index{ii};
    if (ischar(indi) && isequal(indi, ':')) || islogical(indi)
        % No multiple indices possible.
        continue;
    end
    % Get only the last unique instance of each index
    % Example indi = [1 1 2]: uindLHS = [1 2], uindRHS = [2 3].
    %[uindLHS{ii}, uindRHS{ii}] = unique(indi, 'last');
    
    % For performance: Use sparse indexing to get same result
    % as the expression using UNIQUE above.
    sparseIndexVector = sparse([]);
    sparseIndexVector(indi, 1) = 1:numel(indi); %#ok<SPRIX>
    
    if nnz(sparseIndexVector) < numel(indi)
        [~, uindLHS{ii}, uindRHS{ii}] = find(sparseIndexVector);
        needUniqueInd = true;
    end
end

if any(ShapeRHS ~= 1)
    % Example: dRHS([2 3]) = dZ([1 2]). This returns dRHS of the
    % correct size, because the last index in uindRHS is always
    % equal to the length(indi).
    if needUniqueInd
        dRHS(uindRHS{:}) = dZ(uindLHS{:});
    else
        dRHS = dZ(uindLHS{:});
    end
    dRHS = reshape(dRHS, ShapeRHS);
else
    % Note: Only scalar expansion supported for subsassign, no
    % general implicit expansion.
    dRHS = sum(dZ(uindLHS{:}), 'all');
end
end

function index = iRemoveObjectIndex(index)
for ii=1:length(index)
    if isobject(index{ii})
        index{ii} = subsindex(index{ii})+1;
    end
end
end
