function X = instancenorm(X, offset, scale, nvArgs)
%INSTANCENORM   Normalize data across each channel
%
%   Y = INSTANCENORM(X,OFFSET,SCALEFACTOR) normalizes the input data X
%   across each channel and performs an affine transformation with the
%   specified OFFSET and SCALEFACTOR. The function normalizes
%   per-observation and applies the affine transformation per-channel.
%   X must be a formatted dlarray with a 'C' dimension.
%
%   If OFFSET and SCALEFACTOR are unformatted dlarray objects or 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
%   formatted dlarray objects, then they must have 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.
%
%   Y = INSTANCENORM(__,'DataFormat',FMT) computes the instance 
%   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 = INSTANCENORM(__,'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.
%       '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 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 instance normalization
%       Z = instancenorm(X,offset,scaleFactor);
%
%   See also DLARRAY, BATCHNORM, GROUPNORM, RELU, DLCONV, FULLYCONNECT

%   Copyright 2020-2022 The MathWorks, Inc.
arguments
    X {mustBeNumeric, mustBeNonsparse, mustBeReal}
    offset {mustBeNumeric, mustBeNonsparse, mustBeReal, mustBeNonempty}
    scale {mustBeNumeric, mustBeNonsparse, mustBeReal, mustBeNonempty}
    nvArgs.DataFormat {deep.internal.dlarray.mustBeDataFormatType} = ''
    nvArgs.Epsilon {deep.internal.dlarray.mustBeValidEpsilon} = 1e-5 % by convention (including cuDNN)
end

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

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

X = permuteExplicitly(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);

% Retrieve batch index. If data does not have 'B' dimensions, but has 'T'
% and/or 'U', we need to and insert a singleton 'B' dimension. This is
% necessary because the forward and backward methods need to know the
% relative position of 'C' and 'B' dimension in the data tensor.
batchDim = observationIndex(fd);
needsSingletonBatch = batchDim > ndims(xdata);

if needsSingletonBatch
    origSizexdata = size(xdata);

    % Insert singleton dimension for the batch, after 'C', but before 'T'
    % and/or 'U'
    % Get dimensions until 'C'
    origSizexdataBeforeBatch = size(xdata, 1:channelDim);
    % Get dimensions from 'T' and/or 'U'
    origSizexdataAfterBatch = origSizexdata(channelDim+1:end);
    % Insert singleton dimension
    dimsSingleton = [origSizexdataBeforeBatch, 1, origSizexdataAfterBatch];
    xdata = reshape(xdata,dimsSingleton);
end

% Set the number of groups
numGroups = numChannels;

% Extract offset & scale
offset = iExtractLearnableParameter(offset, numChannels, @iOffsetErrorIDs);
scale = iExtractLearnableParameter(scale, numChannels, @iScaleErrorIDs);

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

% Be sure that the learnable parameters have shape compatible with the data
[offset, scale] = ...
    iReshapeParameters(offset, scale, ndims(xdata), channelDim, numChannels);

% Call the internal group norm API
xdata = internal_groupnorm(xdata, offset, scale, numGroups, nvArgs.Epsilon, channelDim);

% Restore the data shape
if needsSingletonBatch
    xdata = reshape(xdata, origSizexdata);
end

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

% For unformatted, 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(offset, scale, ndimsInput, channelDim, numChannels)
% Reshape the parameters

% Get the reshape vector
reshapeVec = ones(1, max(ndimsInput,channelDim));
reshapeVec(channelDim) = numChannels;

% Reshape the parameters if not empty
offset = reshape(offset, reshapeVec);
scale = reshape(scale, reshapeVec);
end

function value = iExtractLearnableParameter(value, numChannels, errorIDFcn)
% This function extract the learnable parameter's data (if wrapped in a
% dlarray)
if isa(value, 'dlarray')
    [fd,value] = extractData(value.FormattedData);
    if ~isempty(fd.DimensionMetadata) % If formatted dlarray
        numChannelsParam = size(value, channelIndex(fd));
        if numChannelsParam ~= numChannels
            error(iGetErrorMessage(errorIDFcn, 'InvalidLabeled', numChannels));
        end
    else % Unformatted dlarray
        numChannelsParam = numel(value);
    end
else
    % The type of the output of the operation could be modified by the
    % computation to a type that is unsupported by dlarray
    deep.internal.dlarray.validateNonDlarray(value);
    numChannelsParam = numel(value);
end
% Parameter must be a vector with as many elements as the number of channels
if numChannelsParam ~= numChannels || ~isvector(value)
    error(iGetErrorMessage(errorIDFcn, 'Invalid', numChannels));
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:GroupnormInvalidLabeledOffset', ...
    'Invalid', 'deep:dlarray:GroupnormInvalidOffset');
end

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