function [X, mean, variance] = batchnorm(X, offset, scale, mean, variance, nvArgs)
%BATCHNORM   Normalize each channel of input data
%
%   [Y,MEAN,VARIANCE] = BATCHNORM(X,OFFSET,SCALE) normalizes each channel
%   of the input data X, using the statistics of X and learnable parameters
%   OFFSET and SCALE. X must be a formatted dlarray with a 'C' dimension.
%
%   If OFFSET and SCALE 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.
%   MEAN and VARIANCE are the per-channel mean and variance of X,
%   respectively. These values are returned as column numeric vectors whose
%   length matches the size of the 'C' dimension of the input X.
%
%   Y = BATCHNORM(X,OFFSET,SCALE,MEAN,VARIANCE) uses MEAN and VARIANCE
%   as statistics for the normalization. MEAN and VARIANCE must be numeric
%   vectors whose lengths match the size of the 'C' dimension of X.
%
%   [Y,DATASETMEAN,DATASETVARIANCE] = BATCHNORM(X,OFFSET,SCALE,DATASETMEAN,
%   DATASETVARIANCE) normalizes each channel of the input data X using the
%   statistics of X and learnable parameters OFFSET and SCALE. Inputs
%   DATASETMEAN and DATASETVAR are the values of mean and variance
%   calculated across multiple mini-batches, respectively.
%   Use this syntax to iteratively update the mean and variance statistics
%   over several batches of data. Outputs DATASETMEAN and DATASETVAR are
%   computed from the statistics of X, DATASETMEAN, and DATASETVAR using a
%   exponentially weighted moving average filter of the form
%   DATASETSTATISTIC = DECAY*CURRENT + (1-DECAY)*DATASETSTATISTIC, where
%   CURRENT is the statistic computed from input X, and parameter DECAY is
%   the decay. The default value of DECAY is 0.1.
%
%   Y = BATCHNORM(__,'DataFormat',FMT) computes the batch 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 = BATCHNORM(__,'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.
%       'MeanDecay'         - Decay value to use for moving average
%                             computation of the mean statistic. The
%                             default value is 0.1.
%       'VarianceDecay'     - Decay value to use for moving average
%                             computation of the variance statistic. The
%                             default value is 0.1.
%       '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 a two spatial, one channel, and
%       % one observation dimension
%       width = 12;
%       height = 12;
%       numChannels = 3;
%       numObservations = 16;
%       X = dlarray(randn(width,height,numChannels,numObservations),'SSCB');
%
%       % Create learnable parameters as unformatted dlarrays
%       offset = dlarray(zeros(numChannels,1));
%       scale = dlarray(ones(numChannels,1));
%
%       % Compute batch normalization and extract the statistics of the
%       % batch
%       [Z,currentMean,currentVariance] = batchnorm(X,offset,scale);
%
%   See also DLARRAY, RELU, DLCONV, FULLYCONNECT

%   Internal API for tracing mean and variance:
%   [Y,DATASETMEAN,DATASETVARIANCE] = BATCHNORM(X,OFFSET,SCALE,DATASETMEAN,
%   DATASETVARIANCE) allows dlarray DATASETMEAN and DATASETVARIANCE. An input
%   mean and variance of dlarray([]) will output dlarray batch statistics.
%   Otherwise dlarray input statistics will output dlarray statistics.
%   Input dlarray statistics must be unformatted.
    
%   Copyright 2019-2022 The MathWorks, Inc.

arguments
    X {mustBeNumeric, mustBeNonsparse, mustBeReal}
    offset {mustBeNumeric, mustBeNonsparse, mustBeReal, mustBeNonempty}
    scale {mustBeNumeric, mustBeNonsparse, mustBeReal, mustBeNonempty}
    mean {mustBeNumeric, mustBeNonsparse, mustBeReal} = []
    variance {mustBeNumeric, mustBeNonsparse, mustBeReal} = []
    nvArgs.DataFormat {deep.internal.dlarray.mustBeDataFormatType} = ''
    nvArgs.MeanDecay (1,1) {deep.internal.dlarray.mustBeNonDlarray, ...
        mustBeNumeric, mustBeFinite, mustBeNonnegative, ...
        mustBeLessThanOrEqual(nvArgs.MeanDecay,1)} = 0.1
    nvArgs.VarianceDecay (1,1) {deep.internal.dlarray.mustBeNonDlarray, ...
        mustBeNumeric, mustBeFinite, mustBeNonnegative, ...
        mustBeLessThanOrEqual(nvArgs.VarianceDecay,1)} = 0.1
    nvArgs.Epsilon {deep.internal.dlarray.mustBeValidEpsilon} = 1e-5 % by convention (including cuDNN)
end


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

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

labelsX = dims(X);
% Locate channel dimension and find number of channels of X
channelDim = find(labelsX=='C');
if ~isempty(channelDim)
    numChannels = size(X, channelDim);
else
    error(message('deep:dlarray:MustHaveOneChannelDimension'))
end

iValidateLearnableParameter(offset, numChannels, @iOffsetErrorIDs);
iValidateLearnableParameter(scale, numChannels, @iScaleErrorIDs);
iValidateAllStatisticParameters(mean, variance, numChannels);

doInference = false;
if nargout <= 1
    doInference = true;
end

[X,mean,variance] = deep.internal.dlarray.batchNormalizeData(X, offset, scale,...
    mean, variance, nvArgs.Epsilon, nvArgs.MeanDecay, nvArgs.VarianceDecay, ...
    doInference);

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

% Validation helper functions
function iValidateLearnableParameter(value, numChannels, errorIDFcn)
% This function validates the learnable parameter's data

if isdlarray(value)
    labels = dims(value);
    if ~isempty(labels) % If formatted dlarray
        channelDimParam = find(labels=='C');
        if isempty(channelDimParam) || size(value, channelDimParam) ~= numChannels
            error(iGetErrorMessage(errorIDFcn, 'InvalidLabeled', numChannels));
        else
            nbChannelsParam = size(value, channelDimParam);
        end
    else % Unformatted dlarray
        nbChannelsParam = numel(value);
    end
else
    deep.internal.dlarray.validateNonDlarray(value);
    nbChannelsParam = numel(value);
end
% Parameter must be a vector with as many elements as the number of channels
if nbChannelsParam ~= numChannels || ~isvector(value)
    error(iGetErrorMessage(errorIDFcn, 'Invalid', numChannels));
end
end

function iValidateAllStatisticParameters(mean, variance, numChannels)
% Validate the trained mean and variance parameters. This return either
% two empty values or two values with numChannels elements, formatted for
% the internal API.
iValidateStatisticParameter(mean, numChannels, @iMeanErrorIDs);
iValidateStatisticParameter(variance, numChannels, @iVarianceErrorIDs);
% Error if the emptiness of the two parameters does not match
meanAndVarAreEmpty = isempty(mean);
if meanAndVarAreEmpty ~= isempty(variance)
    error(message('deep:dlarray:BatchnormInvalidStatistic'))
end
end

function iValidateStatisticParameter(stat, numChannels, errorIDFcn)
% Validate the mean or variance, and reshape it for the builtin
mustBeMember(underlyingType(stat),iStatisticDataTypes)
if ~isempty(stat)
    if numel(stat) ~= numChannels || ~isvector(stat)
        error(iGetErrorMessage(errorIDFcn, 'Invalid', numChannels));
    end
end
if isdlarray(stat) && ~isempty(dims(stat))
    error(iGetErrorMessage(errorIDFcn, 'InvalidLabeled'));
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:BatchnormInvalidLabeledOffset', ...
    'Invalid', 'deep:dlarray:BatchnormInvalidOffset');
end

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

function s = iMeanErrorIDs()
s = struct(...
    'Invalid', 'deep:dlarray:BatchnormInvalidMean',...
    'InvalidLabeled', 'deep:dlarray:BatchnormInvalidFormatMean');
end

function s = iVarianceErrorIDs()
s = struct(...
    'Invalid', 'deep:dlarray:BatchnormInvalidVariance',...
    'InvalidLabeled', 'deep:dlarray:BatchnormInvalidFormatVariance');
end

function c = iStatisticDataTypes
c = {'single', 'double'};
end
