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

%   Copyright 2019-2022 The MathWorks, Inc.

% Set desired precision
precision = nnet.internal.cnn.util.Precision('single');

% Check if we have a single observation
if ~any(this.NetworkInfo.HasSequenceInput) && iSingleNumericObservation( ...
        X, this.NetworkInfo.InputSizes{1}, this.NetworkInfo.IsFeatureInput(1))
    Y = iPredictSingle( ...
        this, X, options.MiniBatchSize, executionEnvironment, precision, ...
        acceleration );
else
    Y = predictBatch( ...
        this, X, options, executionEnvironment, precision, ...
        acceleration );
end

end

function tf = iSingleNumericObservation(X, inputSize, isFeatureInput)
tf = isnumeric(X) && ~issparse(X) && ~isdlarray(X) && ...
    iIsSingleImageOrFeatureObs(X, inputSize, isFeatureInput);
end

function Y = iPredictSingle( net, X, miniBatchSize, executionEnvironment, ...
    precision, acceleration )
% Call predict on a single observation
 
% Initialize inferenceInfo to generate the optimizer
inferenceInfo = nnet.internal.cnn.InferenceInfo( "predict", net.NetworkInfo.InputSizes, ...
    miniBatchSize, 1, executionEnvironment, precision, [], ...
    acceleration, [] );
useMex = inferenceInfo.UseMex;

% Get the prediction network from the inference options
[predictNetwork, gpuShouldBeUsed] = setupPredictionEnvironment(net, inferenceInfo);

X = precision.cast(X);

if net.NetworkInfo.IsFeatureInput(1)
    X = X';
end

% MEX accept a gpuArray for a single input
mustGatherInputsForMex = false;

% Move the input in the selected environment and get the predictions
[X, mustReturnGpuArray] = moveInferenceInputToEnvironment(X, gpuShouldBeUsed, mustGatherInputsForMex);
Y = predictNetwork.predict( X );

% Gather all the outputs if necessary
numOutputs = net.PrivateNetwork.NumOutputs;
if ~mustReturnGpuArray
    Y = cellfun(@gather, Y, 'UniformOutput', false);
end

% MEX does not return gpuArray in the case of multiple outputs
if mustReturnGpuArray && useMex
    Y = cellfun(@gpuArray, Y, 'UniformOutput', false);
end

% Reshape each output as a 2D array. MEX returns an output already reshaped
if ~useMex
    outputSizes = net.NetworkInfo.OutputSizes;
    outputObservationDims = cellfun( @(s)numel(s)+1, outputSizes );
    for i = 1:numOutputs
        Y{i} = iFormatPredictionsAs2DRowResponses( ...
            Y{i}, outputObservationDims(i) );
    end
end
end

function YFormatted = iFormatPredictionsAs2DRowResponses(Y, observationDim)
% iFormatPredictionsAs2DRowResponses   Format predictions according to the
% problem.
% If Y is [1 1 K N]/[1 1 1 K N], then YFormatted will be [N K].
% If Y is [H W (D) K N], with H, W or D not singleton, then YFormatted will
% be the same as Y.
YSize = ones(1, observationDim);
YSize(1:ndims(Y)) = size(Y);
if all(YSize(1:observationDim-2)==1)
    % Get rid of first two/three/ObsDim-2 singleton dimensions for 2D/3D/ND
    YFormatted = shiftdim(Y,observationDim-2);
    % Transpose [K N] -> [N K]
    YFormatted = YFormatted';
else
    YFormatted = Y;
end
end

function tf = iIsSingleImageOrFeatureObs(X, inputSize, isFeatureInput)
if isFeatureInput
    tf = isequal(size(X), [1 inputSize]);
else
    tf = isequal(size(X, 1:(numel(inputSize)+1)), [inputSize 1]);
end
end
