function X = dlconv(X,weights,bias,varargin)
%DLCONV  Deep learning convolution 
%
%    Y = DLCONV(X,WEIGHTS,BIAS) computes the deep learning 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, it must be
%    specified as a filterSize-by-numChannelsPerGroup-by-numFiltersPerGroup
%    -by-numGroups array, where filterSize must have as many entries as
%    the number of spatial dimensions of X. If numGroups is greater than 1,
%    the function performs grouped convolutions. Otherwise, numGroups is
%    singleton and can be omitted. The size of the 'C' dimension of
%    output Y is the product of numFiltersPerGroup and numGroups.
%
%    If WEIGHTS is a formatted dlarray, it can have multiple spatial
%    dimensions labeled 'S', one channel dimension labeled 'C', and up to
%    two other dimensions labeled 'U'. Input WEIGHTS can have a number of
%    'S' dimensions up to the number of 'S' dimensions of the input X. If
%    WEIGHTS has fewer 'S' dimensions than X, the 'S' dimensions of WEIGHTS
%    are expanded to match the number of 'S' dimensions of X.
%
%    If WEIGHTS contains less than two 'U' dimensions, DLCONV performs
%    a convolution with no groups. The size of the 'C' dimension must
%    match the size of the 'C' dimension of input X. The 'U' dimension
%    specifies the number of filters. The size of the 'C' dimension of
%    output Y is equal to the number of filters.
%
%    If WEIGHTS contain two 'U' dimensions, DLCONV performs a grouped
%    convolution. The first 'U' dimension specifies the number of filters
%    per group. The second 'U' dimension specifies the number of groups.
%    The size of the 'C' dimension of WEIGHTS must be equal to the size of
%    the 'C' dimension of input X divided by the number of groups. The size
%    of the 'C' dimension of output Y is the product of the size of the 'U'
%    dimensions of WEIGHTS.
%
%    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 has any dimension labels, the non-
%    -singleton dimension must be labeled 'C'.
%
%    Y = DLCONV(X,WEIGHTS,BIAS,'DataFormat',FMT) computes the deep learning
%    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 = DLCONV(_,'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.
%      'Padding'                 - Padding applied to the input along
%                                  the edges, specified as one of
%                                  the following:
%                                  - 'same'. Padding 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
%                                    ceil(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
%                                    padding 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 padding along the lower
%                                    and upper edges of the ith spatial
%                                    dimension.
%                                  The size of the padding must be less
%                                  than the filter size for each 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.
%
%
% 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 3 96]),'SSCU');
%
%     % Compute deep learning convolution
%     dlconv(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 4 3 2]);
%
%     % Compute deep learning convolution
%     dlconv(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 convolution
%     dlconv(X,weights,0,'DataFormat','SSCB')
%
%   See also BATCHNORM, DLARRAY, DLTRANSPCONV, FULLYCONNECT, MAXPOOL, RELU

%   Copyright 2019-2020 The MathWorks, Inc.

narginchk(3, inf);

% Parse Name-Value pairs with no validation
args = iParseInputArguments(varargin{:});

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

% Get labels of X
labelsXview = dims(X);
nLabels = numel(labelsXview);

% Get internal permutation of X
permXData = iGetInternalPermutation(X, nLabels);

% Currently the dlconv supports only 1D, 2D, and 3D convolution.
nSpatialDimsInX = count(labelsXview, 'S');
if nSpatialDimsInX == 0 || nSpatialDimsInX > 3
    error(message('deep:dlarray:ConvUnsupportedNumberOfSpatialDimensions'))
end

% Locate channel dimension and find number of channels of X
channelDimInX = find(labelsXview == 'C');
isChannelDimInX = ~isempty(channelDimInX);
if isChannelDimInX
    nbChannelsInX = size(X, channelDimInX);
else
    % Add singleton C dimension to the permXData vector
    permXData = [permXData(1:nSpatialDimsInX) nLabels+1 permXData(nSpatialDimsInX+1:end)];
    nbChannelsInX = 1;
end

% Extract data in the format needed by builtin convolution functions.
fd = X.FormattedData;
X.FormattedData = [];
[fd,xdata] = extractData(fd);

% Extract weights & bias
[weights, permWeights, nbFilters, nbGroups, isWeightsLabeled] = iExtractWeights(weights, nSpatialDimsInX, nbChannelsInX);
bias = iExtractBias(bias, nSpatialDimsInX, nbFilters, isWeightsLabeled);

% The dlarray methods should not accept logical inputs 
if islogical(weights) || islogical(bias)
    error(message('deep:dlarray:LogicalsNotSupported'));
end

% Grouped convolution not supported for 3D 
if nbGroups~=1 && nSpatialDimsInX>2
    error(message('deep:dlarray:Conv3DWithGroups'))
end

% Convolution that creates 'C' dimension that was not there in the
% input data is not supported.
if ~isChannelDimInX && nbFilters~=1
    error(message('deep:dlarray:ConvWeightsWithFiltersXWithNoChannel'))
end

% Find the filter size.
filterSize = iSpatialDimsSize(weights, permWeights, nSpatialDimsInX);

% Validate Name-Value pairs
sizesSpatialDimsInX = iSpatialDimsSize(xdata, permXData, nSpatialDimsInX);
args = iValidateInputArguments(args, nSpatialDimsInX, sizesSpatialDimsInX, filterSize);

iValidateFilterSizeAgainstPaddedInput(filterSize, nSpatialDimsInX, ...
    sizesSpatialDimsInX, args.PaddingSize, args.DilationFactor, isWeightsLabeled);

[xdata, permOutput] = internal_dlconv_permuted(xdata,permXData,weights,permWeights,bias,...
    args.PaddingSize,args.Stride,args.DilationFactor,nSpatialDimsInX);

% Remove channel dim if initially not present
if ~isChannelDimInX
    sizeOutput = size(xdata);
    nDimsOutput = numel(sizeOutput);
    channelDim = nSpatialDimsInX+1;
    channelDimInOutput = permOutput(channelDim);
    
    % Remove channel dim from xdata
    if nDimsOutput > 2
        if channelDimInOutput <= nDimsOutput
            sizeOutput(channelDimInOutput) = [];
            xdata = reshape(xdata, sizeOutput);
        end
    elseif channelDimInOutput == 1
        % When channel dim is first dimension out of two dimensions in the
        % output then add singleton dimension at the end
        sizeOutput = [sizeOutput(2) 1];
        xdata = reshape(xdata, sizeOutput);
    end
    
    permOutput = iRemovePermEntry(permOutput, channelDim);
end

if ~isempty(args.DataFormat)
    % Remove dimensions labels
    fd.DimensionMetadata = '';
    
    % Insert the data back in FormattedData
    fd = insertData(fd, xdata);
    
    % Insert the FormattedData back in the dlarray
    X.FormattedData = fd;
    comibinedPerm(permFormat) = permOutput;
    inverseComibinedPerm(comibinedPerm) = 1:nLabels;
    X.ViewToDataMap = comibinedPerm;
    X.DataToViewMap = inverseComibinedPerm;
else
    % Update the order of labels in FormattedData
    inversePermOutput(permOutput) = 1:length(permOutput);
    fd.DimensionMetadata = labelsXview(inversePermOutput);
    
    % Insert the data back in FormattedData
    fd = insertData(fd, xdata);
    
    % Insert the FormattedData back in the dlarray
    X.FormattedData = fd;
    X.ViewToDataMap = permOutput;
    X.DataToViewMap = inversePermOutput;
end
end

function [weights, permWeights, nbFilters, nbGroups, isWeightsLabeled] = iExtractWeights(weights, nSpatialDimsInX, nbChannelsInX)
isWeightsLabeled = false;
if isa(weights, 'dlarray')
    % Extract labels, numerical weights and their permutation
    fd = weights.FormattedData;
    weightsLabel = fd.DimensionMetadata;
    ndimsWeights = length(weightsLabel);
    permWeights = iGetInternalPermutation(weights, ndimsWeights);
    [~,weights] = extractData(fd);
    
    if ~isempty(weightsLabel)
        isWeightsLabeled = true;
        
        % Find number of spatial dims in weights
        nSpatialDimsInW = nnz(weightsLabel == 'S');
        
        % Ensure no 'T' or 'B' dims
        if(any(weightsLabel == 'T') || any(weightsLabel == 'B'))
            error(message('deep:dlarray:ConvWeightsContainBatchDim'))
        end
        
        % If W different number of spatial dims than X
        if(nSpatialDimsInW > nSpatialDimsInX)
            %If W contains too many spatial dims error
            error(message('deep:dlarray:ConvWeightsHaveTooManySpatial', nSpatialDimsInW, nSpatialDimsInX));
        elseif(nSpatialDimsInW < nSpatialDimsInX)
            permWeights = [permWeights(1:nSpatialDimsInW) ...
                (numel(permWeights)+1):(numel(permWeights)+nSpatialDimsInX-nSpatialDimsInW) ...
                permWeights((nSpatialDimsInW+1):end)];
        end
        
        % Retrieve channel index & number of channels
        channelDim = find(weightsLabel == 'C');
        if isempty(channelDim)
            nbChannelsPerGroupInW = 1;
        else
            nbChannelsPerGroupInW = size(weights, channelDim);
        end
        
        % If the channel dimension was not present we need to add it to the
        % permutation vector
        if isempty(channelDim)
            permWeights = [permWeights(1:nSpatialDimsInX) numel(permWeights)+1 ...
                permWeights(nSpatialDimsInX+1:end)];
        end
        
        % Behaviors for different number of 'U' dims
        unspecifiedDims = find(weightsLabel == 'U');
        nPerm = numel(permWeights);
        if(numel(unspecifiedDims) == 0)
            nbGroups = 1;
            nbFilters = 1;
            permWeights = [permWeights (nPerm+1):(nPerm+2)];
        elseif(numel(unspecifiedDims) == 1)
            nbGroups = 1;
            % The 'U' dim encodes number of filters
            nbFilters = size(weights, unspecifiedDims);
            permWeights = [permWeights nPerm+1];
        elseif(numel(unspecifiedDims) == 2)
            % Undo permutation flipping 'U' dims
            if permWeights(end-1) < permWeights(end)
                unspecifiedDims = flip(unspecifiedDims);
            end
            
            % The first 'U' dim encodes number of filters per group
            nbGroups = size(weights, unspecifiedDims(1));
            % The second 'U' dim encodes number of groups
            nbFilters = nbGroups * size(weights, unspecifiedDims(2));
        else
            % If more than 2 'U' dims
            error(message('deep:dlarray:ConvWeightsHaveMultipleOtherDims'))
        end
    else
        % If weights are unformatted dlarray
        permWeights = [permWeights (numel(permWeights)+1):(nSpatialDimsInX+3)];
        [nbChannelsPerGroupInW, nbFilters, nbGroups] = iSizesOfNumericWeights(weights, permWeights, nSpatialDimsInX);
    end
else
    % If weights are numeric array
    deep.internal.dlarray.validateNonDlarray(weights);
    permWeights = 1:(nSpatialDimsInX+3);
    [nbChannelsPerGroupInW, nbFilters, nbGroups] = iSizesOfNumericWeights(weights, permWeights, nSpatialDimsInX);
end

% Channels in W must match channels in X
if (nbChannelsPerGroupInW * nbGroups ~= nbChannelsInX)
    if(nbGroups==1)
        if isWeightsLabeled
            error(message('deep:dlarray:ConvLabeledWeightsChannelMismatch', nbChannelsInX, nbChannelsPerGroupInW))
        else
            error(message('deep:dlarray:ConvChannelMismatch', nbChannelsInX, nbChannelsPerGroupInW, nSpatialDimsInX+1))
        end
    else
        if isWeightsLabeled
            error(message('deep:dlarray:ConvLabeledWeightsChannelsPerGroupMismatch', nbChannelsInX, nbChannelsPerGroupInW, nbGroups))
        else
            error(message('deep:dlarray:ConvChannelsPerGroupMismatch', nbChannelsInX, nbChannelsPerGroupInW, nSpatialDimsInX+1, nbGroups, nSpatialDimsInX+3))
        end
    end
end
end

function [nbChannelsPerGroupInW, nbFilters, nbGroups] = iSizesOfNumericWeights(weights, permWeights, nSpatialDimsInX)
% For the case when weights are unlabeled find the number of channels, 
% filters, and groups from a numeric array containing weights.

if isempty(weights)
    error(message('deep:dlarray:ConvEmptyWeights'))
end

if(ndims(weights) > nSpatialDimsInX+3)
    error(message('deep:dlarray:ConvWeightsHaveTooManyDims', nSpatialDimsInX))
end

% The case with groups encoded in the last dim
nbGroups = size(weights, permWeights(nSpatialDimsInX+3));

% The number of filters dim is just after channel dim
nbFilters = nbGroups * size(weights, permWeights(nSpatialDimsInX+2));

% The channel dim is just after spacial dims
nbChannelsPerGroupInW = size(weights, permWeights(nSpatialDimsInX+1));
end

function iValidateFilterSizeAgainstPaddedInput(filterSize, nSpatialDims, ...
    sizesSpatialDimsInX, paddingSize, dilationFactor, isWeightsLabeled)

% Check the spatial sizes of filter is not larger than those of activations
sizeOfFilters = zeros(1, nSpatialDims);
paddedInputSize = zeros(1, nSpatialDims);
for i = 1:nSpatialDims
    sizeOfFilters(i) = filterSize(i)*dilationFactor(i) - dilationFactor(i) + 1;
    paddedInputSize(i) = sizesSpatialDimsInX(i) + paddingSize(2*i-1) + paddingSize(2*i);
end

if any(sizeOfFilters > paddedInputSize)
    if isWeightsLabeled
        error(message('deep:dlarray:ConvTooLargeLabeledFilter', mat2str(paddedInputSize), mat2str(sizeOfFilters)))
    else
        error(message('deep:dlarray:ConvTooLargeUnlabeledFilter', mat2str(paddedInputSize), mat2str(sizeOfFilters), nSpatialDims))
    end
end
end

function bias = iExtractBias(bias, nSpatialDims, nbFilters, isWeightsLabeled)

% Extract bias if wrapped in a dlarray.
if isa(bias, 'dlarray')
    [fd,bias] = extractData(bias.FormattedData);
    
    % If formatted dlarray
    if ~isempty(fd.DimensionMetadata)
        channelIdx = channelIndex(fd);
        
        % Bias must contain a 'C' dimension and it must not contain other
        % non-singleton dimension apart from the 'C' dimensions
        if ~contains(fd.DimensionMetadata,'C') || ...
            (~isscalar(bias) && size(bias, channelIdx) ~= numel(bias))
            error(message('deep:dlarray:ConvBiasWithWrongDims'));
        end
    end
else
    deep.internal.dlarray.validateNonDlarray(bias);
end

nbElementsBias = numel(bias);
if(nbElementsBias == 1)
    % Scalar bias is interpreted as constant vector.
    bias = bias(ones(1, nbFilters));
elseif nnz(size(bias) ~= 1) > 1
    % Bias cannot have more than one non-singleton dimension.
    error(message('deep:dlarray:ConvNonScalarOrVectorBias'));
elseif(nbElementsBias ~= nbFilters)
    % The number of elements of Bias must equal the number of filters.
    if isWeightsLabeled
        error(message('deep:dlarray:ConvLabeledWeightsMismatchBias', nbElementsBias, nbFilters));
    else
        error(message('deep:dlarray:ConvWeightsMismatchBias', nbElementsBias, nbFilters, nSpatialDims+2));
    end
end
end

function inputArguments = iParseInputArguments(varargin)
varargin = nnet.internal.cnn.layer.util.gatherParametersToCPU(varargin);
parser = iCreateParser();
parser.parse(varargin{:});
inputArguments = parser.Results;
end

function p = iCreateParser()
p = inputParser;
defaultStride = 1;
defaultDilationFactor = 1;
defaultPadding = 0;
defaultFormat = [];

p.addParameter('Stride', defaultStride);
p.addParameter('DilationFactor', defaultDilationFactor);
p.addParameter('Padding', defaultPadding);
p.addParameter('DataFormat', defaultFormat, @(x)validateattributes(x, {'char', 'string'}, {'nonempty'}));
end

function args = iValidateInputArguments(args, dims, sizesSpatialDimsInX, filterSize)
iAssertValidStride(args.Stride, dims);
iAssertValidDilationFactor(args.DilationFactor, dims);
iAssertValidPadding(args.Padding, dims);

args = iConvertToCanonicalForm(args, dims, sizesSpatialDimsInX, filterSize);
end

function iAssertValidStride(value,dims)
validateattributes(value, {'numeric'}, ...
    {'positive', 'real', 'integer', 'nonempty'});

if ~(isscalar(value) || iIsRowVector(value,dims))
    error(message('deep:dlarray:ConvStrideMustBeScalarOrVector'));
end
end

function iAssertValidDilationFactor(value,dims)
validateattributes(value, {'numeric'}, ...
    {'positive', 'real', 'integer', 'nonempty'});

if ~(isscalar(value) || iIsRowVector(value,dims))
    error(message('deep:dlarray:ConvDilationMustBeScalarOrVector'));
end
end

function iAssertValidPadding(padding, dims)
deep.internal.dlarray.assertValidPaddingOrCropping(padding, dims, 'Padding');
end

function inputArguments = iConvertToCanonicalForm(args, dims, sizesSpatialDimsInX, filterSize)
% Make sure integral values are converted to double and strings to char vectors
inputArguments = struct;
inputArguments.Stride = double(iMakeIntoRowVector(args.Stride,dims));
inputArguments.DilationFactor = double(iMakeIntoRowVector(args.DilationFactor,dims));
inputArguments.PaddingSize = double(nnet.internal.cnn.layer.padding.calculatePaddingSize(args.Padding,dims));

% If padding is 'same'
if isempty(inputArguments.PaddingSize)
    % Accont for dilation.
    effectiveFilterSize = (filterSize-1).*inputArguments.DilationFactor + 1;

    if dims == 1
        % Work around calculateSamePadding supporting only 2D and 3D cases
        inputArguments.PaddingSize = nnet.internal.cnn.layer.padding.calculateSamePadding([effectiveFilterSize 1],inputArguments.Stride,[sizesSpatialDimsInX 1]);
        inputArguments.PaddingSize = inputArguments.PaddingSize(1:2);
    else
        inputArguments.PaddingSize = nnet.internal.cnn.layer.padding.calculateSamePadding(effectiveFilterSize,inputArguments.Stride,sizesSpatialDimsInX);
    end    
end

inputArguments.DataFormat = args.DataFormat;
end

function rowVector = iMakeIntoRowVector(scalarOrRowVector,dims)
if(iIsRowVector(scalarOrRowVector,dims))
    rowVector = scalarOrRowVector;
else
    rowVector = repmat(scalarOrRowVector,1,dims);
end
end

function tf = iIsRowVector(x,dims)
tf = isrow(x) && numel(x)==dims;
end

function spatialDimsSize = iSpatialDimsSize(array, permWeights, nSpatialDims)
indicies = permWeights(1:nSpatialDims);
spatialDimsSize = size(array, indicies);
end

function internalPerm = iGetInternalPermutation(X, nDimensions)
% Get the internal permutation of dlarray X with nDimensions
if isempty(X.ViewToDataMap)
    internalPerm = 1:nDimensions;
else
    internalPerm = X.ViewToDataMap;
end
end

function permutation = iRemovePermEntry(permutation, entryToRemove)
% Remove one entry from a permutation vector.
valueRemoved = permutation(entryToRemove);
permutation(entryToRemove) = [];

% Deflate all other entries larger than valueRemoved
permutation = permutation - (permutation>valueRemoved);
end