function X = embed(X,weights,nvArgs)
%EMBED  Embed discrete data
%   
%    Y = EMBED(X,WEIGHTS) returns the embedding vectors in WEIGHTS
%    corresponding to the numeric indices in the formatted dlarray X
%    containing the sequence of discrete values to embed.
%
%    WEIGHTS must be a formatted dlarray matrix, with 'C' and 'U'
%    dimensions of size K and (V+1), respectively, or an unformatted
%    dlarray matrix, or a numeric array of size K-by-(V+1), where K
%    corresponds to the embedding dimension and V is the vocabulary size.
%    The extra column corresponds to the out-of-vocabulary embedding
%    vector.
%
%    Output Y is a formatted dlarray with a 'C' dimension of size K. All
%    the other dimensions of Y have the same size and labels as X.
%
%    Y = EMBED(X,WEIGHTS,'DataFormat',FMT) returns the embedding of an
%    unformatted dlarray or a numeric array X. FMT specifies the dimension
%    labels of the input data X. FMT must be a char array or string. If
%    input X is a numeric array, then WEIGHTS must be a dlarray. If FMT
%    contains a 'C' label, then output Y is an unformatted dlarray with the
%    same dimension order as X. The size of the corresponding 'C' dimension
%    in output Y is K. If FMT does not contain a 'C' label, then Y has size
%    K-by-size(X).
%
%   Example 1:
%       % Map a sequence of integers to vectors using an embedding with
%       % dimension 10 and vocabulary size 200.
%       X = dlarray(randi(200,[1 16 5]),'CBT');
%       weights = rand(10,201);
%       Z = embed(X,weights);
%       % Output: Z = 10(C) × 16(B) × 5(T) dlarray
%
%   Example 2:
%       % Embed an unformatted dlarray containing a semantic map using an
%       % embedding with dimension 10 and vocabulary size 200.
%       X = dlarray(randi(200,[28 28 1 16]));
%       weights = rand(10,201); 
%       Z = embed(X,weights,'DataFormat','SSCB'); 
%       % Output: Y = 28×28×10×16 dlarray
%
%   See also DLARRAY, LSTM, FULLYCONNECT

%   Copyright 2020-2023 The MathWorks, Inc.

arguments
    X {deep.internal.dlarray.mustBeNumericDLData}
    weights {mustBeNumeric, matlab.internal.validation.mustBeMatrix}
    nvArgs.DataFormat {deep.internal.dlarray.mustBeDataFormatType} = ''
end

fmt = nvArgs.DataFormat;

% Validate that DataFormat input is a char/string
deep.internal.dlarray.mustBeDataFormatType(fmt);

% X input may be a dlarray or a raw numeric. If X is numeric, this will
% create a FormattedData, and a dlarray wrapper.
[X,fd] = deep.internal.dlarray.toInternal(X, true);

% Check that DataFormat can be applied to FormattedData 
fd = applyDataFormatArg(fd, fmt);

% Validate and extract weights
weights = iValidateWeights(weights);

% For embed, errorOnOOV is false by default. This parameter is needed for
% wordEmbeddingLayer
errorOnOOV = false;

% Execute the embed operation
fd = deep.internal.dlarray.extractEmbedding(fd, weights, errorOnOOV);

% Put data back. X will always be a dlarray
fd = unapplyDataFormatArg(fd, fmt);
    
X = deep.internal.dlarray.fromInternal(matlab.lang.internal.move(fd), X);

end

function weights = iValidateWeights(weights)

if isdlarray(weights)
    labels = dims(weights);
    if ~isempty(labels) && ~(count(labels,'C') == 1 && count(labels,'U') == 1)
        error(message('deep:dlarray:EmbedWeightWrongFormat')); 
    end
    % Extract weights data in the permuted order
    weights = permuteExplicitly(weights);
    [~, weights] = extractData(weights.FormattedData);
else
    labels = '';
    % The method should not produce an output with unsupported dlarray types 
    deep.internal.dlarray.validateNonDlarray(weights);
end

% Get vocabulary dimension and embedding size (for formatted weights,
% vocabSize is the size of 'U')
[~, vocabSize] = size(weights);
if vocabSize <= 1
    if ~isempty(labels)
        error(message('deep:dlarray:EmbedEmptyFormattedVocabulary')); 
    else
        error(message('deep:dlarray:EmbedEmptyUnformattedVocabulary')); 
    end
end
end
