function Y = calculateActivations(this, X, layerIndex, layerOutputIndex, varargin)
% calculateActivations  Shared internal core of the different activations
% methods

%   Copyright 2018-2022 The MathWorks, Inc.

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

% Validate optional inputs
[options, outputAs, executionEnvironment, ~, acceleration] = ...
    this.parseAndValidateActivationsNameValuePairs(varargin{:});

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

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

% Validate that input size from dispatcher is suitable for computing
% activations
iValidateDispatcherForActivations(dispatcherInputSize, ...
    this.PrivateNetwork.InputSizes, this.InputNames, outputAs);

% Get the prediction network from the inference options
layerIndexInUnsortedLayers = this.PrivateNetwork.TopologicalOrder(layerIndex);

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

inferenceInfo = nnet.internal.cnn.InferenceInfo( ...
    "activations", dispatcherInputSize, options.MiniBatchSize, ...
    numObservations, executionEnvironment, precision, ...
    layerIndexInUnsortedLayers, acceleration, options.SequenceLength  );
[predictNetwork, gpuShouldBeUsed, workerLoad, useParallel] = setupPredictionEnvironment(this, inferenceInfo);
useMex = inferenceInfo.UseMex;

statefulInference = inferenceInfo.RecurrentInferenceRequiresState;

% Figure out activation sizes for specified layerIndex
activationSizes = inferOutputSizesGivenInputSizes(predictNetwork, dispatcherInputSize, layerIndex);
activationSizes = iWrapInCell(activationSizes{1});
activationHasSeqOutput = this.NetworkInfo.LayerHasSequenceOutput(layerIndex);
iAssertOutputAsIsValidForSeqOutput(activationHasSeqOutput, outputAs, this.Layers(layerIndex).Name);

% Get output size for specified layerIndex.
outputSize = activationSizes{layerOutputIndex};

% 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 is not a sequence network
outputIsAllocated = ...
    isfinite(numObservations) && ...
    ~any(this.NetworkInfo.HasSequenceInput) && ...
    ~this.NetworkInfo.HasLayerWithSizePropagationValueSink;

if useParallel
    numOfWorkers = nnz(workerLoad);

    % Create output cell - activations has only one output
    Y = cell(1,numOfWorkers);
    
    % 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, predictNetwork, gpuShouldBeUsed, ...
            layerIndex, layerOutputIndex, outputIsAllocated, outputAs, outputSize, ...
            precision, useMex, statefulInference, activationHasSeqOutput);
        Y{1,idx} = Ysub;
    end
    
    % Concatenate the workers' output following the OutputAs format
    if strcmp(outputAs,"channels") && ~activationHasSeqOutput
        catDim = numel(outputSize)+1; 
    elseif strcmp(outputAs,"columns")
        catDim = 2;
    else
        catDim = 1;
    end
    Y = cat(catDim,Y{:});
    
else
    Y = iDoInference(dispatcher, numObservations, predictNetwork, gpuShouldBeUsed, ...
        layerIndex, layerOutputIndex, outputIsAllocated, outputAs, outputSize, ...
        precision, useMex, statefulInference, activationHasSeqOutput);
end
end

function Y = iDoInference(dispatcher, numObservations, predictNetwork, gpuShouldBeUsed, ...
    layerIndex, layerOutputIndex, outputIsAllocated, outputAs, outputSize, ...
    precision, useMex, statefulInference, activationHasSeqOutput)

% Create the output variable
if outputIsAllocated
    % CNN activations for cases where we can preallocate the output
    [sz, indexFcn, reshapeFcn] = iGetOutputSizeAndIndices(...
        outputAs, numObservations, outputSize);
    Y = precision.zeros(sz);
else
    % Unable to preallocate output buffer when dispatcher does not know
    % number of observations in underlying data, or when the network is RNN
    Y = cell(0,1);
end

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

% Create a flag that tracks if the observation is a new one or the
% continuity of the previous one (used for stateful activations only,
% otherwise remains always true)
currentBatchHasNewObs = true;

% Use the dispatcher to iterate over the data
dispatcher.start();
while ~dispatcher.IsDone
    % Collect the next mini-batch
    [Xb, ~, info] = dispatcher.next();
    [Xb, mustReturnGpuArray] = iPrepDataForNetInput(Xb, gpuShouldBeUsed, useMex);
    
    % Move the allocated output on the GPU if necessary
    if ~allocatedOutputIsInExpectedEnvironment && outputIsAllocated && mustReturnGpuArray
        Y = gpuArray(Y);
    end
    allocatedOutputIsInExpectedEnvironment = true;

    % Compute the activations of this mini-batch & update the network
    [Yb, state] = predictNetwork.activations(Xb, layerIndex, layerOutputIndex);
    predictNetwork = iUpdateNetworkState(predictNetwork, state, layerIndex, statefulInference, info.IsCompleteObservation);
    
    % Activations only support one output
    Yb = Yb{1};
    
    % Gather the activations if needed
    if ~mustReturnGpuArray
        Yb = gather(Yb);
    end
    
    % MEX does not return the output on the GPU if there were multiple
    % inputs. Ensure the result on the device
    if useMex && mustReturnGpuArray
        Yb = gpuArray(Yb);
    end
    
    % Insert the mini-batch output in the output array
    if outputIsAllocated
        % Index the activations in the output array
        indices = indexFcn(info.MiniBatchIdx);
        Y(indices{:}) = reshapeFcn(Yb, numel(info.MiniBatchIdx));
    else
        if currentBatchHasNewObs
            % Grow Y with the new observation
            Y{end+1} = Yb; %#ok<AGROW>
        else
            % Stateful activations and the mini-batch is within the same
            % observation
            if activationHasSeqOutput
                % (seq2seq) Concatenate observations along sequence dimension
                Y{end} = cat(numel(outputSize) + 2, Y{end}, Yb);
            else
                % (seq2one) Overwrite previous observations with new
                Y{end} = Yb;
            end
        end
    end
    
    % Declare if next batch contains new observations
    currentBatchHasNewObs = info.IsCompleteObservation;
end

if ~outputIsAllocated
    % Concatenate the output if just accumulated as a cell array
    Y = iConcatenateCellOutput(Y, outputSize, activationHasSeqOutput);
    Y = iFormatActivations(Y, outputAs, outputSize);
end
end

function iValidateDispatcherForActivations(dispatchedInputSizes, networkInputSizes, inputNames, outputAs)
% iValidateDispatcherForActivations   Throws an error if the dispatcher
% doesn't give input that will work with the 'activations' method.

for n = 1:numel(networkInputSizes)
    if iDispatchesImagesSmallerThan(dispatchedInputSizes{n}, networkInputSizes{n})
        throwAsCaller( iCreateExceptionFromErrorID( ...
            'nnet_cnn:DAGNetwork:SmallerImagesForActivations', ...
            mat2str( dispatchedInputSizes{n} ), ...
            inputNames{n}, ...
            mat2str( networkInputSizes{n} ) ) );        
    elseif iDispatchesImagesMismatchedChannelDimension(dispatchedInputSizes{n}, networkInputSizes{n})
        throwAsCaller( iCreateExceptionFromErrorID( ...
            'nnet_cnn:DAGNetwork:MismatchedChannelDimensions', ...
            string(dispatchedInputSizes{n}(end)), ...
            inputNames{n}, ...
            string(networkInputSizes{n}(end)) ))
    elseif iDispatchesInputOfSize(dispatchedInputSizes{n}, networkInputSizes{n})
        % The 'activations' method can use this data
    elseif iLargerImagesCanBeOutputAs( outputAs )
        % The 'activations' method can use this data
    else
        throwAsCaller( iCreateExceptionFromErrorID( 'nnet_cnn:DAGNetwork:ActivationsLargeImagesOutputAs' ) );
    end
end
end

function iAssertOutputAsIsValidForSeqOutput(layerHasSeqOutput, outputAs, layerName)
if layerHasSeqOutput && ismember(outputAs, {'rows', 'columns'})
    throwAsCaller( iCreateExceptionFromErrorID( 'nnet_cnn:DAGNetwork:InvalidOutputAsForLayerWithSeqOutput', layerName ) );
end
end

function tf = iDispatchesImagesSmallerThan(dispatcherInputSize, inputSize)
% iDispatchesImagesSmallerThan   Check if the dispatcher dispatches images
% smaller than inputSize on the spatial dimensions.
numSpatialDims = numel(inputSize)-1;
if(isempty(dispatcherInputSize))
    tf = true;
    return;
end
tf = any(dispatcherInputSize(1:numSpatialDims) < inputSize(1:numSpatialDims));
end

function tf = iDispatchesImagesMismatchedChannelDimension(dispatcherInputSize, inputSize)
% iDispatchesImagesMismatchedChannelDimension   Check that the channel
% dimensions of the images match.
tf = dispatcherInputSize(end)~=inputSize(end);
end

function predictNetwork = iUpdateNetworkState(predictNetwork, state, layerIndex, statefulInference, needToResetState)
if statefulInference
    % Update the network's state
    predictNetwork = predictNetwork.updateNetworkState(state, layerIndex);
    if needToResetState
        predictNetwork = predictNetwork.resetNetworkState();
    end
else
    % This is no-op for a stateless inference call.
end
end

function sz = iInputSize( x )
% iInputSize   Return the size of x as [H W C] where C is 1 when x is a
% grayscale image.
if iIsDataDispatcher( x )
    sz = x.DataSize;
else
    sz = [size(x,1) size(x,2) size(x,3)];
end
end

function tf = iIsDataDispatcher(x)
tf = isa(x,'nnet.internal.cnn.DataDispatcher');
end

function X = iWrapInCell(X)
if ~iscell(X)
    X = {X};
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 tf = iDispatchesInputOfSize(dispatchedInputSize, inputSize)
tf = isequal( dispatchedInputSize, inputSize );
end

function tf = iLargerImagesCanBeOutputAs( outputAs )
tf = isequal(outputAs, 'channels');
end

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

function [outputBatchSize, indexFcn, reshapeFcn] = iGetOutputSizeAndIndices(outputAs, numObs, outputSize)
% Returns the output batch size, indexing function, and reshaping function.
% The indexing function provides the right set of indices based on the
% 'OutputAs' setting. The reshaping function reshapes channel
% formatted output to the shape required for the 'OutputAs' setting.
switch outputAs
    case 'rows'
        outputBatchSize = [numObs prod(outputSize)];
        indexFcn = @(i){i 1:prod(outputSize)};
        reshapeFcn = @(y,n)transpose(reshape(y, [], n));
    case 'columns'
        outputBatchSize = [prod(outputSize) numObs];
        indexFcn = @(i){1:prod(outputSize) i};
        reshapeFcn = @(y,n)reshape(y, [], n);
    case 'channels'
        outputBatchSize = [outputSize numObs];
        indices = arrayfun(@(x)1:x, outputSize, 'UniformOutput', false);
        indexFcn = @(i)[indices i];
        reshapeFcn = @(y,~)y;
end
end

function Y = iFormatActivations(Y, outputAs, outputSize)
switch outputAs
    case 'channels'
        % no-op
    case 'rows'
        Y = reshape(Y, prod(outputSize), []);
        Y = Y';
    case 'columns'
        Y = reshape(Y, prod(outputSize), []);
end
end

function Y = iConcatenateCellOutput(Y, outputSize, hasSeqOutput)
% Concatenate the input cell array. Assume batch dim is ndim+1.
if hasSeqOutput
    numOutputDims = numel(outputSize);
    Y = cellfun(@(y)iSequenceBatchFcn(y, numOutputDims), Y, 'UniformOutput', false);
    Y = vertcat(Y{:});
else
    batchDim = numel(outputSize)+1;
    Y = cat(batchDim,Y{:});
end
end

function seq = iSequenceBatchFcn(seq, numOutputDims)
seq = permute( seq, [1:numOutputDims (numOutputDims + 2) (numOutputDims + 1)] );
seq = squeeze( num2cell(seq, 1:(numOutputDims +1)) );
end
