function [X, h0, c0] = lstm(X, h0, c0, W, R, B, nvArgs)
%LSTM   Long short-term memory
%
%    Y = LSTM(X,H0,C0,WEIGHTS,RECURRENTWEIGHTS,BIAS) applies a long
%    short-term memory (lstm) calculation to the formatted dlarray X, using
%    initial hidden state H0 and initial cell state C0. 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 lstm operation.
%    
%    Inputs H0, C0, 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.
%
%    Inputs H0 and C0 must be vectors or matrices. If H0 or C0 are
%    formatted dlarrays, then they must contain the label 'C' and
%    optionally the label 'B'. The size of the 'C' dimension determines the
%    number of hidden units, and must be the same for H0 and C0. If H0 or
%    C0 are unformatted dlarrays or numeric arrays, the size of the first
%    dimension determines the number of hidden units.
%
%    Input WEIGHTS must be a matrix of size 4*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
%    4*NumHiddenUnits and a 'U' dimension of size InputSize.
%
%    Input RECURRENTWEIGHTS must be a matrix of size
%    4*NumHiddenUnits-by-NumHiddenUnits. If RECURRENTWEIGHTS is a formatted
%    dlarray, then it must contain a 'C' dimension of size 4*NumHiddenUnits
%    and a 'U' dimension of size NumHiddenUnits.
%
%    Input BIAS must be a vector of size 4*NumHiddenUnits. If BIAS is a
%    formatted dlarray, the non-singleton dimension must be labeled with
%    'C'.
%
%    [Y,HIDDENSTATE,CELLSTATE] =
%    LSTM(X,H0,C0,WEIGHTS,RECURRENTWEIGHTS,BIAS) applies an lstm
%    calculation to the formatted dlarray X and also returns the hidden
%    state and cell state. HIDDENSTATE and CELLSTATE have the same type as
%    H0 and C0, respectively. If the inputs H0 or C0 are formatted
%    dlarrays, outputs HIDDENSTATE and CELLSTATE are labeled 'CB'.
%
%    [Y,HIDDENSTATE,CELLSTATE] =
%    LSTM(X,H0,C0,WEIGHTS,RECURRENTWEIGHTS,BIAS,'DataFormat',FMT) applies
%    an lstm 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, C0, WEIGHTS, RECURRENTWEIGHTS, or BIAS must be a dlarray. The
%    output Y is an unformatted dlarray with the same dimension order as X.
%
%    Y = LSTM(_,'PARAM1',VAL1,'PARAM2',VAL2,...) specifies optional
%    parameter name/value pairs:
%
%      'StateActivationFunction' - Activation function to update the cell 
%                                  and hidden state. Options are:
%                                   - 'tanh'
%                                   - 'softsign'
%                                   - 'relu'
%                                  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 a channel, observation, and time
%       % dimensions.
%       numFeatures = 10;
%       numObservations = 32;
%       sequenceLength = 64;
%       X = dlarray(randn(numFeatures,numObservations,sequenceLength), 'CBT');
%
%       % Create formatted dlarrays for the lstm parameters with three
%       % hidden units.
%       numHiddenUnits = 3;
%       H0 = dlarray(randn(numHiddenUnits,numObservations),'CB');
%       C0 = dlarray(randn(numHiddenUnits,numObservations),'CB');
%       weights = dlarray(randn(4*numHiddenUnits,numFeatures),'CU');
%       recurrent = dlarray(randn(4*numHiddenUnits,numHiddenUnits),'CU');
%       bias = dlarray(randn(4*numHiddenUnits,1),'C');
%
%       % Apply an lstm calculation
%       [Y,hiddenState,cellState] = lstm(X,H0,C0,weights,recurrent,bias);
%
%       % View the size of Y, hiddenState and cellState
%       size(Y)
%       dims(Y)
%       size(hiddenState)
%       dims(hiddenState)
%       size(cellState)
%       dims(cellState)
%
%   See also DLARRAY, GRU, FULLYCONNECT, SOFTMAX

%   Copyright 2019-2023 The MathWorks, Inc.

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

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

[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);

xdata = deep.internal.dlarray.orderRecurrentDataForBuiltin(xdata, channelDim, observationDim, timeDim);

% Extract the hidden units and cell states
[h0, h0data] = extractState(h0, nbObservations, @iHiddenStateErrorIDs);
[c0, c0data] = extractState(c0, nbObservations, @iCellStateErrorIDs);
nbHiddenUnits = size(h0data, 1);
if size(c0data, 1) ~= nbHiddenUnits
    error(message('deep:dlarray:LstmInconsistentStates'));
end

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

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

% Extract & validate bias
bdata = extractParameter(B, @iBiasErrorIDs);
bdata = deep.internal.dlarray.validateBias(bdata, 4*nbHiddenUnits, ...
                                           'deep:dlarray:LstmBiasInconsistentHiddenUnits');

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

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

[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);
% Keep only the last value of c0data
c0data = c0data(:,:,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 and third outputs
h0 = formatOutputState(h0, h0data);
c0 = formatOutputState(c0, c0data);
end

function iMustBeValidDataFormat(arg)
validateattributes(arg, {'char', 'string'}, {'nonempty'});
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 = iCellStateErrorIDs()
s = struct(...
    'InvalidStateObservation', 'deep:dlarray:LstmInvalidCellStateObservation', ...
    'InvalidLabeled', 'deep:dlarray:LstmInvalidLabeledCellState', ...
    'TooManyDimensions', 'deep:dlarray:LstmTooManyDimensionsCellState');
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
