function [Y, state, predictNetwork] = predictStateful(this, X, ...
    dispatcher, ~, predictNetwork, gpuShouldBeUsed)
%

%   Copyright 2020-2022 The MathWorks, Inc.

% Check that the initial state is of the expected size
nnet.internal.cnn.util.assertInitialStateIsValidForPredict(predictNetwork.StatefulLayers, ...
    dispatcher.MiniBatchSize);

% Inference
outputSizes = this.NetworkInfo.OutputSizes;
observationDims = cellfun(@(x)numel(x)+1, outputSizes);
numOutputs = predictNetwork.NumOutputs;
if any(predictNetwork.HasSequenceOutput)
    ndimsMask = nnet.internal.cnn.util.getNDimsMask(this.NetworkInfo.InputFormats);
end
Y = cell(1, numOutputs);
Y(:) = {{}};
dispatcher.start();
while ~dispatcher.IsDone
    [Xb, ~, info] = dispatcher.next();    
    [Xb, mustReturnGpuArray] = iPrepDataForNetInput(Xb, gpuShouldBeUsed, false);
    
    [Yb, state] = predictNetwork.statefulPredict(Xb);
    predictNetwork = predictNetwork.updateNetworkState(state);
    if info.IsCompleteObservation
        predictNetwork = predictNetwork.resetNetworkState();
    end
    
    % Allocate prediction into output
    if ~mustReturnGpuArray
        Yb = cellfun(@(x)gather(x), Yb, 'UniformOutput', false);
    end
    for i = 1:numOutputs
        Ycell = nnet.internal.cnn.util.internalToExternalSequenceData(...
            Yb{i}, outputSizes{i});
        Y{i} = [Y{i}; cell( max(info.MiniBatchIdx) - numel(Y{i}), 1 )];
        iValidateBatchSizes(Y{i}(info.MiniBatchIdx), Ycell);
        if predictNetwork.HasSequenceOutput(i)
            % (seq2seq) Concatenate outputs along sequence dimension
            Ycell = nnet.internal.cnn.util.removeSequencePadding(...
                Ycell, info.Mask, outputSizes{i}, ndimsMask);
            Y{i}(info.MiniBatchIdx) = cellfun( @(a,b)cat(observationDims(i), a, b), ...
                Y{i}(info.MiniBatchIdx), Ycell, 'UniformOutput', false);
        else
            % (seq2one) Concatenate outputs along observation dimension
            Y{i}(info.MiniBatchIdx) = Ycell;
        end
        
    end
end

for i = 1:numOutputs
    % Remove singleton spatial dimensions
    spatialOutputSizes = getSizeForDims(this.NetworkInfo.OutputFormats{i}, 'S');
    hasSingletonSpatialDims = all( spatialOutputSizes ==  1 );
    if this.NetworkInfo.HasSequenceOutput(i)
        if hasSingletonSpatialDims && iIsClassificationOutputLayer(predictNetwork.OutputLayers{i})
            numSpatialDims = numel(spatialOutputSizes);
            for n = 1:numel(Y{i})
                Y{i}{n} = shiftdim(Y{i}{n}, numSpatialDims);
            end
        end
        if any(this.NetworkInfo.HasSequenceInput) && iNeedsNumericOutput(X)
            % Return numeric output if we have a numeric input
            Y{i} = Y{i}{1};
        end

    else
        % Concatenate observations
        Y{i} = cat(observationDims(i), Y{i}{:});

        % Return outputs in shape (N)x(outputSize) if outputSize is
        % scalar or the output size has singleton spatial dimensions
        if isscalar(outputSizes{i})       
            Y{i} = Y{i}';
        elseif hasSingletonSpatialDims
            numSpatialDims = numel(spatialOutputSizes);
            Y{i} = shiftdim(Y{i}, numSpatialDims);
            Y{i} = Y{i}';
        end
    end   
end
end

function tf = iNeedsNumericOutput(X)
% If the input X is a single observation passed as a numeric array, we need
% to un-wrap cell seq2seq outputs and return a numeric array
tf = isnumeric(X);
end

function [X, mustReturnGpuArray] = iPrepDataForNetInput(X, gpuShouldBeUsed, useMex)
X = iWrapInCell(X);

mustGatherInputsForMex = useMex && (numel(X) > 1);

[X, inputIsGpuArray] = cellfun(@(x)moveInferenceInputToEnvironment(x, gpuShouldBeUsed, mustGatherInputsForMex), X, 'UniformOutput', false);
mustReturnGpuArray = any([inputIsGpuArray{:}]);
end

function X = iWrapInCell(X)
if ~iscell(X)
    X = {X};
end
end

function iValidateBatchSizes(Yold, Ynew)
if numel(Yold) ~= numel(Ynew)
    error(message('nnet_cnn:DAGNetwork:IncorrectRNNBatchSize', numel(Yold), numel(Ynew)));
end
end

function tf = iIsClassificationOutputLayer(internalLayer)
tf = isa(internalLayer, 'nnet.internal.cnn.layer.ClassificationLayer');
end
