function varargout = minibatchpredict(net, x, NameValueArgs)
%

%   Copyright 2023-2024 The MathWorks, Inc.

arguments
    net(1,1) {iValidateNetwork(net)}
end
arguments (Repeating)
    x
end
arguments
    NameValueArgs.MiniBatchSize(1,1) {mustBeInteger, mustBePositive} = 128
    NameValueArgs.Outputs { deep.internal.util.validateOutputNames(NameValueArgs.Outputs) } = net.OutputNames
    NameValueArgs.Acceleration(1,1) string { mustBeMember(NameValueArgs.Acceleration, ["none", "auto", "mex"]) } = "auto"
    NameValueArgs.ExecutionEnvironment(1,1) string { mustBeMember(NameValueArgs.ExecutionEnvironment, ["cpu", "auto", "gpu"]) } = "auto"
    NameValueArgs.SequenceLength(1,1) string { mustBeMember(NameValueArgs.SequenceLength, ["longest", "shortest"]) } = "longest"
    NameValueArgs.SequencePaddingValue(1,1) double = 0
    NameValueArgs.SequencePaddingDirection(1,1) string { mustBeMember(NameValueArgs.SequencePaddingDirection, ["right", "left"]) } = "right"
    NameValueArgs.InputDataFormats = "auto"
    NameValueArgs.OutputDataFormats = "auto"
    NameValueArgs.UniformOutput(1,1) logical = true
end

NameValueArgs = iConvertToCanonicalForm(NameValueArgs);

% Verify number of inputs
% NB: Multiple inputs only supported for in-memory data
numInputs = numel(net.InputNames);
[hasDatastoreInput, hasMinibatchqueueInput] = iHasDatastoreOrMinibatchqueueInput(x);
hasDatastoreOrMinibatchqueueInput = hasDatastoreInput || hasMinibatchqueueInput;
if ~hasDatastoreOrMinibatchqueueInput && numel(x) ~= numInputs
    error(message('nnet_cnn:dlnetwork:WrongNumInputs', numInputs, numel(x)));
elseif hasDatastoreInput && ~isscalar(x)
    error(message('deep:train:MultipleDatastoreInputs'));
elseif hasMinibatchqueueInput && ~isscalar(x)
    error(message('deep:train:MultipleMinibatchqueueInputs'));
end

% Verify number of outputs
% NB: State outputs not supported, so only support the number of
% activations requested via the Outputs argument
numOutputs = numel(NameValueArgs.Outputs);
numDataOutputs = max(nargout, 1);
if numDataOutputs > numOutputs
    error(message('nnet_cnn:minibatchpredict:TooManyOutputs', numOutputs, numDataOutputs));
end

% Validate and standardize data formats
inputDataFormats = iValidateAndStandardizeFormats(NameValueArgs.InputDataFormats, "InputDataFormats");
outputDataFormats = iValidateAndStandardizeFormats(NameValueArgs.OutputDataFormats, "OutputDataFormats");

try
    [varargout{1:numDataOutputs}] = iCallMiniBatchPredict(x, net, inputDataFormats, outputDataFormats, numInputs, numDataOutputs, hasDatastoreOrMinibatchqueueInput, NameValueArgs);
catch err
    nnet.internal.cnn.util.rethrowDLExceptions(err);
end
end

function NameValueArgs = iConvertToCanonicalForm(NameValueArgs)
    NameValueArgs.Outputs = cellstr(NameValueArgs.Outputs);
end

function iValidateNetwork(net)
if ~isa(net, 'dlnetwork') || ~net.Initialized
    error(message('nnet_cnn:minibatchpredict:InvalidNetwork'));
end
end

function [hasDatastoreInput, hasMinibatchqueueInput] = iHasDatastoreOrMinibatchqueueInput(data)
hasDatastoreInput = false;
hasMinibatchqueueInput = false;
numInputs = numel(data);
for i = 1:numInputs
    x = data{i};
    if matlab.io.datastore.internal.shim.isDatastore(x)
        hasDatastoreInput = true;
        break
    elseif isa(x, 'minibatchqueue')
        hasMinibatchqueueInput = true;
        break
    end
end
end

function formats = iValidateAndStandardizeFormats(formats, formatName)
formats = deep.internal.util.validateDataFormats(formats);

% Verify formats has a batch dim
if ~isempty(formats)
    for i = 1:numel(formats)
        if ~contains(formats{i},'B')
            error(message('nnet_cnn:minibatchpredict:DataFormatDoesNotContainBDim', formatName));
        end
    end
end
end

function varargout = iCallMiniBatchPredict(x, net, inputDataFormats, outputDataFormats, numInputs, numDataOutputs, hasDatastoreOrMinibatchqueueInput, NameValueArgs)
[mbq, metadata] = iCreateMiniBatchQueue(x, net, hasDatastoreOrMinibatchqueueInput, NameValueArgs);

% Error if gpuArray inputs are provided with CPU execution environment
if metadata.IsGpuArray && strcmp(NameValueArgs.ExecutionEnvironment, "cpu")
    error(message('nnet_cnn:minibatchpredict:ExecutionEnvironmentAndGpuArrayInputs'));
end

inputBatchDims = iGetBatchDimsFromFormats(mbq.MiniBatchFormat(1:numInputs));
predictStrategy = deep.internal.util.createMiniBatchPredictStrategy(net, inputBatchDims, numDataOutputs, mbq.OutputEnvironment, NameValueArgs);

if NameValueArgs.UniformOutput
    concatenationStrategy = deep.internal.util.NumericArrayConcatenationStrategy(metadata, inputDataFormats, outputDataFormats, inputBatchDims, NameValueArgs.Outputs);
else
    concatenationStrategy = deep.internal.util.CellArrayConcatenationStrategy(net, metadata, inputDataFormats, outputDataFormats, inputBatchDims, NameValueArgs.Outputs);
end

outputs = cell(0, numDataOutputs);
while hasdata(mbq)
    [batchInputs{1:numInputs}] = next(mbq);
    batchOutputs = predict(predictStrategy, net, batchInputs);

    [concatenationStrategy, batchOutputs] = postprocessDuringLoop(concatenationStrategy, batchInputs, batchOutputs, net);
    outputs(end+1,:) = batchOutputs; %#ok<AGROW>
end
varargout = postprocessAfterLoop(concatenationStrategy, outputs, net);
end

function [mbq, metadata] = iCreateMiniBatchQueue(data, net, hasDatastoreOrMinibatchqueueInput, NameValueArgs)
inputDataFormats = deep.internal.util.standardizeDataFormats(NameValueArgs.InputDataFormats);
if hasDatastoreOrMinibatchqueueInput
    data = data{1};
end
[mbq, ~, metadata] = deep.internal.train.createMiniBatchQueue( ...
    net, data, ...
    MiniBatchSize = NameValueArgs.MiniBatchSize, ...
    SequenceLength = NameValueArgs.SequenceLength, ...
    SequencePaddingValue = NameValueArgs.SequencePaddingValue, ...
    SequencePaddingDirection = NameValueArgs.SequencePaddingDirection,...
    InputDataFormats = inputDataFormats, ...
    OutputEnvironment = NameValueArgs.ExecutionEnvironment, ...
    IncludeTargets = false);

% Always set PreprocessingEnvironment to serial because other values cannot
% guarantee the order of outputs is preserved.
mbq.PreprocessingEnvironment = "serial";
end

function batchDims = iGetBatchDimsFromFormats(formats)
numFormats = numel(formats);
batchDims = zeros(1, numFormats);
for i = 1:numel(formats)
    sortedFormat = deep.internal.format.orderFormat(formats{i});
    batchDims(i) = find(sortedFormat=='B');
end
end