function Y = predictBatch( this, X, options, executionEnvironment, precision, acceleration )
%

%   Copyright 2019-2022 The MathWorks, Inc.

% Validate data and network
nnet.internal.cnn.data.validatePredictionData(X,this.PrivateNetwork.InputLayers);

% Create dispatcher
dispatcher = this.createPredictionDataDispatcher(X, options, precision, this.NetworkInfo);

% Number of observations in dispatcher
numObservations = dispatcher.NumObservations;

% Initialize inferenceInfo to generate the optimizer
inferenceInfo = nnet.internal.cnn.InferenceInfo( "predict", this.NetworkInfo.InputSizes, ...
    options.MiniBatchSize, numObservations, executionEnvironment, precision, [], ...
    acceleration, options.SequenceLength );

% There are three scenarios in which a network needs stateful predict:
%   1) The network has at least one recurrent layer and the
%   'SequenceLength' option has been specified as an integer
%   2) The network has at least one custom stateful layer
needsStatefulPredict = (this.NetworkInfo.HasRecurrentLayer && inferenceInfo.RecurrentInferenceRequiresState) || ...
    this.NetworkInfo.HasStatefulCustomLayer;

if ~needsStatefulPredict
    % Validate the data size for the network
    if this.PrivateNetwork.NumInputLayers == 1
        iValidateInputSize(dispatcher.DataSize, this.NetworkInfo.InputSizes{1});
    else
        iValidateMultipleInputSizes(dispatcher.DataSize, this.NetworkInfo.InputSizes);
    end
end

% Get the prediction network from the inference options
[predictNetwork, gpuShouldBeUsed, workerLoad, useParallel] = setupPredictionEnvironment(this, inferenceInfo);
useMex = inferenceInfo.UseMex;

if needsStatefulPredict
    % Redirect the predict call to predictStateful and return its output
    Y = this.predictStateful(X, dispatcher, precision, predictNetwork, gpuShouldBeUsed);
    return
end

% Cache the information about the inputs and outputs
networkInfo = this.NetworkInfo;
outputSizes = networkInfo.OutputSizes;
outputLayers = predictNetwork.OutputLayers;
numOutputs = this.PrivateNetwork.NumOutputs;
outputObservationDims = cellfun( @(s)numel(s)+1, outputSizes );

% For networks with stateful layers, we check that the initial state is 
% of the expected size
statefulLayers = predictNetwork.StatefulLayers;
if ~isempty(statefulLayers)
    nnet.internal.cnn.util.assertInitialStateIsValidForPredict(statefulLayers, dispatcher.MiniBatchSize)
end

% The output can be allocated if:
%  - the number of observations is known
%  - no layer has output size dependent on the values of input data
%  - the network does not have any sequence outputs
outputIsAllocated = isfinite(numObservations) && ...
    ~any(networkInfo.HasSequenceOutput) && ...
    ~networkInfo.HasLayerWithSizePropagationValueSink;

if useMex
    % MEX may return the output already as 2D array.
    [outputSizes, outputObservationDims] = cellfun( ...
        @(s, obsDim) iInit2DRowResponsesFormat(s, numObservations, obsDim), ...
        outputSizes, num2cell(outputObservationDims), 'UniformOutput', false);
    outputObservationDims = [outputObservationDims{:}];
end  

if useParallel
    numOfWorkers = nnz(workerLoad);
    Y = cell(numOfWorkers, numOutputs);

    % Distribute data to workers
    distributedData = distributeForInference(dispatcher, numOfWorkers);

    parfor idx = 1:numOfWorkers
        subDispatcher = distributedData{idx};
        subNumObservations = subDispatcher.NumObservations;
        
        % Use the dispatcher to run the network on the data
        Ysub = iDoInference(subDispatcher, subNumObservations, gpuShouldBeUsed, ...
            predictNetwork, numOutputs, useMex, outputIsAllocated, ...
            outputObservationDims, outputSizes, precision, networkInfo);
        
        for i = 1:numOutputs
            Y{idx,i} = Ysub{i};
        end
    end        
    Y = iConcatenateCellOutput(Y, outputObservationDims, networkInfo.HasSequenceOutput);                
else
    Y = iDoInference(dispatcher, numObservations, gpuShouldBeUsed, ...
        predictNetwork, numOutputs, useMex, outputIsAllocated, ...
        outputObservationDims, outputSizes, precision, networkInfo);
end

% Post-process the output array to return as a 2D array.
if ~useMex
    for i = 1:numOutputs
        Y{i} = iFormatPredictionsAs2DRowResponses(X, Y{i}, outputSizes{i}, ...
            networkInfo.OutputFormats{i}, outputLayers{i}, ...
            networkInfo.HasSequenceInput, networkInfo.HasSequenceOutput(i));
    end
end
end


function Y = iDoInference(dispatcher, numObservations, gpuShouldBeUsed, ...
    predictNetwork, numOutputs ,useMex, outputIsAllocated, ...
    outputObservationDims, outputSizes, precision, networkInfo)

% Create the output
if outputIsAllocated
    % Find the output size.
    if ~useMex
        outputSizes = cellfun(@(s)[s numObservations], outputSizes, 'UniformOutput', false);
    end
    % Initialize as a cell array, one cell for each output
    Y = cellfun(@(s)precision.zeros(s), outputSizes, 'UniformOutput', false);
else
    % Output size is unknown.
    batchIndex = 0;
    doSeqDimCat = false;
    Y = cell(batchIndex, numOutputs);
end

% Create a flag to indicate whether the allocated output should be
% moved on the GPU during the first batch
allocatedOutputIsInExpectedEnvironment = false;
dataIsOnGpu = false;

% If the network has sequence output, get the number of dimensions of the 
% padding mask
if any(predictNetwork.HasSequenceOutput)
    ndimsMask = nnet.internal.cnn.util.getNDimsMask(networkInfo.InputFormats);
end

% Use the dispatcher to iterate over the data
dispatcher.start();
while ~dispatcher.IsDone
    % Collect the next mini-batch
    [X, ~, info] = dispatcher.next();
    [X, mustReturnGpuArray] = iPrepDataForNetInput(X, gpuShouldBeUsed, useMex);
    
    % Get the mini-batch output
    YBatch = predictNetwork.predict(X);
        
    % Move the allocated output on the GPU if necessary
    if mustReturnGpuArray && ~allocatedOutputIsInExpectedEnvironment && outputIsAllocated
        Y = cellfun( @gpuArray, Y, 'UniformOutput', false );
        dataIsOnGpu = true;
    elseif mustReturnGpuArray && ~dataIsOnGpu && outputIsAllocated
        Y = cellfun( @gpuArray, Y, 'UniformOutput', false );
        dataIsOnGpu = true;
    end
    allocatedOutputIsInExpectedEnvironment = true;
    
    % Insert the mini-batch output in the output array
    for i = 1:numOutputs
        currentY = YBatch{i};
        
        % Gather the current output if needed
        if ~mustReturnGpuArray
            currentY = gather(currentY);
        end
        
        if outputIsAllocated
            % Index the output array. The following code is equivalent to 
            % this for 4-D input case: Y(:,:,:,indices) = gather(YBatch{1});
            Y{i} = nnet.internal.cnn.util.aggregateArrayFromObservations(Y{i},...
                info.MiniBatchIdx, currentY, outputObservationDims(i));
        else
            if useMex && mustReturnGpuArray
                % MEX does not return the output on the GPU if there
                % were multiple inputs. Ensure the result on the device
                currentY = gpuArray(currentY);
            end

            if networkInfo.HasSequenceOutput(i)
                % Split each observation into its own cell - currentY becomes
                % a batchSize x 1 cell array
                currentY = nnet.internal.cnn.util.internalToExternalSequenceData(...
                    currentY, outputSizes{i});
                
                % Remove sequence padding
                currentY = nnet.internal.cnn.util.removeSequencePadding(...
                    currentY, info.Mask, outputSizes{i}, ndimsMask);

                % Write current batch into output cell array
                if doSeqDimCat
                    % Concatenate currentY with existing observations along
                    % sequence dimension
                    Y{batchIndex, i} = cellfun(@(a,b)cat(outputObservationDims(i),a,b), Y{batchIndex,i}, currentY, UniformOutput=false);
                else
                    % Add currentY into the output Y at a new batch index
                    batchIndex = batchIndex + 1;
                    Y{batchIndex, i} = currentY;
                end
                doSeqDimCat = ~info.IsCompleteObservation;
            else
                % Index the output cell array to add an entry
                batchIndex = batchIndex + 1;
                Y{batchIndex, i} = currentY;
            end
        end
    end
end

if ~outputIsAllocated
    % Concatenate the output cells together
    Y = iConcatenateCellOutput(Y, outputObservationDims, networkInfo.HasSequenceOutput);
end
end

function iValidateInputSize(dataSize, inputSize)

if iscell(dataSize) && numel(dataSize) == 1
    dataSize = dataSize{1};
end

if isequal(dataSize, inputSize)
else
    if isscalar(dataSize)
        exception = iCreateExceptionFromErrorID( ...
            'nnet_cnn:DAGNetwork:WrongSizePredictDataForFeatureInputLayer', ...
            mat2str(inputSize));
        throwAsCaller(exception);
    else
        exception = iCreateExceptionFromErrorID( ...
            'nnet_cnn:DAGNetwork:WrongSizePredictDataForImageInputLayer', ...
            mat2str(inputSize));
        throwAsCaller(exception);
    end
end
end

function iValidateMultipleInputSizes(dispatcherInputSizes, networkInputSizes)
nnet.internal.cnn.util.validateMultipleInputSizes(dispatcherInputSizes, networkInputSizes);
end

function exception = iCreateExceptionFromErrorID(errorID, varargin)
exception = MException(message(errorID, varargin{:}));
end

function YBatched = iConcatenateCellOutput(Y, batchDims, hasSequenceOutput)
numOutputs = numel(batchDims);
YBatched = cell(1, numOutputs);

for i = 1:numOutputs
    if hasSequenceOutput(i)
        % For sequence output, we have a cell array where each cell is a
        % batchSize x 1 cell array, so the dimension we want to concatenate
        % over is the first one.
        catDim = 1;
    else
        catDim = batchDims(i);
    end
    YBatched{i} = cat(catDim, Y{:,i});
end
end

function [newOutputSize, formattedObservationDim] = iInit2DRowResponsesFormat(outputSize, numObservations, observationDim)
if all(outputSize(1:observationDim-2) == 1)
    % New output size will be [N K]
    newOutputSize = [numObservations , outputSize(observationDim-1)];
    formattedObservationDim = 1; % Observation dim becomes the first dimension
else
    % New output size will be [H W (D) K N]
    newOutputSize = [outputSize, numObservations];
    formattedObservationDim = observationDim; % Unchanged
end
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 Y = iFormatPredictionsAs2DRowResponses(X, Y, outputSize, outputFormat, ...
    outputLayer, hasSequenceInput, hasSequenceOutput)
% iFormatPredictionsAs2DRowResponses   Format predictions according to the
% problem.
% If Y is [1 1 C B]/[1 1 1 C B], then YFormatted will be [B C].
% If Y is [S S (S) C B], with all S dims non-singleton, then YFormatted will
% be the same as Y.
% For the classification sequence output case, we remove singleton 
% dimension but we don't transpose the contents of the cell arrays.

% Check if the output has all singleton spatial dimensions
spatialOutputSizes = getSizeForDims(outputFormat, 'S');
hasSingletonSpatialDims = all( spatialOutputSizes ==  1 );

if hasSequenceOutput
    % For networks with sequence output, we only remove the singleton
    % dimensions in the classification case
    if hasSingletonSpatialDims && iIsClassificationOutputLayer(outputLayer)
        numSpatialDims = numel(spatialOutputSizes);
        Y = cellfun(@(y)shiftdim(y, numSpatialDims), Y, 'UniformOutput', false);
    end
    
    % Return numeric output if we have a numeric array (single sequence) at
    % input
    if any(hasSequenceInput) && iNeedsNumericOutput(X)
        Y = Y{1};
    end
else
    % Return outputs in shape (N)x(outputSize) if outputSize is
    % scalar or the output size has singleton spatial dimensions
    if hasSingletonSpatialDims
        numSpatialDims = numel(spatialOutputSizes);
        Y = shiftdim(Y, numSpatialDims);
        Y = Y';
    elseif isscalar(outputSize)
        Y = Y';
    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 tf = iIsClassificationOutputLayer(internalLayer)
tf = isa(internalLayer, 'nnet.internal.cnn.layer.ClassificationLayer');
end
