function [X, h0] = gru(X, h0, W, R, B, nvArgs)
%GRU   Gated recurrent unit
%
%    Y = GRU(X,H0,WEIGHTS,RECURRENTWEIGHTS,BIAS) applies a gated recurrent
%    unit (gru) calculation to the formatted dlarray X using initial hidden
%    state H0. WEIGHTS, RECURRENTWEIGHTS, and BIAS are learnable parameters
%    for the operation.
%
%    X must be a formatted dlarray. If X has any dimensions labeled 'S',
%    they are flattened into the 'C' dimension. If X has any dimensions
%    labeled 'U', they must be singleton. If X has a dimension labeled 'B',
%    it is unaffected by the gru operation.
%    
%    Inputs H0, WEIGHTS, RECURRENTWEIGHTS, and BIAS can be formatted
%    dlarrays, unformatted dlarrays, or numeric arrays. Output Y is a
%    formatted dlarray. Except for any 'S' dimensions, Y has the same
%    dimension labels as X.
%
%    Input H0 must be a vector or matrix. If H0 is a formatted dlarray,
%    then it must contain the label 'C' and optionally the label 'B'. The
%    size of the 'C' dimension determines the number of hidden units. If H0
%    is an unformatted dlarray or a numeric array, the size of the first
%    dimension determines the number of hidden units.
%
%    Input WEIGHTS must be a matrix of size 3*NumHiddenUnits-by-InputSize,
%    where NumHiddenUnits is given by the size of the 'C' dimension of H0
%    and InputSize is the size of the 'C' dimension of X. If WEIGHTS is a
%    formatted dlarray, then it must contain a 'C' dimension of size
%    3*NumHiddenUnits and a 'U' dimension of size InputSize.
%
%    Input RECURRENTWEIGHTS must be a matrix of size
%    3*NumHiddenUnits-by-NumHiddenUnits. If RECURRENTWEIGHTS is a formatted
%    dlarray, then it must contain a 'C' dimension of size 3*NumHiddenUnits
%    and a 'U' dimension of size NumHiddenUnits.
%
%    Input BIAS must be a vector of size 3*NumHiddenUnits. If BIAS is a
%    formatted dlarray, the non-singleton dimension must be labeled 'C'.
%
%    [Y,HIDDENSTATE] = GRU(X,H0,WEIGHTS,RECURRENTWEIGHTS,BIAS) applies a
%    GRU calculation to the formatted dlarray X and also returns the hidden
%    state. HIDDENSTATE has the same type as H0. If the input H0 is a
%    formatted dlarray, output HIDDENSTATE is labeled 'CB'.
%
%    [Y,HIDDENSTATE] =
%    GRU(X,H0,WEIGHTS,RECURRENTWEIGHTS,BIAS,DataFormat=FMT) applies a GRU
%    calculation to the unformatted dlarray or numeric array X. FMT
%    specifies the dimension labels of the input X. FMT must be a char
%    array or string. If the input X is a numeric array, at least one of
%    H0, WEIGHTS, RECURRENTWEIGHTS, or BIAS must be a dlarray. The output Y
%    is an unformatted dlarray with the same dimension order as X.
%
%    ___= GRU(___,Name=Value) specifies additional options using one or
%    more name-value arguments:
%
%     DataFormat                - Dimension labels of the input data X,
%                                 specified as a string scalar or character
%                                 vector. Required if X is not a formatted
%                                 dlarray.
%
%     ResetGateMode             - Reset gate mode of the GRU operation,
%                                 specified as one of the following:
%                                   - "after-multiplication". Apply reset
%                                      gate after matrix multiplication.
%                                      This option uses the cuDNN library
%                                      when running on GPU.
%                                   - "before-multiplication". Apply 
%                                      reset gate before matrix
%                                      multiplication.
%                                   - "recurrent-bias-after-multiplication". 
%                                      Apply reset gate after matrix
%                                      multiplication and use recurrent
%                                      bias.
%                                 The default value is "after-multiplication".
%
%     StateActivationFunction   - Activation function to update the hidden 
%                                 state. Options are:
%                                   - "tanh"
%                                   - "softsign"
%                                 The default value is "tanh".
%
%     GateActivationFunction    - Activation function to apply to the 
%                                 gates. Options are:
%                                   - "sigmoid"
%                                   - "hard-sigmoid"
%                                 The default value is "sigmoid".
%
%   Example
%
%       % Create a dlarray with channel, observation, and time dimensions.
%       numFeatures = 10;
%       numObservations = 32;
%       seqLength = 64;
%       X = dlarray(randn(numFeatures,numObservations,seqLength,'single'), 'CBT');
%
%       % Create formatted dlarrays for the GRU parameters with three
%       % hidden units.
%       numHiddenUnits = 3;
%       H0 = dlarray(randn(numHiddenUnits,numObservations),'CB');
%       weights = dlarray(randn(3*numHiddenUnits,numFeatures),'CU');
%       recurrent = dlarray(randn(3*numHiddenUnits,numHiddenUnits),'CU');
%       bias = dlarray(randn(3*numHiddenUnits,1),'C');
%
%       % Apply a GRU calculation
%       [Y,hiddenState] = gru(X,H0,weights,recurrent,bias);
%
%       % View the size of Y and hiddenState
%       size(Y)
%       dims(Y)
%       size(hiddenState)
%       dims(hiddenState)
%
%   See also DLARRAY, LSTM, FULLYCONNECT, SOFTMAX

%   Copyright 2019-2023 The MathWorks, Inc.

arguments
    X                       {mustBeNumeric, mustBeNonsparse, mustBeReal}
    h0                      {mustBeNumeric, mustBeNonsparse, mustBeReal}
    W                       {mustBeNumeric, mustBeNonsparse, mustBeReal}
    R                       {mustBeNumeric, mustBeNonsparse, mustBeReal}
    B                       {mustBeNumeric, mustBeNonsparse, mustBeReal}
    nvArgs.DataFormat       {iMustBeValidDataFormat}
    nvArgs.ResetGateMode    {iMustBeValidResetGateMode} = "after-multiplication"
    nvArgs.StateActivationFunction {mustBeMember(nvArgs.StateActivationFunction,["tanh","softsign"])} = "tanh"
    nvArgs.GateActivationFunction {mustBeMember(nvArgs.GateActivationFunction,["sigmoid","hard-sigmoid"])} = "sigmoid"
end

% Handle data format
dataFormat = iGetDataFormat(nvArgs);
X = deep.internal.dlarray.validateDataFormatArg(X, dataFormat);

resetGateMode = nvArgs.ResetGateMode;

[xdata, fd, XHasObservationDim, XHasTimeDim, channelDim, observationDim, timeDim] = flattenDimOfRecurrentInput(X);

% Clear the formatted data
X.FormattedData = [];

% Retrieve channel and observation dimension index and size
nbChannels = size(xdata, channelDim);
nbObservations = size(xdata, observationDim);

% Order the data for the builtin
xdata = deep.internal.dlarray.orderRecurrentDataForBuiltin(xdata, channelDim, observationDim, timeDim);

% Extract the hidden units
[h0, h0data] = extractState(h0, nbObservations, @iHiddenStateErrorIDs);
nbHiddenUnits = size(h0data, 1);

% Extract & validate weights
wdata = extractParameter(W, @iWeightsErrorIDs);
deep.internal.dlarray.validateWeights(wdata, 3*nbHiddenUnits, nbChannels, ...
                                      'deep:dlarray:GruWeightsInconsistentFirstDim', ...
                                      'deep:dlarray:LstmWeightsInconsistentSecondDim');

% Extract & validate recurrent weights
rdata = extractParameter(R, @iRecurrentWeightsErrorIDs);
deep.internal.dlarray.validateWeights(rdata, 3*nbHiddenUnits, nbHiddenUnits, ...
                                      'deep:dlarray:GruRecurrentWeightsInconsistentFirstDim', ...
                                      'deep:dlarray:LstmRecurrentWeightsInconsistentSecondDim');

% Extract & validate bias
bdata = extractParameter(B, @iBiasErrorIDs);

% Compute trace properties
isTraining = deep.internal.recording.tracingProperties();

if resetGateMode == "after-multiplication"
    bdata = deep.internal.dlarray.validateBias(bdata, 3*nbHiddenUnits, ...
        'deep:dlarray:GruBiasInconsistentHiddenUnits');

    % Call the internal API
    [xdata, h0data] = ...
        internal_gru(xdata, h0data, wdata, rdata, bdata, ...
        nvArgs.StateActivationFunction, nvArgs.GateActivationFunction, ...
        isTraining);
elseif resetGateMode == "recurrent-bias-after-multiplication"
    bdata = deep.internal.dlarray.validateBias(bdata, 6*nbHiddenUnits, ...
        'deep:dlarray:GruBiasInconsistentHiddenUnits');

    % Call the internal API
    [xdata, h0data] = ...
        internal_gru_recurrentBiasAfterMultiplication(xdata, h0data, wdata, rdata, bdata, ...
        nvArgs.StateActivationFunction, nvArgs.GateActivationFunction, ...
        isTraining);
else
    bdata = deep.internal.dlarray.validateBias(bdata, 3*nbHiddenUnits, ...
        'deep:dlarray:GruBiasInconsistentHiddenUnits');

    % Call the internal API
    [xdata, h0data] = ...
        internal_gru_beforeMultiplication(xdata, h0data, wdata, rdata, bdata, ...
        nvArgs.StateActivationFunction, nvArgs.GateActivationFunction);
end

[xdata, xfmt] = deep.internal.dlarray.reshapeRecurrentOutput(xdata, XHasObservationDim, XHasTimeDim, nbHiddenUnits);

% Keep only the last value of h0data to match the GPU output
h0data = h0data(:,:,end);

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

% Remove any permutation from the output dlarray
X.ViewToDataMap = [];

% Remove dimensions from the dlarray if it was passed as unformatted.
if ~isempty(dataFormat)
    X = permuteUnformattedRecurrentOutput(X,~XHasObservationDim,dataFormat);
end

% Create the second output
h0 = formatOutputState(h0, h0data);
end

function iMustBeValidDataFormat(arg)
validateattributes(arg, {'char', 'string'}, {'nonempty'});
end

function iMustBeValidResetGateMode(str)
mustBeMember(str,["after-multiplication","before-multiplication","recurrent-bias-after-multiplication"]);
end

function dataFormat = iGetDataFormat(nameValueArgs)
if isfield(nameValueArgs, 'DataFormat')
    dataFormat = nameValueArgs.DataFormat;
else
    dataFormat = [];
end
end

function s = iHiddenStateErrorIDs()
s = struct(...
    'InvalidStateObservation', 'deep:dlarray:LstmInvalidHiddenStateObservation', ...
    'InvalidLabeled', 'deep:dlarray:LstmInvalidLabeledHiddenState', ...
    'TooManyDimensions', 'deep:dlarray:LstmTooManyDimensionsHiddenState');
end

function s = iWeightsErrorIDs()
s = struct(...
    'InvalidLabeled', 'deep:dlarray:LstmInvalidLabeledWeights', ...
    'TooManyDimensions', 'deep:dlarray:LstmTooManyDimensionsWeights');
end

function s = iRecurrentWeightsErrorIDs()
s = struct(...
    'InvalidLabeled', 'deep:dlarray:LstmInvalidLabeledRecurrentWeights', ...
    'TooManyDimensions', 'deep:dlarray:LstmTooManyDimensionsRecurrentWeights');
end

function s = iBiasErrorIDs()
s = struct(...
    'InvalidLabeled', 'deep:dlarray:LstmInvalidLabeledBias', ...
    'TooManyDimensions', 'deep:dlarray:LstmTooManyDimensionsBias');
end
