function X = layernorm(X, offset, scale, nvArgs)
%LAYERNORM   Normalize data across channel dimension
%
%   Y = LAYERNORM(X,OFFSET,SCALEFACTOR) normalizes the input data X
%   across all channels and performs an affine transformation with the
%   specified offset and scale factor. The function normalizes
%   per-observation and, according to the dimensions of OFFSET and
%   SCALEFACTOR, applies an affine transformation:
%       - channel-wise, if OFFSET and SCALEFACTOR are vectors
%       - general element-wise, if OFFSET and SCALEFACTOR are dlarray
%         objects or numeric arrays with size and format compatible with
%         input data X.
%
%   X must be a formatted dlarray with a 'C' dimension.
%
%   If OFFSET and SCALEFACTOR are numeric arrays, then they must be vectors
%   with the same number of elements as the 'C' dimension of the input X.
%
%   If OFFSET and SCALEFACTOR are unformatted dlarray objects, then they
%   must be accompanied by the 'OffsetFormat' and 'ScaleFormat' options.
%   The options must have a 'C' dimension with the same size as the
%   'C' dimension of the input X and they can contain as many as 'S', 'T'
%   and 'U' dimensions as the input data X. For each type of dimension
%   label, the dimensions with that label must all be singleton or must
%   all have size matching the corresponding dimension in the input data.
%
%   If OFFSET and SCALEFACTOR are formatted dlarray objects, then they must
%   have a 'C' dimension with the same size as the 'C' dimension of the
%   input X. They can contain as many as 'S', 'T' and 'U' dimensions as the
%   input data X. For each type of dimension label, the dimensions with
%   that label must all be singleton or must all have size matching the
%   corresponding dimension in the input data.
%
%   Output Y is a formatted dlarray with the same dimension labels as X.
%
%   Y = LAYERNORM(__,'DataFormat',FMT) computes the layer 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 SCALEFACTOR must be a dlarray. The
%   output Y is an unformatted dlarray with the same dimension order as X.
%
%   Y = LAYERNORM(__,'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.
%
%     'OperationDimension'  - Dimension to normalize over, specified as one
%                             of these values: 
%                             "auto"            - If X is feature,
%                                                 sequence, 1-D image, or
%                                                 spatial-temporal data,
%                                                 then normalize X over the
%                                                 channel dimension.
%                                                 Otherwise, normalize X
%                                                 over the spatial and
%                                                 channel dimensions.
%
%                             "channel-only"    - Normalize X over the
%                                                 channel dimension.
%
%                             "spatial-channel" - Normalize X over the
%                                                 spatial and channel
%                                                 dimensions.
%
%                             "batch-excluded"  - Normalize X over all
%                                                 dimensions excluding the
%                                                 batch dimension.
%                             The default value is "auto".
%
%     'DataFormat'          - Dimension labels of the input data X,
%                             specified as a char array or string. Required
%                             if X is not a formatted dlarray.
%     'OffsetFormat'        - Dimension labels of the offset, specified as
%                             a character vector or string. Required if
%                             offset is unformatted.
%     'ScaleFormat'         - Dimension labels of the scale factor,
%                             specified as a character vector or string.
%                             Required if scaleFactor is unformatted.
%
%   Example 1:
%       % Create a formatted dlarray with two spatial, one channel, and
%       % one observation dimensions
%       width = 12;
%       height = 12;
%       numChannels = 6;
%       numObservations = 16;
%       X = dlarray(randn(width,height,numChannels,numObservations),'SSCB');
%
%       % Create learnable parameters as unformatted dlarray objects
%       offset = dlarray(zeros(numChannels,1));
%       scaleFactor = dlarray(ones(numChannels,1));
%
%       % Compute layer normalization
%       Z = layernorm(X,offset,scaleFactor);
%
%   Example 2:
%       % Create a formatted dlarray with two spatial, one channel, and
%       % one observation dimensions
%       width = 12;
%       height = 12;
%       numChannels = 6;
%       numObservations = 16;
%       X = dlarray(randn(width,height,numChannels,numObservations),'SSCB');
%
%       % Create learnable parameters as unformatted dlarray objects
%       offset = dlarray(randn(width,height,numChannels));
%       scaleFactor = dlarray(randn(width,height,numChannels));
%
%       % Compute layer normalization with spatial affine transformation
%       Z = layernorm(X,offset,scaleFactor,'OffsetFormat','SSC','ScaleFormat','SSC');
%
%   See also DLARRAY, BATCHNORM, GROUPNORM, RELU, DLCONV, FULLYCONNECT

%   Copyright 2020-2023 The MathWorks, Inc.

arguments
    X                         {mustBeNumeric, mustBeNonsparse, mustBeReal}
    offset                    {mustBeNumeric, mustBeNonsparse, mustBeReal, mustBeNonempty}
    scale                     {mustBeNumeric, mustBeNonsparse, mustBeReal, mustBeNonempty}
    nvArgs.Epsilon            {deep.internal.dlarray.mustBeValidEpsilon} = 1e-5
    nvArgs.OperationDimension {mustBeMember(nvArgs.OperationDimension, ["auto", "channel-only", "spatial-channel", "batch-excluded"])} = "auto"
    nvArgs.DataFormat         {deep.internal.dlarray.mustBeDataFormatType} = ''
    nvArgs.ScaleFormat        {deep.internal.dlarray.mustBeDataFormatType} = ''
    nvArgs.OffsetFormat       {deep.internal.dlarray.mustBeDataFormatType} = ''
end

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

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

X = permuteExplicitly(X);

% Create placeholderArray for X for validation purposes
placeholderX = deep.internal.PlaceholderArray(X);

% Extract the input data
fd = X.FormattedData;
X.FormattedData = [];
[fd,xdata] = extractData(fd);

% Retrieve channel index & number of channels
channelDim = channelIndex(fd);

if numel(channelDim) > 1 || channelDim > length(fd.DimensionMetadata)
    error(message('deep:dlarray:MustHaveOneChannelDimension'))
end
numChannels = size(xdata, channelDim);

% Obtain numeric equivalent of the normalization dimension
normalizationDim = iConvertNormalizationDimToNumeric(nvArgs.OperationDimension, fd, channelDim);

% Validate offset and scale with respect to each other and with respect to
% the input data
[offset, scale, spatialSizeParam, temporalSizeParam, unspecifiedSizeParam] = ...
    iValidateLearnablesAndData(offset, scale, nvArgs.OffsetFormat, nvArgs.ScaleFormat, numChannels, ...
    placeholderX);

% Extract offset & scale
offset = iExtractLearnableParameters(offset);
scale = iExtractLearnableParameters(scale);

% Get observation and time dim of data
batchDim = observationIndex(fd);
hasObservationDim = batchDim <= ndims(xdata);

timeDim = indexOfLabels(fd, 'T');
hasTimeDim = ~isempty(timeDim);

% Be sure that the learnable parameters have shape compatible with the data
[offset, scale] = iReshapeParameters({offset, scale}, ndims(xdata), channelDim, numChannels, ...
    spatialSizeParam, temporalSizeParam, unspecifiedSizeParam, hasObservationDim, hasTimeDim);

% Call the internal API
xdata = internal_layernorm(xdata, offset, scale, nvArgs.Epsilon, normalizationDim, channelDim);

% Insert the FormattedData back into the dlarray
fd = insertData(fd, xdata);
X.FormattedData = fd;

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

function [offset, scale] = iReshapeParameters(cellParam, ndimsInput, channelDim, numChannels, spatialSizeParam, temporalSizeParam, unspecifiedSizeParam, hasBDimInput, hasTDimInput)
% Reshape the parameters

% Get the reshape vector
reshapeVec = ones(1, max(ndimsInput,channelDim));
if ~isempty(spatialSizeParam)
    reshapeVec(1:channelDim-1) = spatialSizeParam;
end
if ~isempty(temporalSizeParam)
    reshapeVec(channelDim+hasBDimInput+1) = temporalSizeParam;
end
if ~isempty(unspecifiedSizeParam)
    reshapeVec(channelDim+hasBDimInput+hasTDimInput+1:ndimsInput) = unspecifiedSizeParam;
end
reshapeVec(channelDim) = numChannels;

% Reshape the learnable parameters
offset = reshape(cellParam{1}, reshapeVec);
scale = reshape(cellParam{2}, reshapeVec);
end

function value = iExtractLearnableParameters(value)
% Extract the learnable parameter's data (if wrapped in a
% dlarray)
value = deep.internal.dlarray.toNumeric(value);
end

function iValidateChannelDimLearnablesAndData(learnable,numChannels)
% Validate with respect to numChannels in data
if isvector(learnable)
    % If learnarable is a vector, it must have the dimension of channels.
    % This is true for numeric vector or dlarray vector ('C')
    channelSizeInLearnable = numel(learnable);
else
    % Learnable can be a tensor only if it is a labelled dlarray object.
    channelSizeInLearnable = size(learnable,finddim(learnable,'C'));
end
if numChannels ~= channelSizeInLearnable
    error(message('deep:dlarray:LayernormInconsistentChannelDimLearnablesAndData'));
end
end

function iValidateSizeWeightsAndData(spatialSizeWeights,temporalSizeWeights,unspecifiedSizeWeights,spatialSizeX,temporalSizeX,unspecifiedSizeX)
% This validation assumes that the learnable parameter is not a vector. In
% particular, this validation checks that the spatial, temporal and
% unspecified dimensions of the learnable parameter are compatible with the
% ones in the data.

% Validate spatial dimension
validSpatialParamDataOnes = ones(size(spatialSizeWeights));
validSpatialParamDataAsX = spatialSizeX;

isValidSpatialSize = isequal(spatialSizeWeights,validSpatialParamDataOnes) || ...
    isequal(spatialSizeWeights,validSpatialParamDataAsX);

isValidTemporalSize = isempty(temporalSizeWeights) || temporalSizeWeights==1 || ...
    temporalSizeWeights==temporalSizeX;

% Validate unspecified dimension
isValidUnspecifiedSize = isequal(unspecifiedSizeWeights,ones(size(unspecifiedSizeWeights))) || ...
    isequal(unspecifiedSizeWeights,unspecifiedSizeX);

% Throw an error for invalid dimension size    
if ~isValidSpatialSize
    error(message('deep:dlarray:LayernormInconsistentLearnableSizeAndDataSize','S'))
elseif ~isValidTemporalSize
    error(message('deep:dlarray:LayernormInconsistentLearnableSizeAndDataSize','T'))
elseif ~isValidUnspecifiedSize
    error(message('deep:dlarray:LayernormInconsistentLearnableSizeAndDataSize','U'))
end
end

function tf = iAreVectorsWithIncompatibleSizes (param1, param2)
tf = isvector(param1) && isvector(param2) && numel(param1) ~= numel(param2);
end

function [offset, scale, spatialSizeParam, temporalSizeParam, unspecifiedSizeParam] = iValidateLearnablesAndData(offset, scale, offsetFormat, scaleFormat, numChannels, placeholderX)

spatialSizeX = placeholderX.getSizeForDims('S');
temporalSizeX = placeholderX.getSizeForDims('T');
unspecifiedSizeX = placeholderX.getSizeForDims('U');
labelsX = dims(placeholderX);

[offset, spatialSizeOffset, temporalSizeOffset, unspecifiedSizeOffset] = ...
    deep.internal.dlarray.validateNormalizationLearnables(offset, offsetFormat, labelsX, "Offset");
[scale, spatialSizeScale, temporalSizeScale, unspecifiedSizeScale] = ...
    deep.internal.dlarray.validateNormalizationLearnables(scale, scaleFormat, labelsX, "Scale");

% Validate labels of offset and scale. They must be the same
if isdlarray(offset) && isdlarray(scale)
    labelsOffset = dims(offset);
    labelsScale = dims(scale);
    if ~isequal(labelsOffset,labelsScale)
        error(message('deep:dlarray:LayernormInconsistentLearnablesLabels'));
    end
end

% Validate size of offset and scale. They must be the same
if iAreVectorsWithIncompatibleSizes (offset, scale) || ...
        ~isequal(spatialSizeOffset, spatialSizeScale) || ...
        ~isequal(temporalSizeOffset, temporalSizeScale) || ...
        ~isequal(unspecifiedSizeOffset, unspecifiedSizeScale)
    error(message('deep:dlarray:LayernormInconsistentLearnablesSize'));
end

% Now we are sure that spatial, time, and unspecified sizes are the same in
% offset and scale, we generalize the names to improve readability
spatialSizeParam = spatialSizeOffset;
temporalSizeParam = temporalSizeOffset;
unspecifiedSizeParam = unspecifiedSizeOffset;

% Validate number of channels of learnables with respect to data
iValidateChannelDimLearnablesAndData(offset,numChannels);
iValidateChannelDimLearnablesAndData(scale,numChannels);

% Validate size of offset and scale with respect to data if the learnables
% were not specified as vectors. We already validated their size, so we
% just check one of the two.
if ~isvector(offset)
    iValidateSizeWeightsAndData( ...
        spatialSizeParam,temporalSizeParam,unspecifiedSizeParam, ...
        spatialSizeX,temporalSizeX,unspecifiedSizeX);
end

% The dlarray methods should not accept logical inputs
if islogical(offset) || islogical(scale)
    error(message('deep:dlarray:LogicalsNotSupported'));
end

end

function normDim = iConvertNormalizationDimToNumeric(normDim, fd, channelDim)
spatialDim = indexOfLabels(fd, 'S');
numSLabels = numel(spatialDim);

switch normDim
    case "auto"
        if numSLabels < 2
            normDim = channelDim;
        else
            normDim = [spatialDim, channelDim];
        end
    case "channel-only"
        normDim = channelDim;
    case "spatial-channel"
        normDim = [spatialDim, channelDim];
    case "batch-excluded"
        timeDim = indexOfLabels(fd, 'T');
        unspecifiedDim = indexOfLabels(fd, 'U');
        normDim = [spatialDim, channelDim, timeDim, unspecifiedDim];
end
end
