function layer = lstmProjectedLayer(numHiddenUnits, outputProjectorSize, inputProjectorSize, nameValueArgs)
%

%   Copyright 2022-2023 The MathWorks, Inc.

% Parse the input arguments.
arguments
    numHiddenUnits (1,1) {mustBeNumeric, mustBeInteger, mustBePositive}
    outputProjectorSize (1,1) {mustBeNumeric, mustBeInteger, mustBePositive}
    inputProjectorSize (1,1) {mustBeNumeric, mustBeInteger, mustBePositive}
    nameValueArgs.Name  {mustBeText, nnet.internal.cnn.layer.paramvalidation.validateLayerName} = ""
    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.InputProjector = []
    nameValueArgs.OutputProjector = []
    nameValueArgs.InputWeightsLearnRateFactor (1,:) {iAssertValidFactor} = 1
    nameValueArgs.RecurrentWeightsLearnRateFactor (1,:) {iAssertValidFactor} = 1
    nameValueArgs.BiasLearnRateFactor (1,:) {iAssertValidFactor} = 1
    nameValueArgs.InputProjectorLearnRateFactor (1,:) {iAssertValidFactor} = 1
    nameValueArgs.OutputProjectorLearnRateFactor (1,:) {iAssertValidFactor} = 1
    nameValueArgs.InputWeightsL2Factor (1,:) {iAssertValidFactor} = 1
    nameValueArgs.RecurrentWeightsL2Factor (1,:) {iAssertValidFactor} = 1
    nameValueArgs.BiasL2Factor (1,:) {iAssertValidFactor} = 0
    nameValueArgs.InputProjectorL2Factor (1,:) {iAssertValidFactor} = 1
    nameValueArgs.OutputProjectorL2Factor (1,:) {iAssertValidFactor} = 1
    nameValueArgs.InputWeightsInitializer = "glorot"
    nameValueArgs.RecurrentWeightsInitializer = "orthogonal"
    nameValueArgs.BiasInitializer = "unit-forget-gate"
    nameValueArgs.InputProjectorInitializer = "orthogonal"
    nameValueArgs.OutputProjectorInitializer = "orthogonal"
    nameValueArgs.HiddenState = iDefaultState()
    nameValueArgs.CellState = iDefaultState()
end

% Canonicalization.
numHiddenUnits = iConvertToDouble(numHiddenUnits);
outputProjectorSize = iConvertToDouble(outputProjectorSize);
inputProjectorSize = iConvertToDouble(inputProjectorSize);
nameValueArgs.InputSize = [];
nameValueArgs = nnet.internal.cnn.layer.util.gatherParametersToCPU(nameValueArgs);
nameValueArgs.Name = convertStringsToChars(nameValueArgs.Name);
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.LSTMProjected(nameValueArgs.Name, ...
    numHiddenUnits, ...
    outputProjectorSize, ...
    nameValueArgs.InputSize, ...
    inputProjectorSize, ...
    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.LSTMProjectedLayer(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;

layer.InputProjector = nameValueArgs.InputProjector;
layer.InputProjectorL2Factor = nameValueArgs.InputProjectorL2Factor;
layer.InputProjectorLearnRateFactor = nameValueArgs.InputProjectorLearnRateFactor;
layer.InputProjectorInitializer = nameValueArgs.InputProjectorInitializer;

layer.OutputProjector = nameValueArgs.OutputProjector;
layer.OutputProjectorL2Factor = nameValueArgs.OutputProjectorL2Factor;
layer.OutputProjectorLearnRateFactor = nameValueArgs.OutputProjectorLearnRateFactor;
layer.OutputProjectorInitializer = nameValueArgs.OutputProjectorInitializer;

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

function tf = iGetReturnSequence( mode )
tf = ~strcmp(mode, 'last');
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 value = iConvertToDouble(value)
value = nnet.internal.cnn.layer.util.convertToDouble(value);
end
