function layer = batchNormalizationLayer(varargin)
%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:
%                                     - []
%                                     - 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.
%       'TrainedMean'             - Per-channel mean of the layer input
%                                   data. 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 [].
%       'TrainedVariance'         - Per-channel variance of the layer input
%                                   data. 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 [].
%
%   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-2018 The MathWorks, Inc.

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

% Create an internal representation of the layer.
internalLayer = nnet.internal.cnn.layer.BatchNormalization(args.Name, [], 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.BatchNormalizationLayer(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', 'TrainedMean', 'TrainedVariance', ...
     '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 = [];
defaultTrainedMean = [];
defaultTrainedVariance = [];
defaultLearnRateFactor = 1;
defaultL2Factor = 1;
defaultName = '';
defaultOffsetInitializer = 'zeros';
defaultScaleInitializer = 'ones';

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('TrainedMean', defaultTrainedMean);
p.addParameter('TrainedVariance', defaultTrainedVariance);
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);
end

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

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

function iAssertValidFactor(value)
nnet.internal.cnn.layer.paramvalidation.validateLearnFactor(value);
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