function varargout = predictAndUpdateState(this, varargin)
% predictAndUpdateState   Make predictions on data with a network and update the network state
%
%   [net, Y] = predictAndUpdateState(net, X) computes predictions of the
%   network net on the data X, and updates the state of the network net.
%
%   This method only supports networks with sequence input layers.
%
%   For a network with a sequence input layer, X must be:
%       - A numeric array representing a single sequence.
%       - A cell array of sequences.
%       - A datastore, where the output of the datastore read function is
%         a cell array or a table containing sequences in the first
%         column.
%   Each sequence must be either:
%       - A vector sequence, specified by a C-by-S matrix, where C is the
%         number of features and S is the number of time steps.
%       - A 2-D image sequence, specified by an H-by-W-by-C-by-S array,
%         where H-by-W-by-C is the 2-D image size and S is the number of
%         time steps.
%       - A 3-D image sequence, specified by an H-by-W-by-D-by-C-by-S
%         array, where H-by-W-by-D-by-C is the 3-D image size and S is the
%         number of time steps.
%
%   For a classification problem, Y will contain the predicted scores,
%   arranged in an N-by-K matrix, where N is the number of observations,
%   and K is the number of classes. If the output of the network is a
%   sequence, scores will be an N-by-1 cell array and each element of
%   scores will be a K-by-S numeric array, where S is the sequence length
%   of the output. If the number of observations in X is one, scores will
%   be a K-by-S numeric array.
%
%   [net, Y1, ..., YM] = predictAndUpdateState(net, X) computes predictions
%   with the data X, where M is the number of network outputs. The output
%   layers determine the format of the output data as described above.
%
%   __ = predictAndUpdateState(net, cds) computes predictions for a
%   multiple input network with the combined datastore cds. The combined
%   datastore must return a 1-by-N cell array, where N is the number of
%   network inputs.
%
%   __ = predictAndUpdateState(net, X1, ..., XN) computes predictions for a
%   multiple input network with N inputs, where the input data is X1, ...,
%   XN.
%
%   __ = predictAndUpdateState(__, 'PARAM1', VAL1, ...) will compute
%   predictions with the following optional name/value pairs:
%
%       'MiniBatchSize'         - The size of the mini-batches for
%                                 computing predictions. Larger mini-batch
%                                 sizes lead to faster predictions, at the
%                                 cost of more memory. The default is 128.
%       'ExecutionEnvironment'  - The execution environment for the
%                                 network. This determines what hardware
%                                 resources will be used to run the
%                                 network.
%                                   - 'auto' - Use a GPU if it is
%                                     available, otherwise use the CPU.
%                                   - 'gpu' - Use the GPU. To use a
%                                     GPU, you must have Parallel Computing
%                                     Toolbox(TM) and a supported GPU
%                                     device. If a suitable GPU is not
%                                     available, predictAndUpdateState
%                                     returns an error message.
%                                   - 'cpu' - Use the CPU.
%                                 The default is 'auto'.
%       'Acceleration'          - Optimizations that can improve
%                                 performance at the expense of some
%                                 overhead on the first call and possible
%                                 additional memory usage.
%                                   - 'auto' - Automatically select 
%                                     optimizations suitable for the input 
%                                     network and environment. 
%                                   - 'none' - Disable all acceleration.
%                                 The default is 'auto'.
%       'ReturnCategorical'     - Option to return outputs of
%                                 classification layers as categorical,
%                                 specified as true or false. The default
%                                 is false.
%       'SequenceLength'        - Strategy to determine the length of the
%                                 sequences used per mini-batch. Options
%                                 are:
%                                   - 'longest' to pad all sequences in a
%                                     batch to the length of the longest
%                                     sequence.
%                                   - 'shortest' to truncate all sequences
%                                     in a batch to the length of the
%                                     shortest sequence.
%                                   - Positive integer - For each
%                                     mini-batch, pad the sequences to the
%                                     nearest multiple of the specified
%                                     length that is greater than the
%                                     longest sequence length in the
%                                     mini-batch, and then split the
%                                     sequences into smaller sequences of
%                                     the specified length. If splitting
%                                     occurs, then the software creates
%                                     extra mini-batches.
%                                 The default is 'longest'.
%       'SequencePaddingValue'  - Scalar value used to pad sequences where 
%                                 necessary. The default is 0.
%       'SequencePaddingDirection' 
%                               - Direction of padding or truncation,
%                                 specified as:
%                                  - 'right' - Pad or truncate sequences
%                                  on the right.
%                                  - 'left' - Pad or truncate sequences on
%                                  the left.
%                                 The default is 'right'.
%
%   See also DAGNetwork/classifyAndUpdateState.

%   Copyright 2018-2022 The MathWorks, Inc.

if ~isscalar(this)
    error(message('nnet_cnn:DAGNetwork:NonScalarNetwork'));
end

% This method is only available for networks with a sequence input layer
if ~any(this.NetworkInfo.HasSequenceInput)
    error(message('nnet_cnn:DAGNetwork:InvalidNetworkForPredictAndUpdateState'));
end

if this.NumInputLayers == 1
    % Network has one input
    X = varargin{1};
    nameValuePairs = varargin(2:end);
elseif iDataForMultipleInputNetworkPassedAsOneArgument(nargin-1, varargin) && ...
        iIsCombinedOrTransformedDatastore(varargin)
    % Network has multiple inputs, but the inputs are passed in as one
    % argument.
    X = varargin{1};
    nameValuePairs = varargin(2:end);
elseif iDataForMultipleInputNetworkPassedAsMultipleArguments(nargin-1, varargin, this.NumInputLayers)
    % Network has multiple inputs, and the inputs are passed in as multiple
    % arguments.
    X = varargin(1:this.NumInputLayers);
    nameValuePairs = varargin((this.NumInputLayers+1):end);
else
    error( message( 'nnet_cnn:DAGNetwork:InvalidPredictDataForMultipleInputNetwork', this.NumInputLayers) );
end

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

statefulPredict = true;
[options, executionEnvironment, acceleration, returnCategorical] = ...
    this.parseAndValidatePredictNameValuePairs(statefulPredict, nameValuePairs{:});

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

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

% Cache variables that are slow to recalculate to optimize
% calling of predict in a tight loop
inferenceInfo = nnet.internal.cnn.InferenceInfo( ...
    "predictAndUpdateState", this.NetworkInfo.InputSizes, options.MiniBatchSize, ...
    dispatcher.NumObservations, executionEnvironment, precision, [], ...
    acceleration, options.SequenceLength );

[predictNetwork, gpuShouldBeUsed] = setupPredictionEnvironment(this, inferenceInfo);

% Call predictStateful method
try
[Y, finalState, predictNetwork] = this.predictStateful(X, dispatcher, ...
    precision, predictNetwork, gpuShouldBeUsed);
catch ex
    nnet.internal.cnn.util.rethrowDLExceptions(ex)
end

% Update the network to the final state
finalState = iGatherState(finalState);
predictNetwork = predictNetwork.updateNetworkState(finalState);
this.PrivateNetwork = predictNetwork;

% Empty the layers cache
this.LayersCache = nnet.internal.cnn.layer.learnable.CacheHandle();

% Create a copy of the PredictEnvironment cache handle
predictEnvironment = this.PredictEnvironment.Value;
predictEnvironment.PredictNetwork = predictNetwork;
newPredictEnvironmentHandle = nnet.internal.cnn.layer.learnable.CacheHandle();
newPredictEnvironmentHandle.fillCache( predictEnvironment );
this.PredictEnvironment = newPredictEnvironmentHandle;

% Convert outputs to categorical if necessary
numOutputs = predictNetwork.NumOutputs;
varargout = cell(1,numOutputs+1);
if returnCategorical
    Y = this.convertMultiOutputToCategorical(Y);
end
varargout(1) = {this};
varargout(2:numOutputs+1) = Y;

end

function state = iGatherState(state)
for ii = 1:numel(state)
    if ~isempty(state{ii})
        state{ii} = cellfun(@gather, state{ii}, 'UniformOutput', false);
    end
end
end

function tf = iDataForMultipleInputNetworkPassedAsOneArgument( ...
    numArguments, arguments )
tf = (numArguments == 1) || ...
    iNameValuePairsShouldStartAtThisPosition(numArguments, arguments, 2);
end

function tf = iIsCombinedOrTransformedDatastore(arguments)
tf = isa(arguments{1},'matlab.io.datastore.CombinedDatastore') || ...
    isa(arguments{1},'matlab.io.datastore.TransformedDatastore');
end

function tf = iDataForMultipleInputNetworkPassedAsMultipleArguments( ...
    numArguments, arguments, numInputs )
tf = (numArguments == numInputs) || ...
    iNameValuePairsShouldStartAtThisPosition(numArguments, arguments, numInputs+1);
end

function tf = iNameValuePairsShouldStartAtThisPosition(numArguments, arguments, position)
% Note that this check takes place before we parse arguments, so it does
% not definitively prove that name/value pairs start at this position (i.e.
% the test is necessary but not sufficient). However, in the event that
% there is a string in this position, but it is not an N-V pair, the parser
% will find the error.
tf = (numArguments >= position) && ...
    iFirstStringOrArrayIsAtThisPosition(arguments, position);
end

function tf = iFirstStringOrArrayIsAtThisPosition(inputCell, position)

firstStringOrCharArray = find( ...
    cellfun(@iIsValidStringOrCharArray, inputCell), 1 );

if isempty(firstStringOrCharArray)
    tf = false;
else
    tf = ( position == firstStringOrCharArray );
end

end

function tf = iIsValidStringOrCharArray(value)
tf = nnet.internal.cnn.layer.paramvalidation.isValidStringOrCharArray(value);
end