function X = groupnorm(X, numGroups, offset, scale, nvArgs)
%GROUPNORM   Normalize data across groups of channels
%
%   Y = GROUPNORM(X,NUMGROUPS,OFFSET,SCALEFACTOR) normalizes the input data
%   X across groups of channels using the statistics of X, and performs an
%   affine transformation with the specified offset and scale factor.
%   The normalization is applied 'per-observation per-group', the affine
%   transformation is applied 'per-channel'.
%   X must be a formatted dlarray with a 'C' dimension.
%
%   NUMGROUPS can be one of the following:
%        - positive integer       - Divide all incoming channels into number
%                                   of groups specified. Specified number
%                                   of groups must exactly divide the number
%                                   of channels.
%        - "all-channels"         - Gather all channels into one group and
%                                   perform the normalization. Also known
%                                   as "channel-wise layer normalization".
%        - "channel-wise"         - Perform channel-wise normalization.
%                                   Number of groups is equal to the number
%                                   of incoming channels, determined during
%                                   training. Also known as "instance
%                                   normalization".
%
%   If OFFSET and SCALEFACTOR are unformatted dlarrays or numeric arrays,
%   they must be vectors with the same number of elements as the 'C'
%   dimension of input X. If OFFSET and SCALE are formatted dlarrays, they
%   must contain a 'C' dimension with the same size as the 'C' dimension of
%   the input X. Any other dimension must be singleton.
%
%   Output Y is a formatted dlarray with the same dimension labels as X.
%
%   Y = GROUPNORM(__,'DataFormat',FMT) computes the group normalization
%   of an unformatted dlarray or a numeric array X. FMT specifies the
%   dimension labels of the input data X. If the input X is a numeric
%   array, at least one of OFFSET or SCALE must be a dlarray. The output Y
%   is an unformatted dlarray with the same dimension order as X.
%
%   Y = GROUPNORM(__,'PARAM',VAL) specifies optional parameter
%   name/value pairs:
%
%       'Epsilon'           - Offset for the variances, to avoid
%                             divide-by-zero errors, specified as a 
%                             positive scalar. 
%                             The default value is 1e-5.
%       'DataFormat'        - Dimension labels of the input data X,
%                             specified as a char array or string. Required
%                             if X is not a formatted dlarray.
%
%   Example
%       % Create a formatted dlarray with two spatial, one channel, and
%       % one observation dimensions
%       width = 12;
%       height = 12;
%       numChannels = 6;
%       numObservations = 16;
%       X = dlarray(randn(width,height,numChannels,numObservations),'SSCB');
%
%       % Create learnable parameters as unformatted dlarray objects
%       offset = dlarray(zeros(numChannels,1));
%       scaleFactor = dlarray(ones(numChannels,1));
%
%       % Specify number of groups. Normalize over three groups of two
%       % channels each.
%       numGroups = 3;
%
%       % Compute group normalization
%       Z = groupnorm(X,numGroups,offset,scaleFactor);
%
%   See also DLARRAY, BATCHNORM, LAYERNORM, RELU, DLCONV, FULLYCONNECT

%   Copyright 2020-2022 The MathWorks, Inc.

arguments
    X                    {mustBeNumeric, mustBeNonsparse, mustBeReal}
    numGroups
    offset               {mustBeNumeric, mustBeNonsparse, mustBeReal}
    scale                {mustBeNumeric, mustBeNonsparse, mustBeReal}
    nvArgs.Epsilon       {deep.internal.dlarray.mustBeValidEpsilon} = 1e-5  % by convention (including cuDNN)
    nvArgs.DataFormat    {deep.internal.dlarray.mustBeDataFormatType} = ''
end

% Only data defines execution environment, not parameters
numGroups = gather(numGroups);
nvArgs.Epsilon = gather(nvArgs.Epsilon);

% Handle data format
[X,perm] = deep.internal.dlarray.validateDataFormatArg(X, nvArgs.DataFormat);

X = permuteExplicitly(X);

% Extract the input data
fd = X.FormattedData;
X.FormattedData = [];
[fd,xdata] = extractData(fd);

% Retrieve channel index & number of channels
channelDim = channelIndex(fd);
if numel(channelDim) > 1 || channelDim > length(fd.DimensionMetadata)
    error(message('deep:dlarray:MustHaveOneChannelDimension'))
end
numChannels = size(xdata, channelDim);

% Retrieve batch index. If data does not have 'B' dimensions, but has 'T'
% and/or 'U', we need to and insert a singleton 'B' dimension. This is
% necessary because the forward and backward methods need to know the
% relative position of 'C' and 'B' dimension in the data tensor.
batchDim = observationIndex(fd);
needsSingletonBatch = batchDim > ndims(xdata);

if needsSingletonBatch
    origSizexdata = size(xdata);

    % Insert singleton dimension for the batch, after 'C', but before 'T'
    % and/or 'U'
    % Get dimensions until 'C'
    origSizexdataBeforeBatch = size(xdata, 1:channelDim);
    % Get dimensions from 'T' and/or 'U'
    origSizexdataAfterBatch = origSizexdata(channelDim+1:end);
    % Insert singleton dimension
    dimsSingleton = [origSizexdataBeforeBatch, 1, origSizexdataAfterBatch];
    xdata = reshape(xdata,dimsSingleton);
end

% Validate numGroups
numGroups = iAssertValidNumGroups(numGroups, numChannels);

% Extract offset & scale
offset = iExtractLearnableParameter(offset, numChannels, @iOffsetErrorIDs);
scale = iExtractLearnableParameter(scale, numChannels, @iScaleErrorIDs);

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

% Be sure that the learnable parameters have shape compatible with the data
[offset, scale] = ...
    deep.internal.dlarray.reshapeParameterForNormalizationFunction({offset, scale}, ...
    ndims(xdata), channelDim, numChannels);

% Call the internal API
xdata = internal_groupnorm(xdata, offset, scale, numGroups, nvArgs.Epsilon, channelDim);

% Restore the data shape
if needsSingletonBatch
    xdata = reshape(xdata, origSizexdata);
end

% Insert the FormattedData back in the dlarray
fd = insertData(fd, xdata);
X.FormattedData = fd;

% For unformatted, revert the internal permutation vector to its original
% state
if ~isempty(nvArgs.DataFormat)
    X = stripdims(X);
    X = ipermute(X,perm);
end
end

function value = iExtractLearnableParameter(value, numChannels, errorIDFcn)
% This function extract the learnable parameter's data (if wrapped in a
% dlarray)
if isa(value, 'dlarray')
    [fd,value] = extractData(value.FormattedData);
    if ~isempty(fd.DimensionMetadata) % If formatted dlarray
        numChannelsParam = size(value, channelIndex(fd));
        if numChannelsParam ~= numChannels
            error(iGetErrorMessage(errorIDFcn, 'InvalidLabeled', numChannels));
        end
    else % Unformatted dlarray
        numChannelsParam = numel(value);
    end
else
    % The type of the output of the operation could be modified by the
    % computation to a type that is unsupported by dlarray
    deep.internal.dlarray.validateNonDlarray(value);
    numChannelsParam = numel(value);
end
% Parameter must be a vector with as many elements as the number of channels
if numChannelsParam ~= numChannels || ~isvector(value)
    error(iGetErrorMessage(errorIDFcn, 'Invalid', numChannels));
end
end

% Helpers to the throw a more specific error for each input
function msg = iGetErrorMessage(errorIDFcn, IDName, varargin)
s = errorIDFcn();
msg =  message(s.(IDName), varargin{:});
end

function s = iOffsetErrorIDs()
s = struct(...
    'InvalidLabeled', 'deep:dlarray:GroupnormInvalidLabeledOffset', ...
    'Invalid', 'deep:dlarray:GroupnormInvalidOffset');
end

function s = iScaleErrorIDs()
s = struct(...
    'InvalidLabeled', 'deep:dlarray:GroupnormInvalidLabeledScale', ...
    'Invalid', 'deep:dlarray:GroupnormInvalidScale');
end

function value = iAssertValidNumGroups(value, numChannels)
if iIsValidStringOrCharArray(value)
    validatedString = validatestring(value,[iChannelWiseString, iAllChannelsString]);
    if startsWith(validatedString,'a')
        % Assume the string is "all-channels"
        value = 1;
    else
        % Assume the string is "channel-wise"
        value = numChannels;
    end
else
    iAssertScalarPositiveInteger(value)
    
    numChannelsPerGroup = numChannels./value;
    % Error if numGroups is greater than numChannels or if numGroups is not
    % a divisor of numChannels
    if value > numChannels || ~iIsInteger(numChannelsPerGroup)
        error(message('deep:dlarray:GroupnormIncompatibleNumGroups', numChannels, value ));
    end
end
end

function str = iChannelWiseString()
str = "channel-wise";
end

function str = iAllChannelsString()
str = "all-channels";
end

function iAssertScalarPositiveInteger(value)
deep.internal.dlarray.mustBeNonDlarray(value);
validateattributes(value, {'numeric'}, {'scalar','integer','positive'});
end

function tf = iIsInteger(x)
tf = ~mod(x,1);
end

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