function X = maxunpool(X, indices, outputSize, nvArgs)
%MAXUNPOOL   Unpool the output of maxpool
%
%   Y = MAXUNPOOL(X,INDICES,OUTPUTSIZE) unpools the input feature map X.
%   Input X must be a formatted dlarray. INDICES is a dlarray with the same
%   size as X containing the indices of the maximum values in each pooled
%   region. If INDICES is a formatted dlarray, it must have the same
%   dimension labels as X. OUTPUTSIZE is a numeric array that represents
%   the size of the output feature map after the unpool operation. INDICES
%   and OUTPUTSIZE are obtained from the maxpool function. Output Y is a
%   formatted dlarray with the same dimension labels as X.
%
%   Y = MAXUNPOOL(X,INDICES,OUTPUTSIZE,'DataFormat',FMT) unpools the input
%   feature map X, where X is an unformatted dlarray. FMT specifies the
%   dimension labels of the input data X. FMT must be a char array or
%   string. If INDICES is a formatted dlarray, it must have the same size
%   and format as X. If INDICES is an unformatted dlarray, its size and
%   dimension order must match the size and dimension order of X permuted
%   to match the following order - all dimensions with label 'S', then
%   dimensions with labels 'C', 'B', and 'T', then all dimensions with
%   label 'U'. OUTPUTSIZE represents the size of the output feature map Y,
%   with dimensions permuted according to the same order. Output Y is an
%   unformatted dlarray with the same dimension order as X.
%
%   Example 1:
%       % Create input data X as an unformatted dlarray with 2 spatial
%       % dimensions of size 4 and 1 channel.
%       numChannels = 1;
%       height = 4; width = 4;
%       X = dlarray(randn(height,width,numChannels));
%       % Apply the maxpool operation to input data X and obtain the
%       % indices of the maximum of each pooling region and output size.
%       [X,indices,outputSize] = maxpool(X,2,'Stride',2,'DataFormat','SSC');
%       % Apply the maxunpool operation with the vector of indices and the
%       % outputSize
%       Y = maxunpool(X,indices,outputSize,'DataFormat','SSC');
%       % Show the size of the output
%       size(Y)
%       % The expected output is a 4x4 dlarray
%
%   Example 2:
%       % Create input data X as a formatted dlarray with a channel
%       % dimension and a batch dimension, each of size 2, and a temporal
%       % dimension of size 4.
%       numChannels = 2; numObservations = 2;
%       sequenceLength = 4;
%       X = dlarray(randn(numChannels, numObservations, sequenceLength), 'CBT');
%       % Apply the maxpool operation to input data X and obtain the
%       % indices of the maximum of each pooling region and output size.
%       [X,indices,outputSize] = maxpool(X,2,'Stride',2,'PoolFormat','T');
%       % Apply the maxunpool operation with the vector of indices and the
%       % outputSize
%       Y = maxunpool(X,indices,outputSize);
%       % Show the size of the output
%       size(Y)
%       % The expected output is a 2x2x4 dlarray
%
%   See also DLARRAY, MAXPOOL

%   Copyright 2019-2021 The MathWorks, Inc.
arguments
    X {mustBeNonempty}
    indices {iMustBeValidIndices}
    outputSize (1,:) {deep.internal.dlarray.mustBeNonDlarray, ...
        mustBeNumeric, mustBePositive, mustBeInteger, mustBeNonempty}
    nvArgs.DataFormat {deep.internal.dlarray.mustBeDataFormatType} = ''
end

% Only data defines execution environment, not parameters
indices = gather(indices);
outputSize = gather(outputSize);

% Handle input label and check the validity of DataFormat against the input
% data X
[X, perm] = deep.internal.dlarray.validateDataFormatArg(X, nvArgs.DataFormat); 
dataFormat = dims(X);

% Find and validate the number of spatial dimensions
numSpatialDims = iValidateUnpoolingDims(dataFormat);

% Mutually validate arguments 
outputSize = iMutuallyValidateInputArguments(X, indices, outputSize, numSpatialDims, dataFormat);

% Internal API
X = deep.internal.dlarray.unpoolData(X, numSpatialDims, ...
                                     indices, outputSize);

% Remove dimensions from the dlarray if it was passed as unformatted.
if ~isempty(nvArgs.DataFormat)
    X = stripdims(X);
    X = ipermute(X, perm); % Permutes back in the same order as DataFormat
end
end

function numSpatialDims = iValidateUnpoolingDims(dataFormat)
maxSDims = 3;
numSpatialDims = count(dataFormat, 'S');

if ~(numSpatialDims > 0 || contains(dataFormat, 'T'))
    error(message('deep:dlarray:UnpoolingInvalidInputData'));
end

if numSpatialDims > maxSDims
    error(message('deep:dlarray:UnpoolingTooManySpatialDimensions'));
end
end

function outputSize = iMutuallyValidateInputArguments(X, indices, outputSize, numSpatialDims,dataFormat)
xsize = size(X);
outputSize = iAssertValidOutputSize(outputSize,xsize,numSpatialDims,dataFormat);
iAssertValidIndices(indices,outputSize,xsize,numSpatialDims,dataFormat);
end

function value = iAssertValidOutputSize(value,xsize,numSpatialDims,dataFormat)
% VALUE is the output size and should not be smaller than input size. VALUE
% is always ordered
if numel(xsize) ~= numel(value)
    [xsize, value] = iExpandInputOutputDimensions(xsize,value);
end
if ~contains(dataFormat,'C')
    xsize = [xsize(1:numSpatialDims) 1 xsize(numSpatialDims+1:end)];
    value = [value(1:numSpatialDims) 1 value(numSpatialDims+1:end)];
end
% Check the size of each dimension in output is greater than or
% equal to the size of the corresponding dimension in input.
if ~all(value>=xsize)
    error(message('deep:dlarray:UnpoolingWrongOutputSize'));
end
end

function [xsize, outputSize] = iExpandInputOutputDimensions(xsize, outputSize)
% Add singleton dimensions to xsize or outputSize to make the number of
% elements of the two size vectors match
maxLength = max(numel(xsize), numel(outputSize));
xsize(end+1:maxLength) = 1;
outputSize(end+1:maxLength) = 1;
end

function iAssertValidIndices(indices,outputSize,inputSize,numSpatialDims,datafmt)

if isdlarray(indices)
    % Check formats match, if formatted
    idxfmt = dims(indices);
    idxSpatialDims = count(idxfmt,'S');
    indicesSize = size(indices);
    if ~isempty(idxfmt)
        if idxSpatialDims ~= numSpatialDims || count(datafmt, 'T') ~= count(idxfmt, 'T')
            error(message('deep:dlarray:UnpoolingIndicesWrongFormat'))
        else
            uniqueIndicesLabels = ~ismember(idxfmt,datafmt);
            uniqueDataLabels = ~ismember(datafmt,idxfmt);
            if any(uniqueIndicesLabels) && ...
                    any(indicesSize(uniqueIndicesLabels == 1) ~= 1) || ...
                    any(uniqueDataLabels) && ...
                    any(inputSize(uniqueDataLabels == 1) ~= 1)
                error(message('deep:dlarray:UnpoolingIndicesWrongFormat'))
            end
        end
    end
    indices = extractdata(indices);
end

% Check sizes match
iValidateIndicesSize(indices, inputSize);

% Check that all indices are below or equal the numel of the output.
maxIndex = prod(outputSize);
if isdlarray(indices)
    % For dlarray it is faster to check on the underlying data
    indicesData = deep.internal.dlarray.toNumeric(indices, true);
else
    indicesData = indices;
end
if max(indicesData,[],'all') > maxIndex
    error(message('deep:dlarray:UnpoolingIndicesOutOfBounds', maxIndex))
end
end

function iValidateIndicesSize(indices,inputSize)
if isvector(indices)
    % If INDICES is a vector, it must have L elements, where L = S*...*S*C*B*T*U.
    numOfIndices = prod(inputSize);
    if numel(indices) ~= numOfIndices
        error(message('deep:dlarray:UnpoolingIndicesWrongNumElements',numOfIndices));
    end
else
    idxSizeForInputDims = size(indices, 1:numel(inputSize));
    if ~(all(idxSizeForInputDims==inputSize) && numel(indices)==prod(idxSizeForInputDims))
        error(message('deep:dlarray:UnpoolingIndicesWrongSize'))
    end
end
end

function iMustBeValidIndices(x)
mustBeNumeric(x);

% All checks can be performed on the underlying data if x is a dlarray
if isdlarray(x)
    xData = deep.internal.dlarray.toNumeric(x, true);
else
    xData = x;
end

if ~isreal(xData)
    throwAsCaller(MException(message('MATLAB:validators:mustBeReal')));
end

if isempty(xData)
    throwAsCaller(MException(message('MATLAB:validators:mustBeNonempty')));
end

if ~all(xData > 0, 'all')
    throwAsCaller(MException(message('MATLAB:validators:mustBePositive')));
end

if ~all(isfinite(xData), 'all') || ~all(xData == floor(xData), 'all')
    throwAsCaller(MException(message('MATLAB:validators:mustBeInteger')));
end
end
