function [X, W] = attention(Q, K, V, numHeads, options)
%ATTENTION   Dot product attention
%
%   Y = ATTENTION(QUERIES,KEYS,VALUES,NUMHEADS) computes the dot product
%   attention of QUERIES, KEYS, and VALUES using the number of attention
%   heads NUMHEADS. QUERIES must be a formatted dlarray. KEYS and VALUES
%   can be dlarray or numeric arrays. If KEYS or VALUES is a formatted
%   dlarray, its format must match the format of QUERIES. If KEYS or VALUES
%   is not a formatted dlarray, the function uses the same format as
%   QUERIES.
%
%   QUERIES can have at most one dimension labeled 'S' or 'T'. If QUERIES
%   has any dimensions labeled 'U', they must be singleton. The size of the
%   'C' dimension in KEYS must match the size of the corresponding
%   dimension in QUERIES. The size of any 'S' or 'T' dimension in KEYS must
%   match the size of the corresponding dimension in VALUES. The size of
%   any 'B' dimension in QUERIES, KEYS, and VALUES must match.
%
%   NUMHEADS must be a positive integer scalar and must be a factor of the
%   size of the channel dimensions of QUERIES, KEYS, and VALUES. Output Y
%   is a formatted dlarray with the same dimension labels as QUERIES. The
%   size of the 'C' dimension in Y is the same as the size of the
%   corresponding dimension in VALUES and the size of the 'S' or 'T'
%   dimension in Y is the same as the size of the corresponding dimension
%   in QUERIES.
%
%   [Y, WEIGHTS] = ATTENTION(QUERIES,KEYS,VALUES,NUMHEADS) computes the dot
%   product attention and also returns the attention weights. WEIGHTS is a
%   Nk-by-Nq-by-numHeads-by-numObservations unformatted dlarray, where Nk
%   is the size of the 'S' or 'T' dimension in KEYS, Nq is the size of the
%   corresponding dimension in QUERIES, numHeads is the number of attention
%   heads NUMHEADS, and numObservations is the size of the 'B' dimension in
%   the input data.
%
%   [Y, WEIGHTS] = ATTENTION(QUERIES,KEYS,VALUES,NUMHEADS,DataFormat=FMT)
%   computes the dot product attention when QUERIES is not a formatted
%   dlarray. FMT specifies the dimension labels of the input QUERIES. FMT
%   must be a character array or string scalar. If KEYS or VALUES is a
%   formatted dlarray, its format must match the format FMT. If KEYS or
%   VALUES is not a formatted dlarray, the function uses the format FMT. If
%   QUERIES is a numeric array, KEYS or VALUES must be a dlarray. Output Y
%   is an unformatted dlarray with the same dimension order as QUERIES.
%
%   [Y, WEIGHTS] = ATTENTION(__,Name=Value) specifies additional options
%   using one or more name-value arguments:
%
%      Scale             - Multiplicative factor for scaling the dot
%                          product of QUERIES and KEYS, specified as one of
%                          the following:
%                          'auto' - The function multiplies the dot product
%                                   by 1/sqrt(D), where D is size of the
%                                   'C' dimension of KEYS divided by
%                                   NUMHEADS. This is the value proposed in
%                                   [1].
%                          Scalar - Multiply the dot product by the
%                                   specified scalar value. Value 1 implies
%                                   no scaling.
%                          The default value is 'auto'.
%
%     PaddingMask        - Mask preventing attention to padding elements
%                          in input data key-value pairs, specified as a
%                          dlarray, logical array, or numeric array
%                          consisting of 0 and 1. The function prevents
%                          and allows attention to elements of input data
%                          key-value pairs when the corresponding element
%                          in PaddingMask is 0 and 1 respectively. If
%                          PaddingMask is a formatted dlarray, its format
%                          must match that of KEYS. If PaddingMask is not
%                          a formatted dlarray, the function uses the same
%                          format as KEYS. The size of the 'S', 'T', and
%                          'B' dimensions in PaddingMask must match the
%                          size of the corresponding dimensions in input
%                          data key-value pairs.
%                          The default value is a logical array of ones.
%
%     AttentionMask      - Mask preventing attention to elements in certain
%                          positions in input data key-value pairs,
%                          specified as one of the following:
%                          'none'   - Do not prevent attention to elements
%                                     with respect to their positions. If
%                                     PaddingMask is specified, the
%                                     function prevents attention to
%                                     padding elements only.
%                          'causal' - Prevent an element in a position M
%                                     in the 'S' or 'T' dimension of
%                                     QUERIES from providing attention to
%                                     an element in position N > M in the
%                                     corresponding dimension of input
%                                     key-value pairs. Use this option for
%                                     auto-regressive models.
%                          Array    - Logical or numeric array or
%                                     unformatted dlarray consisting of 0
%                                     and 1. The function prevents
%                                     attention to elements of input data
%                                     key-value pairs in positions
%                                     corresponding to the positions of 0
%                                     in the array. The array must be a
%                                     Nk-by-Nq matrix or a
%                                     Nk-by-Nq-by-numObservations array,
%                                     where Nk is the size of the 'S' or
%                                     'T' dimension in KEYS, Nq is the size
%                                     of the corresponding dimension in
%                                     QUERIES, and  numObservations is the
%                                     size of the 'B' dimension in the
%                                     input data.
%                            The default value is 'none'.
%
%     DropoutProbability - Dropout probability for attention weights,
%                          specified as a nonnegative scalar less than 1.
%                          The default value is 0.
%
%     DataFormat         - Dimension labels of QUERIES, specified as a
%                          character vector or string scalar. If QUERIES
%                          is not a formatted dlarray, then you must
%                          specify the data format.
%
%   Example:
%       % Specify parameters for the queries, keys, and values.
%       queryDim = 100;
%       valueDim = 120;
%       numQueries = 64;
%       numValues = 80;
%       numObservations = 32;
%       % Create the queries as formatted dlarray with channel,
%       % observation, and time dimensions.
%       Q = dlarray(rand(queryDim, numObservations, numQueries), 'CBT');
%       % Create the keys and values as unformatted dlarray
%       K = dlarray(rand(queryDim, numObservations, numValues));
%       V = dlarray(rand(valueDim, numObservations, numValues));
%
%       % Specify number of attention heads
%       numHeads = 5;
%
%       % Compute dot product attention
%       [Y, W] = attention(Q, K, V, numHeads);
%       % Outputs:
%       % Y - 120(C) x 32(B) x 64(T) dlarray double
%       % W - 80 x 64 x 5 x 32 dlarray double
%
%   1. Vaswani et al. Attention is all you need. Advances in Neural
%      Information Processing Systems (2017).
%
%   See also DLARRAY, LSTM, DLCONV

%   Copyright 2022 The MathWorks, Inc.

arguments
    Q {iMustBeSingleOrDouble, mustBeNonsparse, mustBeReal}
    K {iMustBeSingleOrDouble, mustBeNonsparse, mustBeReal}
    V {iMustBeSingleOrDouble, mustBeNonsparse, mustBeReal}
    numHeads (1,1) {iMustBeSingleOrDouble, mustBeInteger, mustBePositive, mustBeNonsparse, deep.internal.dlarray.mustBeNonDlarray}
    options.Scale
    options.PaddingMask {deep.internal.dlarray.mustBeBinaryMask}
    options.AttentionMask = "none"
    options.DropoutProbability (1,1) {iMustBeSingleOrDouble, mustBeNonsparse, deep.internal.dlarray.mustBeNonDlarray, ...
        mustBeInRange(options.DropoutProbability, 0, 1, "exclude-upper")} = 0
    options.DataFormat {deep.internal.dlarray.mustBeDataFormatType} = ''
end

% Ensure data format is in the canonical form
dataFormat = char(options.DataFormat);

% Extract and validate data
[qdata, kdata, vdata, X, fd, dataOrder, hasSpatialDim, hasObservationDim, hasTimeDim, hasChannelDim, positionDim, observationDim, ...
    permFormat, hasNoDataFormat, labels, numLabels, numObservations] = deep.internal.dlarray.extractAttentionData(Q, K, V, dataFormat);

% numHeads should not determine execution environment
numHeads = gather(numHeads);
% Assert numHeads is a factor of sizes of channel dimensions of data
keySize = size(kdata);
keyNumChannelsPerHead = keySize(1)/numHeads;
valueNumChannelsPerHead = size(vdata, 1)/numHeads;
if mod(keyNumChannelsPerHead, 1) ~= 0 || mod(valueNumChannelsPerHead, 1) ~= 0
    error(message('deep:dlarray:AttentionInvalidNumHeads'));
end

% Validate scale
scale = 1/sqrt(keyNumChannelsPerHead);
if isfield(options, 'Scale')
    scale = iValidateScale(options.Scale, scale);
end

% Validate padding mask
keyPositionLength = keySize(2);
queryPositionLength = size(qdata, 2);
if isfield(options, 'PaddingMask')
    paddingMask = deep.internal.dlarray.validateAttentionPaddingMask( ...
        options.PaddingMask, keyPositionLength, queryPositionLength, numObservations, numLabels, labels, permFormat, ...
        hasNoDataFormat, hasSpatialDim, hasTimeDim, positionDim, observationDim, dataOrder);
else
    paddingMask = 1;
end

% Construct attention mask
attentionMask = deep.internal.dlarray.constructAttentionMask( ...
    options.AttentionMask, keyPositionLength, queryPositionLength, numObservations, paddingMask);

% Call the internal API
dropoutProb = gather(options.DropoutProbability);
[xdata, wdata] = internal_attention( ...
    qdata, kdata, vdata, attentionMask, numHeads, scale, keyNumChannelsPerHead, valueNumChannelsPerHead, ...
    queryPositionLength, keyPositionLength, dropoutProb);

% Insert output
if nargout > 1
    fdW = fd;
    dlW = X;
    wfmt = '';
    perm = [];
    W = iInsertFormattedData(wdata, dlW, fdW, wfmt, perm);
end

if hasTimeDim
    X = iInsertOutputWithTimeDimension(xdata, fd, X, hasObservationDim, hasNoDataFormat, dataFormat, permFormat, hasChannelDim);
elseif hasSpatialDim
    X = iInsertOutputWithSpatialDimension(xdata, fd, X, hasObservationDim, hasNoDataFormat, dataFormat, permFormat, hasChannelDim);
else
    X = iInsertOutput(xdata, fd, X, hasObservationDim, hasNoDataFormat, dataFormat, permFormat, hasChannelDim);
end

end

function iMustBeSingleOrDouble(val)
if ~isfloat(val)
    error(message('deep:dlarray:AttentionInvalidDatatype'));
end
end

function scale = iValidateScale(scale, defaultScale)
if ischar(scale) || isstring(scale)
    if ~strcmp(scale, "auto")
        error(message('deep:dlarray:AttentionInvalidScale'));
    else
        scale = defaultScale;
    end
elseif ~(isfloat(scale) && isscalar(scale) && isfinite(scale) && isreal(scale) && ~issparse(scale))
    error(message('deep:dlarray:AttentionInvalidScale'));
else
    % Only data and mask defines execution environment
    scale = gather(scale);
end
end

function X = iInsertOutputWithTimeDimension(xdata, fd, X, hasObservationDim, hasNoDataFormat, dataFormat, permFormat, hasChannelDim)
if hasNoDataFormat
    % Set format to match data order and specify permutation order for
    % permuting data to the canonical form
    if hasObservationDim
        xfmt = 'CTB';
        perm = [1 3 2];
    else
        xfmt = 'CT';
        perm = [];
    end
else
    xfmt = '';
    if dataFormat(1) == 'C' && (contains("CTBU", dataFormat) || strcmp(dataFormat, "CU"))
        % DataFormat is in the same order as data. Only need to convert
        % data to dlarray
        perm = [];
    else
        if hasObservationDim
            perm = [1 3 2];
        else
            perm = [1 2];
        end
        % Get the permutation order for permuting data to the same order as
        % DataFormat
        perm = iGetPermutationOrderForUnformattedOutput(perm, permFormat, hasChannelDim);
    end
end
% Convert data to dlarray
X = iInsertFormattedData(xdata, X, fd, xfmt, perm);
end

function X = iInsertOutputWithSpatialDimension(xdata, fd, X, hasObservationDim, hasNoDataFormat, dataFormat, permFormat, hasChannelDim)
if hasNoDataFormat
    % Set format to match data order and specify permutation order for
    % permuting data to the canonical form
    if hasObservationDim
        xfmt = 'CSB';
        perm = [2 1 3];
    else
        xfmt = 'CS';
        perm = [2 1];
    end
else
    xfmt = '';
    if dataFormat(1) == 'C' && (contains("CSBU", dataFormat) || strcmp(dataFormat, "CU"))
        % DataFormat is in the same order as data. Only need to convert
        % data to dlarray
        perm = [];
    else
        if hasObservationDim
            perm = [2 1 3];
        else
            perm = [2 1];
        end
        % Get the permutation order for permuting data to the same order as
        % DataFormat
        perm = iGetPermutationOrderForUnformattedOutput(perm, permFormat, hasChannelDim);
    end
end
% Convert data to dlarray
X = iInsertFormattedData(xdata, X, fd, xfmt, perm);
end

function X = iInsertOutput(xdata, fd, X, hasObservationDim, hasNoDataFormat, dataFormat, permFormat, hasChannelDim)
if hasNoDataFormat
    perm = [];
    % Remove singleton "position" dimension from data
    sz = size(xdata);
    xdata = reshape(xdata,sz(1),sz(end));
    if hasObservationDim
        xfmt = 'CB';
    else
        xfmt = 'CU';
    end
else
    xfmt = '';
    if strcmp(dataFormat, "CU")
        % DataFormat is in the same order as data. Only need to convert
        % data to dlarray
        perm = [];
    else
        % Remove singleton "position" dimension from data
        sz = size(xdata);
        xdata = reshape(xdata,sz(1),sz(end));
        % Get the permutation order for permuting data to the same order as
        % DataFormat
        perm = iGetPermutationOrderForUnformattedOutput([1 2], permFormat, hasChannelDim);
    end
end
% Convert data to dlarray
X = iInsertFormattedData(xdata, X, fd, xfmt, perm);
end

function X = iInsertFormattedData(xdata, X, fd, xfmt, perm)
% Insert FormattedData
fd.DimensionMetadata = xfmt;
fd = insertData(fd, xdata);
X.FormattedData = fd;

% Remove any permutation from the output dlarray
X.ViewToDataMap = perm;
end

function perm = iGetPermutationOrderForUnformattedOutput(perm, permFormat, hasChannelDim)
% Input perm is the permutation order of data returned by the internal
% function to the canonical order. Output perm is the permutation order of
% data returned by the internal function to the user-sepcified data order.
% permFormat is the permutation order of the user-specified data order to
% the canonical order.

if ~hasChannelDim
    % Place the channel dimension at the end
    channelViewIdx = numel(permFormat)+1;
    permFormat = [channelViewIdx permFormat];
end
numPermFormat = numel(permFormat);
numPerm = numel(perm);
if numPerm < numPermFormat
    % Add dimensions corresponding to trailing singleton dimensions in
    % permFormat to perm
    perm = [perm, numPerm+1:numPermFormat];
end
% Obtain the inverse permutation order for permFormat. This is the order
% that can be used to permute back from canonical order to user-specified
% data order
inversePermFormat(permFormat) = 1:numPermFormat;
% Use the inverse permutation order to sort input perm so that the
% resulting order can be used to directly permute data returned by the
% internal function to user-specified data order
perm = perm(inversePermFormat);
end
