function [X, mean, invSqrtVar, runningMean, runningVariance] = ...
    internal_batchnorm(X, offset, scale, epsilon, channelDim, mode, trainedMean, trainedVariance, decayMean, decayVar)
% This function is undocumented and will change in a future release.

%   [Y,mean,invSqrtVar,runningMean,runningVariance] = internal_batchnorm(X,offset,scale,epsilon,trainedMean,trainedVariance,dim)
%   Compute the batch normalization in training or inference modes.
%
%   Inputs:
%     X               Batch data.
%                     numeric array
%     offset          Offset. Also called beta. Shape is only non-trivial in channel dimension.
%                     numeric array 
%     scale           Scale. Also called gamma. Shape is only non-trivial in channel dimension.
%                     numeric array 
%     epsilon         Constant to avoid division by zero.
%                     numeric scalar
%     channelDim      Channel dimension.
%                     double
%     mode            Mode of computation for batchnorm, specified as a string. 
%                     Possible values: "inference", "training", "trainingWithMovingAverage"
%     trainedMean     Trained mean, non empty if in inference mode or moving average computation mode.
%                     Input running average if in training with moving average mode.
%                     Shape is only non-trivial in channel dimension.
%                     numeric array | []
%     trainedVariance Trained variance, non empty if in inference mode or moving average computation mode. 
%                     Input running average if in training with moving average mode.
%                     Shape is only non-trivial in channel dimension.
%                     numeric array  | []
%     decayMean       Decay for exponential moving average computation of
%                     the mean. Non-empty for moving average computation.
%                     double | []
%     decayVar        Decay for exponential moving average computation of
%                     the variance. Non-empty for moving average computation.
%                     double | []
%  
%   Outputs:
%     Y               Normalized batch. Shape is only non-trivial in channel dimension.
%                     numeric array
%     mean            Batch mean, if in training mode. Shape is only non-trivial in channel dimension.
%                     In inference mode, mean is the trainedMean.
%                     numeric array | []
%     invSqrtVar      Batch inverse-sqrt-variance, if in training mode. Shape is only non-trivial in channel dimension.
%                     In inference mode, invSqrtVar is defined as 1./sqrt(trainedVariance + epsilon).
%                     numeric array | []
%     runningMean     Running average of the mean updated, if in training with moving average computation mode. 
%                     Shape is only non-trivial in channel dimension.
%                     numeric array | []
%     runningVariance Running average of the variance updated, if in training with moving average computation mode. 
%                     Shape is only non-trivial in channel dimension.
%                     numeric array | []

%   Copyright 2019-2021 The MathWorks, Inc.

if ~isreal(X) || ~isreal(offset) || ~isreal(scale)
    error(message('deep:dlarray:ComplexNotSupported'))
end

runningMean = [];
runningVariance = [];
switch mode
    case "inference" % Inference mode
        if ~isreal(trainedMean) || ~isreal(trainedVariance)
            error(message('deep:dlarray:ComplexNotSupported'))
        end
        trainedVariance = max(trainedVariance,0);
        X = nnet.internal.cnnhost.batchNormalizationForwardPredict(X,offset,scale,epsilon,trainedMean,trainedVariance,channelDim);
        mean = trainedMean;
        invSqrtVar = 1./sqrt(trainedVariance + epsilon);
    case "training" % Training mode
        [X,mean,invSqrtVar] = nnet.internal.cnnhost.batchNormalizationForwardTrain(X,offset,scale,epsilon,channelDim);
    case "trainingWithMovingAverage" % Training mode with moving average computation
        if ~isreal(trainedMean) || ~isreal(trainedVariance)
            error(message('deep:dlarray:ComplexNotSupported'))
        end
        [X,mean,invSqrtVar] = nnet.internal.cnnhost.batchNormalizationForwardTrain(X,offset,scale,epsilon,channelDim);
        runningMean = nnet.internal.cnn.statistics.exponentialMovingAverage(mean,trainedMean,decayMean);
        variance = 1./(invSqrtVar).^2-epsilon;
        runningVariance = nnet.internal.cnn.statistics.exponentialMovingAverage(variance,trainedVariance,decayVar);
        runningVariance = nnet.internal.cnn.util.boundAwayFromZero(runningVariance);
end

end

