function layer = lstmLayer(numHiddenUnits, nameValueArgs)
%

%   Copyright 2017-2023 The MathWorks, Inc.

% Parse the input arguments.
arguments
    numHiddenUnits (1,1) {mustBeNumeric, mustBeInteger, mustBePositive}
    nameValueArgs.Name  {mustBeText, iAssertValidLayerName} = ""
    nameValueArgs.OutputMode {mustBeText} = "sequence"
    nameValueArgs.StateActivationFunction {mustBeText} = "tanh"
    nameValueArgs.GateActivationFunction {mustBeText} = "sigmoid"
    nameValueArgs.HasStateInputs (1,1) {iAssertBinary} = false
    nameValueArgs.HasStateOutputs (1,1) {iAssertBinary} = false
    nameValueArgs.InputWeights = []
    nameValueArgs.RecurrentWeights = []
    nameValueArgs.Bias = []
    nameValueArgs.InputWeightsLearnRateFactor (1,:) {iAssertValidFactor} = 1
    nameValueArgs.RecurrentWeightsLearnRateFactor (1,:) {iAssertValidFactor} = 1
    nameValueArgs.BiasLearnRateFactor (1,:) {iAssertValidFactor} = 1
    nameValueArgs.InputWeightsL2Factor (1,:) {iAssertValidFactor} = 1
    nameValueArgs.RecurrentWeightsL2Factor (1,:) {iAssertValidFactor} = 1
    nameValueArgs.BiasL2Factor (1,:) {iAssertValidFactor} = 0
    nameValueArgs.InputWeightsInitializer = "glorot"
    nameValueArgs.RecurrentWeightsInitializer = "orthogonal"
    nameValueArgs.BiasInitializer = "unit-forget-gate"
    nameValueArgs.HiddenState = iDefaultState()
    nameValueArgs.CellState = iDefaultState()
end

% Canonicalization.
numHiddenUnits = iConvertToDouble(numHiddenUnits);
nameValueArgs = nnet.internal.cnn.layer.util.gatherParametersToCPU(nameValueArgs);
nameValueArgs.Name = convertStringsToChars(nameValueArgs.Name);
nameValueArgs.InputSize = [];
nameValueArgs.OutputMode = iAssertAndReturnValidOutputMode(nameValueArgs.OutputMode);
nameValueArgs.StateActivationFunction = iAssertAndReturnValidStateActivation(nameValueArgs.StateActivationFunction);
nameValueArgs.GateActivationFunction = iAssertAndReturnValidGateActivation(nameValueArgs.GateActivationFunction);
nameValueArgs.HasStateInputs = logical(nameValueArgs.HasStateInputs);
nameValueArgs.HasStateOutputs = logical(nameValueArgs.HasStateOutputs);
nameValueArgs.HiddenState = iAssertAndReturnValidState(nameValueArgs.HiddenState, nameValueArgs.HasStateInputs, "'HiddenState'");
nameValueArgs.CellState = iAssertAndReturnValidState(nameValueArgs.CellState, nameValueArgs.HasStateInputs, "'CellState'");

% Create an internal representation of the layer.
internalLayer = nnet.internal.cnn.layer.LSTM(nameValueArgs.Name, ...
    nameValueArgs.InputSize, ...
    numHiddenUnits, ...
    true, ...
    true, ...
    iGetReturnSequence(nameValueArgs.OutputMode), ...
    nameValueArgs.StateActivationFunction, ...
    nameValueArgs.GateActivationFunction, ...
    nameValueArgs.HasStateInputs, ...
    nameValueArgs.HasStateOutputs);

% Use the internal layer to construct a user visible layer.
layer = nnet.cnn.layer.LSTMLayer(internalLayer);

% Set learnable parameters, learn rate, L2 factors and initializers.
layer.InputWeights = nameValueArgs.InputWeights;
layer.InputWeightsL2Factor = nameValueArgs.InputWeightsL2Factor;
layer.InputWeightsLearnRateFactor = nameValueArgs.InputWeightsLearnRateFactor;
layer.InputWeightsInitializer = nameValueArgs.InputWeightsInitializer;

layer.RecurrentWeights = nameValueArgs.RecurrentWeights;
layer.RecurrentWeightsL2Factor = nameValueArgs.RecurrentWeightsL2Factor;
layer.RecurrentWeightsLearnRateFactor = nameValueArgs.RecurrentWeightsLearnRateFactor;
layer.RecurrentWeightsInitializer = nameValueArgs.RecurrentWeightsInitializer;

layer.Bias = nameValueArgs.Bias;
layer.BiasL2Factor = nameValueArgs.BiasL2Factor;
layer.BiasLearnRateFactor = nameValueArgs.BiasLearnRateFactor;
layer.BiasInitializer = nameValueArgs.BiasInitializer;

% Set hidden state and cell state.
if ~layer.HasStateInputs
    layer.HiddenState = nameValueArgs.HiddenState;
    layer.CellState = nameValueArgs.CellState;
end
end

function tf = iGetReturnSequence( mode )
tf = true;
if strcmp( mode, 'last' )
    tf = false;
end
end

function state = iDefaultState()
state = [];
end

function validString = iAssertAndReturnValidOutputMode(value)
validString = validatestring(value, {'sequence', 'last'});
end

function validString = iAssertAndReturnValidStateActivation(value)
validString = validatestring(value, {'tanh', 'softsign', 'relu'});
end

function validString = iAssertAndReturnValidGateActivation(value)
validString = validatestring(value, {'sigmoid', 'hard-sigmoid'});
end

function iAssertBinary(value)
nnet.internal.cnn.options.OptionsValidator.assertBinary(value);
end

function iAssertValidFactor(value)
validateattributes(value, {'numeric'},  {'vector', 'real', 'nonnegative', 'finite'});
end

function state = iAssertAndReturnValidState(state, hasStateInputs, name)
% If hasStateInputs is true, then only the default state is valid
if hasStateInputs && ~isequal(state, iDefaultState())
    error( message('nnet_cnn:layer:LSTMLayer:SettingStateWithStateInputs', name) );
end
end

function iAssertValidLayerName(name)
iEvalAndThrow(@()...
    nnet.internal.cnn.layer.paramvalidation.validateLayerName(name));
end

function iEvalAndThrow(func)
% Omit the stack containing internal functions by throwing as caller
try
    func();
catch exception
    throwAsCaller(exception)
end
end

function value = iConvertToDouble(value)
value = nnet.internal.cnn.layer.util.convertToDouble(value);
end
