function layer = groupNormalizationLayer(varargin)
% groupNormalizationLayer   Group normalization layer
%
%   layer = groupNormalizationLayer(numGroups) creates a
%   group normalization layer. This type of layer divides the channels into
%   numGroups groups, normalizes across each group and applies an affine
%   transformation with learnable parameters.
%   Allowed values for numGroups are:
%        - 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 "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"
%
%   layer = groupNormalizationLayer(numGroups, 'PARAM1', VAL1, ...)
%   specifies optional parameter name/value pairs for creating the layer:
%
%       'Name'                    - A name for the layer. Default is ''.
%       'Offset'                  - Initial value for learnable parameter
%                                   Offset (also called beta). This can be:
%                                     - []
%                                     - a 1-by-1-by-C array (for 2-D images)
%                                     - a 1-by-1-by-1-by-C array (for 3-D
%                                       images)
%                                     - a C-by-1 array (for sequence and
%                                       feature vectors)
%                                   where C is the number of channels.
%                                   The default is [].
%       'Scale'                   - Initial value for learnable parameter
%                                   Scale (also called gamma). This can be:
%                                     - []
%                                     - a 1-by-1-by-C array (for 2-D images)
%                                     - a 1-by-1-by-1-by-C array (for 3-D
%                                       images)
%                                     - a C-by-1 array (for sequence and
%                                       feature vectors)
%                                   where C is the number of channels.
%                                   The default is [].
%       'Epsilon'                 - Offset for the variance to avoid
%                                   divide-by-zero errors. Must be at least
%                                   1e-5. Default is 1e-5.
%       'OffsetInitializer'       - Function to initialize the Offset,
%                                   specified as 'narrow-normal', 'zeros',
%                                   'ones' or a function handle. The
%                                   default is 'zeros'.
%       'ScaleInitializer'        - Function to initialize the Scale,
%                                   specified as 'narrow-normal', 'ones' or
%                                   a function handle. The default is
%                                   'ones'.
%       'OffsetLearnRateFactor'   - Multiplier for the learning rate of
%                                   Offset. Default is 1.
%       'ScaleLearnRateFactor'    - Multiplier for the learning rate of
%                                   Scale. Default is 1.
%       'OffsetL2Factor'          - Multiplier for the L2 weight regulariser
%                                   for Offset. Default is 1.
%       'ScaleL2Factor'           - Multiplier for the L2 weight regulariser
%                                   for Scale. Default is 1.
%
%   Example:
%     % Create a group normalization layer that divides the channels
%     % into 32 groups.
%
%     layer = groupNormalizationLayer(32);
%
%   See also nnet.cnn.layer.groupNormalizationLayer.
%
%   <a href="matlab:helpview('deeplearning','list_of_layers')">List of Deep Learning Layers</a>

%   Copyright 2019-2020 The MathWorks, Inc.

% Parse the input arguments.
args = iParseInputArguments(varargin{:});

% Create an internal representation of the layer.
internalLayer = nnet.internal.cnn.layer.GroupNormalization(args.Name, [], args.NumGroups, args.Epsilon);

internalLayer.Offset.L2Factor = args.OffsetL2Factor;
internalLayer.Offset.LearnRateFactor = args.OffsetLearnRateFactor;

internalLayer.Scale.L2Factor = args.ScaleL2Factor;
internalLayer.Scale.LearnRateFactor = args.ScaleLearnRateFactor;

% Use the internal layer to construct a user visible layer.
layer = nnet.cnn.layer.GroupNormalizationLayer(internalLayer);

% Set these properties (if supplied) using the visible layer so that the
% number of channels is correctly checked.
layer = iSetPropertyIfNonEmpty(layer, args, ...
    {'Offset', 'Scale' ...
    'OffsetInitializer', 'ScaleInitializer'});
end

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

function p = iCreateParser()
p = inputParser;

defaultEpsilon = 1e-5;
defaultOffset = [];
defaultScale = [];
defaultLearnRateFactor = 1;
defaultL2Factor = 1;
defaultName = '';
defaultOffsetInitializer = 'zeros';
defaultScaleInitializer = 'ones';

p.addRequired('NumGroups', @iAssertValidNumGroups);
p.addParameter('Name', defaultName, @iAssertValidLayerName);
p.addParameter('Epsilon', defaultEpsilon);
p.addParameter('Offset', defaultOffset);
p.addParameter('Scale', defaultScale);
p.addParameter('OffsetLearnRateFactor', defaultLearnRateFactor, @iAssertValidFactor);
p.addParameter('ScaleLearnRateFactor', defaultLearnRateFactor, @iAssertValidFactor);
p.addParameter('OffsetL2Factor', defaultL2Factor, @iAssertValidFactor);
p.addParameter('ScaleL2Factor', defaultL2Factor, @iAssertValidFactor);
p.addParameter('OffsetInitializer', defaultOffsetInitializer);
p.addParameter('ScaleInitializer', defaultScaleInitializer);
end

function inputArguments = iConvertToCanonicalForm(params)
inputArguments = params;
inputArguments.Epsilon = iValidateAndConvertEpsilon(inputArguments.Epsilon);
inputArguments.Name = convertStringsToChars(inputArguments.Name);
inputArguments.NumGroups = double( iComputeInternalNumGroups(params.NumGroups) );
end

function iAssertValidLayerName(name)
nnet.internal.cnn.layer.paramvalidation.validateLayerName(name)
end

function epsilon = iValidateAndConvertEpsilon(value)
minValue = 1e-5;
epsilon = nnet.internal.cnn.layer.paramvalidation...
    .validateAndConvertEpsilon(value, minValue);
end

function iAssertValidFactor(value)
nnet.internal.cnn.layer.paramvalidation.validateLearnFactor(value);
end

function iAssertValidNumGroups(value)
if iIsValidStringOrCharArray(value)
    validatestring(value,[iChannelWiseString, iAllChannelsString]);
else
    iAssertScalarPositiveInteger(value)
end
end

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

function layer = iSetPropertyIfNonEmpty(layer, args, propertyNames)
for i=1:numel(propertyNames)
    name = propertyNames{i};
    if ~isempty(args.(name))
        layer.(name) = args.(name);
    end
end
end

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

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

function y = iComputeInternalNumGroups(x)
% The NumGroups value has previously been validated. It must either be a
% string or a scalar integer. An empty value for NumGroups indicates to the
% internal layer that the size still needs to be inferred.
if iIsValidStringOrCharArray(x)
    if startsWith(x,'a')
        y = 1;
    else
        % Assume the string is "channel-wise"
        y = [];
    end
else
    y = x;
end
end


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