function varargout = activationsMIMO(this, X, varargin)
% activationsMIMO   Get activations from a MIMO network
%
%   [out1, ...] = activationsMIMO(net, X, layerOut) returns
%   network activations using the single input network net and
%   the data X. layerOut is a cell array of character vectors.
%   Each character vector is either the name of a layer, or if
%   the layer has multiple outputs, it is the name of the layer
%   followed by the '/' character, followed by the name of the
%   layer output. If layerOut has M elements, the function will
%   return M outputs.
%
%   For a network with a single image input layer, specify X as
%   one of the following:
%       - A 4-D array of images, where the first three
%         dimensions correspond to the height, width, and
%         channels of an image, and the fourth dimension
%         correspond to the image index.
%       - An image datastore.
%       - A matlab.io.datastore.MiniBatchDatastore.
%       - A table, where the first column contains either image
%         paths or images.
%
%   [out1, ...] = activationsMIMO(net, ds, dataMap, layerOut)
%   returns network activations using the network net (which
%   may have any number of inputs), a datastore ds, and a
%   table dataMap. layerOut is a cell array of character
%   vectors. Each character vector is either the name of a
%   layer, or if the layer has multiple outputs, it is the name
%   of the layer followed by the '/' character, followed by the
%   name of the layer output. If layerOut has M elements, the
%   function will return M outputs. dataMap is a table with two
%   columns "Data" and "Destination" describing the mapping
%   from the datastore to the layers of the network.
%
%   [out1, ...] = activationsMIMO(..., 'PARAM1', VAL1, ...)
%   will compute activations 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,
%                           activations returns an error
%                           message.
%                         - 'cpu' - Use the CPU.
%                       The default is 'auto'.
%
%       'MiniBatchesInCell'
%                     - Logical value that specifies whether
%                       each mini-batch should be returned in
%                       a cell or not. The default is false.
%
%       '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'.
%
%   See also DAGNetwork/predict, DAGNetwork/classify.

%   Copyright 2018-2022 The MathWorks, Inc.

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

% Deduce whether a mapping is being used
mappingIsUsed = istable(varargin{1}) || isstruct(varargin{1});
if mappingIsUsed
    layerOutputs = varargin{2};
    nameValuePairs = varargin(3:end);
else
    layerOutputs = varargin{1};
    nameValuePairs = varargin(2:end);
end

% Validate the layer outputs
iValidateLayerOutputs(layerOutputs);

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

% Convert the source character vectors
if ~iscell(layerOutputs)
    layerOutputs = {layerOutputs};
end
[~, layerIndices, ~, layerOutputIndices] = ...
    iGetSourceInformationFromCellArray(layerOutputs, this.Layers);

% Get the layer indices in topological order.
[ord, layerIndices] = find(reshape(this.getTopologicalOrder(),1,[]) == layerIndices);

% If a mapping was used, we must create a special kind of
% dispatcher which can handle multiple inputs and outputs.
if mappingIsUsed
    if isstruct(varargin{1})
        % If the mapping is a struct, use GeneralDatastoreDispatcher.
        mapping = varargin{1};
        dispatcher = nnet.internal.cnn.dispatcher.GeneralDatastoreDispatcher( X,...
            options.MiniBatchSize, 'truncateLast', 'never', precision,...
            mapping.DataTranslator, mapping.CollateFunctions, ...
            mapping.InputFormats, mapping.OutputFormats, mapping.UnpackCellInputOutput);
    else
        translator = nnet.internal.cnn.util.TableTranslator( ...
            varargin{1}, this.PrivateNetwork);
        dispatcher = nnet.internal.cnn.dispatcher.MultiInputMultiOutputMiniBatchDatastoreDispatcher( ...
            X, options.MiniBatchSize, 'truncateLast', 'never', precision, translator);
    end
else
    dispatcher = this.createPredictionDataDispatcher(X, options, precision, this.NetworkInfo);
end

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

% Assertion is safe because activationsMIMO is internal only (Hidden)
assert(~inferenceInfo.UseMex, 'MEX Acceleration is not supported for activationsMIMO');

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

% Allocate space for the layer outputs
numActivations = numel(layerOutputs);
numMiniBatches = ceil(dispatcher.NumObservations/options.MiniBatchSize);
if isinf(numMiniBatches)
    % GeneralDatastore dispatcher will set the NumObservations to inf, since
    % the number of observations are unknown for any datastore.
    % Set the preallocation to 1 for now. There maybe some optimizations that
    % could be done with preallocation here.
    numMiniBatches = 1;
end
Z = cell(numMiniBatches, numActivations);

% Get all of the outputs for this network
dispatcher.start();
miniBatchIndex = 1;
while ~dispatcher.IsDone
    X = dispatcher.next();
    if ~iscell(X)
        X = {X};
    end
    if gpuShouldBeUsed
        X = iMoveCellToGPU(X);
    end
    
    activationsBufferBatch = predictNetwork.activations( ...
        X, layerIndices, layerOutputIndices, false);
    
    for i = 1:numActivations
        Z{miniBatchIndex,ord(i)} = gather( activationsBufferBatch{i} );
    end
    miniBatchIndex = miniBatchIndex + 1;
end

varargout = cell(1,numActivations);
if miniBatchesInCell
    for i = 1:numActivations
        varargout{i} = Z(:,i)';
    end
else
    for i = 1:numActivations
        varargout{i} = cat(4,Z{:,i});
    end
end
end

function iValidateLayerOutputs(layerOutputs)
if iIsValidStringOrCharArray(layerOutputs) || iIsValidVectorCellOfStringOrCharArray(layerOutputs)
else
    throwAsCaller( iCreateExceptionFromErrorID( ...
        'nnet_cnn:DAGNetwork:ActivationMIMOLayerOutputMustBeString' ) );
end
end

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

function tf = iIsValidVectorCellOfStringOrCharArray(value)
if iscell(value) && isvector(value)
    tf = all(cellfun(@(x)iIsValidStringOrCharArray(x), value));
else
    tf = false;
end
end

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

function [startLayerNames, startLayerIndices, layerOutputNames, layerOutputIndices] = ...
    iGetSourceInformationFromCellArray(sCell, layers)
cellfun(@(x)iValidateThereAreNotMultipleForwardSlashesInSource(x), sCell);
sSplits = cellfun(@(x)strsplit(x, '/'), sCell, 'UniformOutput', false);
startLayerNames = cellfun(@(x)x{1}, sSplits, 'UniformOutput', false);
cellfun(@(x)iValidateLayerName(x, layers), startLayerNames);
startLayerIndices = cellfun(@(x)iConvertLayerNameToIndex(x, layers), startLayerNames);
layerOutputNames = repmat({'out'},numel(sCell),1);
for i = 1:numel(sCell)
    if numel(sSplits{i}) == 2
        layerOutputNames{i} = sSplits{i}{2};
    else
        iThrowErrorIfLayerHasMultipleOutputs( layers(startLayerIndices(i)) );
    end
end
startLayerIndices = num2cell(startLayerIndices');
layerOutputNames = iWrapInCell(layerOutputNames);
layerOutputIndices = cellfun(@(x,y)iGetLayerOutputIndex(layers(x),y), ...
    startLayerIndices, layerOutputNames);
startLayerIndices = cell2mat(startLayerIndices);
end

function iValidateThereAreNotMultipleForwardSlashesInSource(s)
if iContainsMoreThanOneBackslash(s)
    error(message('nnet_cnn:DAGNetwork:ActivationLayerOutputCannotHaveMultipleForwardSlashes'));
end
end

function tf = iContainsMoreThanOneBackslash(x)
tf = length(strfind(x, '/')) > 1;
end

function iValidateLayerName(layerName, externalLayers)
if ~iLayerNameExists(layerName, externalLayers)
    error(message('nnet_cnn:DAGNetwork:LayerDoesNotExist', layerName))
end
end

function tf = iLayerNameExists(layerName, externalLayers)
layerNames = {externalLayers.Name}';
tf = any(strcmp(layerNames, layerName));
end

function layerIndex = iConvertLayerNameToIndex(layerName, externalLayers)
layerNames = {externalLayers.Name}';
layerIndex = find(strcmp(layerNames, layerName));
end

function iThrowErrorIfLayerHasMultipleOutputs(layer)
if iLayerHasMultipleOutputs(layer)
    error(message('nnet_cnn:DAGNetwork:MustSpecifyOutputForMultipleOutputLayer', layer.Name));
end
end

function tf = iLayerHasMultipleOutputs(layer)
% Note: this is a temporary fix and will be removed once all external
% layers have the OutputNames property
isCustomMimoLayer = isa(layer,'nnet.layer.Layer') && layer.NumOutputs>1;
isMaxPoolWithUnpooling = isa(layer, 'nnet.cnn.layer.MaxPooling2DLayer') && layer.HasUnpoolingOutputs;
tf = isCustomMimoLayer || isMaxPoolWithUnpooling;
end

function X = iWrapInCell(X)
if ~iscell(X)
    X = {X};
end
end

function layerOutputIndex = iGetLayerOutputIndex(layer, layerOutputName)
% Note: this is a temporary fix and will be removed once all external
% layers have the OutputNames property
if isa(layer,'nnet.cnn.layer.MaxPooling2DLayer')
    maxPoolingOutputNames = {'out', 'indices', 'size'};
    iValidateLayerOutputName(layer.Name, layerOutputName, maxPoolingOutputNames);
    layerOutputIndex = find(strcmp(maxPoolingOutputNames, layerOutputName));
elseif isa(layer,'nnet.layer.Layer')
    iValidateLayerOutputName(layer.Name, layerOutputName, layer.OutputNames);
    layerOutputIndex = find(strcmp(layer.OutputNames, layerOutputName));
else
    iValidateLayerOutputName(layer.Name, layerOutputName, 'out');
    layerOutputIndex = 1;
end
end

function iValidateLayerOutputName(layerName, layerOutputName, validOutputNames)
if ~any(strcmp(layerOutputName, validOutputNames))
    error(message('nnet_cnn:DAGNetwork:NonExistentLayerOutput', layerName, layerOutputName));
end
end

function X = iMoveCellToGPU(X)
X = cellfun(@(x)gpuArray(x), X, 'UniformOutput', false);
end
