function X = dltranspconv(X, weights, bias, nvArgs)
%DLTRANSPCONV  Deep learning transposed convolution
%
%    Y = DLTRANSPCONV(X,WEIGHTS,BIAS) computes the deep learning transposed
%    convolution of the input data X using filters defined by WEIGHTS and
%    adds a constant BIAS. X must be a formatted dlarray. Inputs WEIGHTS
%    and BIAS can be formatted dlarrays, unformatted dlarrays, or numeric
%    arrays. Output Y is a formatted dlarray with the same dimension labels
%    as X.
%
%    If WEIGHTS is an unformatted dlarray or a numeric array, then its size 
%    depends on the type of convolution:
%    -  For standard transposed convolution, WEIGHTS must have 
%       numSpatialDimensions + 2 dimensions, where numSpatialDimensions is 
%       the number of spatial dimensions of X. The first numSpatialDimensions 
%       correspond to the spatial dimensions of the filters and have 
%       matching sizes. The next two dimensions correspond to the filters 
%       and channels, respectively, and have matching sizes. 
%    -  For grouped transposed convolution, WEIGHTS must have an extra 
%       trailing dimension that corresponds to the groups. The size of the 
%       channel dimension must match the number of channels in the input 
%       data divided by the number of groups.
%    The size of the 'C' dimension of output Y is the product of the number 
%    of filters per group and number of groups.
%
%    If WEIGHTS is a formatted dlarray, then it can have multiple dimensions
%    labeled 'S' (spatial), one dimension labeled 'C' (channel), up to one 
%    dimension labeled 'T' (time), and up to two dimensions labeled 'U' 
%    (unspecified). WEIGHTS must have up to as many 'S' dimensions as X. If 
%    WEIGHTS has fewer 'S' dimensions than X, then the function 
%    automatically adds singleton 'S' dimensions to match the number of 'S' 
%    dimensions of X. If WEIGHTS has a 'T' dimension, then X must also
%    contain a 'T' dimension.
%
%    The format of WEIGHTS depends on the type of convolution:
%    -  For standard transposed convolution, WEIGHTS must have one 'C' 
%       dimension and at most one 'U' dimension. The size of the 'U' 
%       dimension must match the size of the 'C' dimension of input X. The 
%       size of 'C' dimension corresponds to the number of filters. If 
%       WEIGHTS specifies a single filter, then you can omit the 'C' 
%       dimension. The size of the 'C' dimension of output Y is equal to 
%       the number of filters.
%    -  For grouped transposed convolution, WEIGHTS must have one 'C' 
%       dimension and two 'U' dimensions. The first 'U' dimension specifies 
%       the number of input channels per group. Its size must be equal to 
%       the size of the 'C' dimension of input X divided by the number of 
%       groups. The second 'U' dimension specifies the number of groups. 
%       The size of the 'C' dimension of output Y is the product of the 
%       number of filters per group and the number of groups.
%
%    BIAS can be a numeric array or a formatted or unformatted dlarray with
%    up to one non-singleton dimension. If BIAS is a scalar or has only
%    singleton dimensions, the same bias is applied to each output.
%    Otherwise, the number of elements must match the number of filters
%    specified by WEIGHTS. If BIAS is a formatted dlarray, the non-singleton 
%    dimension must be labeled 'C'.
%
%    Y = DLTRANSPCONV(X,WEIGHTS,BIAS,'DataFormat',FMT) computes the deep
%    learning transposed convolution of an unformatted dlarray or a numeric
%    array X. FMT specifies the dimension labels of the input X. If
%    the input X is a numeric array, at least one of WEIGHTS or BIAS must
%    be a dlarray. The output Y is an unformatted dlarray with the same
%    dimension order as X.
%
%    Y = DLTRANSPCONV(_,'PARAM1',VAL1,'PARAM2',VAL2,...) specifies optional
%    parameter name/value pairs:
%
%      'Stride'                  - Step size for traversing the input,
%                                  specified as one of the following
%                                  - a scalar, in which case the same value
%                                     is used for all spatial dimensions.
%                                  - a vector [s1,...,sd] where si is
%                                    the stride along the ith dimension.
%                                  The default value is 1.
%      'DilationFactor'          - Step size for sampling the input or
%                                  equivalently the up-sampling factor of
%                                  the filter. Corresponds to an effective
%                                  filter size of filterSize +
%                                  + (filterSize-1) * (dilationFactor-1).
%                                  Specified as one of the following
%                                  - a scalar, in which case the same value
%                                     is used for all spatial dimensions.
%                                  - a vector [c1,...,cd] where ci is the
%                                    dilation along the ith dimension.
%                                  The default value is 1.
%      'Cropping'                - Amount to trim the edges of the full
%                                  transposed convolution, specified as
%                                  one of the following:
%                                  - 'same'. Cropping is set so that
%                                    the output size is the same as
%                                    the input size when the stride is 1.
%                                    More generally, the output size is
%                                    inputSize*stride, where
%                                    inputSize is the size of the input
%                                    along a spatial dimension.
%                                  - a scalar, in which case the same value
%                                     is used for all spatial dimensions.
%                                  - a vector [p1,...,pd] where pi is the
%                                    cropping applied to the start and end
%                                    along the ith spatial dimension.
%                                  - a matrix [low1,...,lowd;high1,...,
%                                    highd] of size 2-by-d where lowi and
%                                    highi are the cropping along the lower
%                                    and upper edges of the ith spatial
%                                    dimension.
%                                  The default value is 0.
%      'DataFormat'              - Dimension labels of the input data X,
%                                  specified as a char array or string.
%                                  Required if X is not a formatted dlarray.
%      'WeightsFormat'           - Dimension labels of the weights. 
%                                  Use this option when doing transposed 
%                                  convolution over a 'T' dimension and 
%                                  WEIGHTS is an unformatted dlarray or 
%                                  numeric array.
%
%
% Example 1:
%     % Create input data as formatted dlarray with 3 channels
%     X = dlarray(ones([64 64 3]),'SSC');
%
%     % Create weights as formatted dlarray representing 96 filters
%     % with height and width of 11
%     weights = dlarray(ones([11 11 96 3]),'SSCU');
%
%     % Compute deep learning transposed convolution
%     dltranspconv(X, weights, 0);
%
% Example 2:
%     % Create input data as formatted dlarray representing a batch of 7
%     % 64-by-64 activations with 8 channels
%     X = dlarray(ones([64 64 8 7]),'SSCB');
%
%     % Create weights as numeric array representing 2 groups of filters
%      % with 3 filters per group
%      weights = ones([5 5 3 4 2]);
%
%      % Compute deep learning transposed convolution
%      dltranspconv(X,weights,0);
%
% Example 3:
%      % Create input data as numeric array representing a batch of 7
%      % 64-by-64 activations with 8 channels
%      X = ones([64 64 8 7]);
%
%      % Create weights as formatted dlarray representing 8 groups of
%      % filters with 1 filter per group
%      weights = dlarray(ones([5 5 1 1 8]),'SSCUU');
%
%      % Compute deep learning transposed convolution
%      dltranspconv(X,weights,0,'DataFormat','SSCB');
%
% Example 4:
%     % Create a formatted dlarray representing a batch of 7 sequences with 
%     % 20 time steps and 8 channels
%     X = dlarray(ones([8 7 20]),'CBT');
%
%     % Create weights as numeric array representing 96 filters of size 5
%     weights = ones([5 96 8]);
%
%     % Compute deep learning transposed convolution, specifying the 
%     % 'WeightsFormat' option to carry out temporal transposed
%     % convolution.
%     dltranspconv(X,weights,0,'WeightsFormat','TCU');
%
%   See also AVGPOOL, DLARRAY, DLCONV, MAXPOOL, MAXUNPOOL

%   Copyright 2019-2021 The MathWorks, Inc.

arguments
    X
    weights
    bias
    nvArgs.DataFormat {deep.internal.dlarray.mustBeDataFormatType} = ''
    nvArgs.Stride {mustBeNumeric, mustBePositive, mustBeInteger, mustBeNonempty, deep.internal.dlarray.mustBeNonDlarray} = 1
    nvArgs.DilationFactor {mustBeNumeric, mustBePositive, mustBeInteger, mustBeNonempty, deep.internal.dlarray.mustBeNonDlarray} = 1
    nvArgs.Cropping {deep.internal.dlarray.mustBeNonDlarray} = 0
    nvArgs.WeightsFormat {deep.internal.dlarray.mustBeDataFormatType} = ''
end

% Ensure X is a dlarray and include DataFormat
[X, fmtPerm] = deep.internal.dlarray.validateDataFormatArg(X, nvArgs.DataFormat);

% Get labels of X
labelsX = dims(X);

% Validate the format of the weights
[weights, numSpatialDimsInWeights, convolveOverTime, isWeightsDlarray, isWeightsLabeled] = ...
    deep.internal.dlarray.validateWeightsFormat(weights, nvArgs.WeightsFormat);

% Get information about dimensions to convolve over
[convolutionDimsIdx, numConvolutionDims, numSpatialDimsToConvolve, numSpatialDimsInX] = ...
    deep.internal.dlarray.getConvolutionDimensions(labelsX, numSpatialDimsInWeights, convolveOverTime);

% Locate channel dimension and find number of channels of X
isChannelDimInX = any(labelsX=='C');
if isChannelDimInX
    nbChannelsInX = size(X, numSpatialDimsInX+1);
else
    % A singleton C dimension is assumed
    nbChannelsInX = 1;
end

% Extract and validate weights
[weights, permWeights, filterSize, nbFiltersPerGroup, nbChannelsPerGroup, nbGroups] = ...
    deep.internal.dlarray.extractConvolutionFilter(weights, isWeightsDlarray, numConvolutionDims, numSpatialDimsInX, numSpatialDimsInWeights, convolveOverTime);

% Check that the channel dimension in the weights matches that of the input
% data
isConvolution = false;
deep.internal.dlarray.validateConvolutionChannelDimension(numConvolutionDims+2, numConvolutionDims+3, ...
    isChannelDimInX, nbChannelsInX, nbChannelsPerGroup, nbFiltersPerGroup, nbGroups, isWeightsLabeled, isConvolution);

% Extract and validate bias
bias = deep.internal.dlarray.extractConvolutionBias(...
    bias, numConvolutionDims, numConvolutionDims+1, nbFiltersPerGroup*nbGroups, isWeightsLabeled, isConvolution);

% Find the spatial size of X
sizesConvolutionDimsInX = size(X,convolutionDimsIdx);

% Validate Name-Value pairs and convert them to canonical values.
[nvArgs, effectiveFilterSize, sizesConvolutionDimsInX] = deep.internal.dlarray.validateConvolutionArguments(nvArgs, numConvolutionDims, sizesConvolutionDimsInX, filterSize, 'Cropping', isConvolution);
[nvArgs.Cropping, outputSizeOffset] = ...
    iCalculateSameCroppingAndOutputSizeOffsetIfNeeded(sizesConvolutionDimsInX, filterSize, effectiveFilterSize, nvArgs.Cropping, nvArgs.Stride, nvArgs.DilationFactor);

internalArgs = {nvArgs.Cropping, nvArgs.Stride, nvArgs.DilationFactor, nbGroups, outputSizeOffset};

X = deep.internal.dlarray.convolveData(X, weights, permWeights, bias, internalArgs, numConvolutionDims, numSpatialDimsToConvolve, convolveOverTime, isChannelDimInX, isConvolution);

% Remove dimensions from the dlarray if it was passed as unformatted.
if ~isempty(nvArgs.DataFormat)
    X = stripdims(X);

    % Permutes back in the same order as DataFormat
    X = ipermute(X, fmtPerm);
end
end

function [croppingSize, outputSizeOffset] = iCalculateSameCroppingAndOutputSizeOffsetIfNeeded(sizesSpatialDimsInX, filterSize, effectiveFilterSize, croppingSize, stride, dilation)
if isempty(croppingSize)
    % Same cropping case: Set cropping size and outputSizeOffset for
    %    a given stride and filter size in such a way that
    %    1) The outputSize is stride*inputSize
    %    2) Cropping is minimum
    KMinusS = effectiveFilterSize - stride;
    outputSizeOffset = max(0, -KMinusS);
    totalCropping = KMinusS + outputSizeOffset;

    croppingSize = [floor(totalCropping/2); ceil(totalCropping/2)];
    croppingSize = croppingSize(:).';
else
    % Manual cropping case: find minimum outputSizeOffset such that the
    % output size is positive.
    outputSizeOffset = iComputeMinOutputSizeOffset(sizesSpatialDimsInX, filterSize, croppingSize, stride, dilation);
end
end

function outputSizeOffset = iComputeMinOutputSizeOffset(sizesSpatialDimsInX, filterSize, croppingSize, stride, dilation)
% If cropping not same, compute dynamically
% outputSizeOffset, as the minimum allowed value.
% This is the minimum value such that
% 1) 0 <= outputSizeOffset < stride
% 2) The output size is positive.
totalCropping = croppingSize(1:2:end) + croppingSize(2:2:end);

% If filterSize is scalar, we are doing 1D transposed convolution and need 
% to add atrailing singleton dimension
if isscalar(filterSize)
    filterSize = [filterSize 1];
end

outputSizeWithoutOffset = stride .* (sizesSpatialDimsInX-1) + filterSize.*dilation - (dilation-1) - totalCropping;
minOutputSizeOffset = max(1 - outputSizeWithoutOffset, 0);

if any(minOutputSizeOffset >= stride)
    % If such value does not exist, error.
    error(message('deep:dlarray:TranspConvInvalidInputSize'));
else
    outputSizeOffset = minOutputSizeOffset;
end
end
