function layer = batchNormalizationLayer(nameValueArgs)
%batchNormalizationLayer   Batch normalization layer
%
%   layer = batchNormalizationLayer() creates a batch normalization layer.
%   This type of layer normalizes each channel across a mini-batch. This
%   can be useful in reducing sensitivity to variations within the data.
%
%   layer = batchNormalizationLayer('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 
%                                   [] or a vector with one nonsingleton 
%                                   dimension of size C, where C is the 
%                                   number of channels. The default is [].
%       'Scale'                   - Initial value for learnable parameter
%                                   Scale (also called gamma). This can be 
%                                   [] or a vector with one nonsingleton 
%                                   dimension of size C, where C is the 
%                                   number of channels. The default is [].
%       'Epsilon'                 - Offset for the variance to avoid
%                                   divide-by-zero errors, specified as a
%                                   positive scalar.
%                                   The 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.
%       'MeanDecay'               - Decay value to use for moving average
%                                   computation of the mean statistic.
%                                   Default value is 0.1.
%       'VarianceDecay'           - Decay value to use for moving average
%                                   computation of the variance statistic.
%                                   Default value is 0.1.
%       'TrainedMean'             - Per-channel mean of the layer input
%                                   data. This can be 
%                                   [] or a vector with one nonsingleton 
%                                   dimension of size C, where C is the 
%                                   number of channels. The default is [].
%       'TrainedVariance'         - Per-channel variance of the layer input
%                                   data. This can be 
%                                   [] or a vector with one nonsingleton 
%                                   dimension of size C, where C is the 
%                                   number of channels. The default is [].
%
%   Example:
%     % Create a batch normalization layer.
%
%     layer = batchNormalizationLayer();
%
%   See also nnet.cnn.layer.BatchNormalizationLayer.
%
%   <a href="matlab:helpview('deeplearning','list_of_layers')">List of Deep Learning Layers</a>

%   Copyright 2017-2023 The MathWorks, Inc.

arguments
    nameValueArgs.Name                   {iAssertValidLayerName} = ""
    nameValueArgs.Epsilon                (1,1) = 1e-5
    nameValueArgs.Offset                 = []
    nameValueArgs.Scale                  = []
    nameValueArgs.TrainedMean            = []
    nameValueArgs.TrainedVariance        = []
    nameValueArgs.OffsetLearnRateFactor  (1,1) {iAssertValidFactor} = 1
    nameValueArgs.ScaleLearnRateFactor   (1,1) {iAssertValidFactor} = 1
    nameValueArgs.OffsetL2Factor         (1,1) {iAssertValidFactor} = 1
    nameValueArgs.ScaleL2Factor          (1,1) {iAssertValidFactor} = 1
    nameValueArgs.OffsetInitializer      = 'zeros'
    nameValueArgs.ScaleInitializer       = 'ones'
    nameValueArgs.MeanDecay              (1,1) {mustBeNumeric, mustBeNonnegative, ...
                                         mustBeLessThanOrEqual(nameValueArgs.MeanDecay,1)} = 0.1
    nameValueArgs.VarianceDecay          (1,1) {mustBeNumeric, mustBeNonnegative, ...
                                         mustBeLessThanOrEqual(nameValueArgs.VarianceDecay,1)} = 0.1
end

nameValueArgs.Epsilon = iValidateAndConvertEpsilon(nameValueArgs.Epsilon);
nameValueArgs.Name = convertStringsToChars(nameValueArgs.Name);

nameValueArgs = nnet.internal.cnn.layer.util.gatherParametersToCPU(nameValueArgs);

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

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

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

iValidateDecayFactorAndTrainedStatistic(nameValueArgs.MeanDecay, ...
    nameValueArgs.TrainedMean, iGetInvalidMeanDecayErrorID)
iValidateDecayFactorAndTrainedStatistic(nameValueArgs.VarianceDecay, ...
    nameValueArgs.TrainedVariance, iGetInvalidVarianceDecayErrorID)

internalLayer.MeanDecay = iConvertToFloat(nameValueArgs.MeanDecay);
internalLayer.VarianceDecay = iConvertToFloat(nameValueArgs.VarianceDecay);

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

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

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

function epsilon = iValidateAndConvertEpsilon(value)
epsilon = nnet.internal.cnn.layer.paramvalidation.validateAndConvertEpsilon(value);
end

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

function iValidateDecayFactorAndTrainedStatistic(valueDecay,valueTrainedStat,errorID)
if isequal(valueDecay,0) && isempty(valueTrainedStat)
    error(message(errorID))
end
end

function errorID = iGetInvalidMeanDecayErrorID()
errorID = 'nnet_cnn:layer:BatchNormalizationLayer:InvalidMeanDecay';
end

function errorID = iGetInvalidVarianceDecayErrorID()
errorID = 'nnet_cnn:layer:BatchNormalizationLayer:InvalidVarianceDecay';
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 value = iConvertToFloat(value)
value = nnet.internal.cnn.layer.util.convertToFloat(value);
end
