function varargout = internal_dlfor(bodyInfo,varargin)
% This function is undocumented and will change in a future release.

% Applies a for loop operation. All inputs and outputs are numeric.
%
% [z1,...,zN,intermediatesCell] = internal_dlfor(bodyInfo,range1,...,rangeM,x1,...,xN)
%
% Inputs:
%   bodyInfo          = A struct containing information about the body
%                       function.
%   range1,...,rangeM = Vectors specifying the ranges of the for loop.
%                       Length of range i is Ri.
%   x1,...,xN         = Inputs to the first iteration of the for loop.
%
% Outputs:
%   z1,...,zN         = Outputs from the final iteration of the for loop.
%   intermediatesCell = Variables needed for backward calculations. This is
%                       a cell array of length prodR = prod(Ri over i)
%                       where element k of the cell array holds variables
%                       needed for backward calculations for iteration k of
%                       the nested for loop.
%
% Suppose bodyForwardFcn is the loop body forward function. Then
% bodyForwardFcn should be callable like this:
%
%    [x1,...,xN,y1,...,yL] = bodyForwardFcn(i1,...,iM,x1,...,xN)
%
% where i1,...,iM are the loop indices for the M ranges, N is the number of
% loop body outputs and L is the number of intermediates produced by the
% loop body forward function.
%
% Let [x1_k,...,xN_k] be the loop body outputs for the kth iteration and
% let [y1_k,...,yL_k] be the loop body intermediates for the kth iteration.
% Let:
%
%   u_k = {i1,...,im,x1_(k-1),...,xN_(k-1),x1_k,...,xN_k,y1_k,...,yL_k};
%   usubset_k = u_k(bodyMask_k);
%
% Here u_k are the loop body inputs, outputs and intermediates for
% iteration k and bodyMask_k is the loop body mask that specifies the
% subset usubset_k needed for backward of iteration k. intermediatesCell is
% a cell array defined like this:
%
%   intermediatesCell = [usubset_1,...,usubset_prodR];

%   Copyright 2022 The MathWorks, Inc.

% 1. Number of ranges, outputs and intermediates from the loop body.
numRanges = bodyInfo.RangeInfo.NumRanges;
numBodyOutputs = bodyInfo.NumBodyOutputs;
numBodyIntermediates = bodyInfo.NumBodyIntermediates;

% 2. Subset of inputs/outputs to the loop body needed for backward.
bodyMask = bodyInfo.BodyAccOp.Memory;
[loopIndexMask,XMask,ZMask,YMask] = iSplitMask(bodyMask,numBodyOutputs,numRanges);

% 3. Forward function of the loop body.
bodyForwardFcnWithConstantsInput = bodyInfo.BodyForwardFcnWithConstantsInput;
forwardConstants = bodyInfo.ForwardConstants;

% 4. Loop indices of the flattened loop.
ranges = varargin(1:numRanges);
[subsMat,prodRangeSizes] = deep.internal.recording.rangesToIndexMatrix(ranges);

% 5. Inputs to the loop.
X = varargin(numRanges+1:end);
varargin = []; %#ok<NASGU>

% 6. Forward through the loop body and save intermediates for backward.
Y = cell(1,numBodyIntermediates);
intermediatesCell = cell(1,prodRangeSizes);
loopIndex = cell(1,numRanges);
for k = 1:prodRangeSizes
    for r = 1:numRanges
        loopIndex{r} = subsMat(k,r);
    end
    savedLoopIndex = loopIndex(loopIndexMask);
    savedX = X(XMask);
    [X{:},Y{:}] = bodyForwardFcnWithConstantsInput(loopIndex{:},X{:},forwardConstants);
    savedZ = X(ZMask);
    savedY = Y(YMask);
    intermediatesCell{k} = [matlab.lang.internal.move(savedLoopIndex),matlab.lang.internal.move(savedX),matlab.lang.internal.move(savedZ),matlab.lang.internal.move(savedY)];
end

% 7. Return op outputs and intermediates for backward.
varargout = [X,{intermediatesCell}];
end

function [loopIndexMask,XMask,ZMask,YMask] = iSplitMask(bodyMask,numBodyOutputs,numRanges)
% Signature of the loop body forward function looks like this:
% [x1,...,xN,y1,...,yL] = bodyForwardFcn(i1,...,iM,x1,...,xN)
loopIndexMask = bodyMask(1:numRanges);
count = numRanges;
XMask = bodyMask(count+1:count+numBodyOutputs);
count = count + numBodyOutputs;
ZMask = bodyMask(count+1:count+numBodyOutputs);
count = count + numBodyOutputs;
YMask = bodyMask(count+1:end);
end