function X = ctc( X, T, XMask, TMask, nvArgs )
%CTC    Connectionist temporal classification (CTC) loss for unaligned 
%       sequence classification.
%
%   LOSS = CTC(Y,TARGETS,YMASK,TARGETSMASK) computes the CTC loss of the
%   prediction sequence Y and the target sequence TARGETS for
%   classification problems, where Y and TARGETS can have different
%   sequence lengths. 
%   
%   Y must be a formatted dlarray with dimensions corresponding to
%   classes, observations, and time steps. The dimension corresponding
%   to the classes must have dimension label 'C' (channel). For example,
%   specify Y as a numClasses-by-numObservations-by-lengthPredictions
%   dlarray with format 'CBT' (channel, batch, time).
%
%   TARGETS is an array of integers between 1 and K, where K is the
%   number of channels of Y, with dimensions corresponding to
%   observations and time steps. For example, specify TARGETS as a
%   numObservations-by-lengthTargets dlarray or numeric array. For each
%   observation, the number of time steps of TARGETS lengthTargets is
%   less than or equal lengthPredictions + numRepeats, where
%   lengthPredictions is the number of time steps of Y, and numRepeats
%   is the number of adjacent repeat elements in TARGETS. TARGETS must
%   not contain the blank index, which is 1 by default.
% 
%   YMASK and TARGETSMASK are dlarray objects, logical arrays, or
%   numeric arrays with the same size as Y and TARGETS, respectively.
%   The function includes and excludes elements of Y and TARGETS for
%   loss computation when the corresponding value in their respective
%   masks are 1 and 0, respectively.
%
%   LOSS is an unformatted dlarray scalar computed as the average CTC
%   loss across the batch dimension.
%
%   LOSS = CTC(_,'PARAM1',VAL1,'PARAM2',VAL2,...) specifies optional
%   parameter name/value pairs:
%
%      'BlankIndex'    - Index of blank character, specified as one of the 
%                        following:
%                            - 'last'  - Use the last element of the
%                                        vocabulary as the blank character.
%                            - integer - Use the character in the 
%                                        vocabulary with the specified 
%                                        index as the blank character. The 
%                                        specified value must be between 1 
%                                        and the number of channels.
%                        The default value is 1.
%      'DataFormat'    - Dimension labels of the input data Y, specified as 
%                        a char array or string. Required if Y is not a 
%                        formatted dlarray.
%
%   LOSS = CTC(Y,TARGETS,YMASK,TARGETSMASK,'DataFormat',FMT) computes the CTC
%   loss of an unformatted dlarray or numeric array Y. FMT specifies the
%   dimension labels of the input data Y. FMT must be a char array or
%   string. For the CTC loss, FMT must contain a dimension labeled 'C', and
%   must not contain a dimension labeled 'S' or a nonsingleton dimension
%   labeled 'U'.
%
%   Example:
%       
%       numChannels = 26;
%       numObservations = 64;
%       lengthPredictions = 6;
%       blankIdx = 1;
%       Y = softmax(dlarray(randn(numChannels, numObservations, lengthPredictions), 'CBT'));
% 
%       % Create a target sequence of indices which doesn't include the 
%       % blank index.
%       lengthTargets = 3;
%       targets = randi([blankIdx+1 numChannels], [numObservations lengthTargets]);
%       
%       % Sequence masks
%       YMask = true(numChannels, numObservations, lengthPredictions);
%       targetsMask = true(numObservations, lengthTargets);
% 
%       LOSS = ctc(Y, targets, YMask, targetsMask);
%
%   See also CROSSENTROPY, L1LOSS, L2LOSS, MSE

%   Copyright 2020-2022 The MathWorks, Inc.

arguments
    X                   {mustBeNumeric, mustBeNonsparse, deep.internal.dlarray.mustBeDlarray, mustBeNonempty}
    T                   {mustBeNumeric, mustBeNonsparse, mustBeInteger, matlab.internal.validation.mustBeMatrix}
    XMask               {deep.internal.dlarray.mustBeBinaryMask}
    TMask               {deep.internal.dlarray.mustBeBinaryMask}
    nvArgs.DataFormat   {deep.internal.dlarray.mustBeDataFormatType}
    nvArgs.BlankIndex   {iMustBeValidBlankIndex} = 1
end

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

% Extract the dimensions of X
[X, numLabels, numObs, maxSeqX] = iValidateXDimension(X);

% Validate and extract the numeric value blankIdx from BlankIndex argument
blankIdx = iValidateBlankIndex(nvArgs.BlankIndex, numLabels);

% Validate the indices T
[T, TMask] = iValidateTargetsAndTMask(T, TMask, numLabels, numObs, blankIdx);

% Validate the masks
XMask = iValidateAndExtractPredictionMask(XMask, numObs, maxSeqX, numLabels, 'XMask');

% Calculate how many repeats
repeats = iCalculateRepeats(T, TMask, numObs);

% Validate the sequence lengths
iValidateSequenceDims(XMask, TMask, repeats);

% Extract formatted data from X (fd.Data is explicitly permuted)
[X,fd] = deep.internal.dlarray.toInternal(X);

% Extract permuted data
[fd, xdata] = extractData(fd);

% Extract the true values from T
tdata = iGetExtractedData(T);

logicalXMask = logical(XMask);
logicalTMask = logical(TMask);

% Overwrite values with mask set to zero. They are disregarded in the loss
% function computation (so their actual value does not count) but we assign
% a new value which doesn't produce any low-level error.
tdata(~logicalTMask) = 1;

xdata = internal_ctc(xdata, tdata, logicalXMask, logicalTMask, blankIdx);

X = deep.internal.dlarray.fromNumeric(xdata, fd, X);

% Remove dimensions from the dlarray
X = stripdims(X);

end

function iMustBeValidBlankIndex(arg)
deep.internal.dlarray.mustBeNonDlarray(arg);
validateattributes(arg,{'numeric', 'char', 'string'},{'nonempty'}, 'ctc', 'Blank index')
end

function dataFormat = iGetDataFormat(nameValueArgs)
if isfield(nameValueArgs, 'DataFormat')
    dataFormat = nameValueArgs.DataFormat;
else
    dataFormat = [];
end
end

function [X, nbLabels, nbObservations, nbMaxSeqX] = iValidateXDimension(X)

labelsX = dims(X);

% Extract the input data
[~,fd] = deep.internal.dlarray.toInternal(X);

% Error out if X does not have 'C'
if ~any(labelsX == 'C')
    error(message('deep:dlarray:MustHaveOneChannelDimension'))
end

if any(labelsX == 'S')
    error(message('deep:dlarray:UnsupportedSpatialDimensions'))
end

if any(labelsX == 'U') && ( size(X, finddatadim(X, 'U')) > 1 )
    error(message('deep:dlarray:NonsingletonUnspecifiedDim'))
end

% Retrieve channel, time and observation dimensions
timeDim = indexOfLabels(fd, 'T');
observationDim = observationIndex(fd);
channelDim = channelIndex(fd);

% Retrieve number of labels
nbLabels = size(X, channelDim);

if ~isempty(observationDim)
    nbObservations = size(X, observationDim);
else
    nbObservations = 1;
end

if ~isempty(timeDim)
    nbMaxSeqX = size(X, timeDim);
else
    nbMaxSeqX = 1;
end

X = dlarray(reshape(X, nbLabels, nbObservations, nbMaxSeqX), 'CBT');

end

function blankIdx = iValidateBlankIndex(value, channelDim)

if isstring(value) || ischar(value)
    validatestring(value, "last");
    blankIdx = channelDim; 
else
    validateattributes(value, {'numeric'}, {'scalar', 'integer', 'positive', '<=', channelDim})
    blankIdx = value;
end

end

function [T, validTMask] = iValidateTargetsAndTMask(T, TMask, numLabels, numObsX, blankIdx)

% Validate T of size [numObsX, maxSeqT] 

TisDlarray = isdlarray(T);
if ~TisDlarray || (TisDlarray && isempty(dims(T)))
    [numObsT, maxSeqT] = size(T);
else
    labelsT = dims(T);

    if ~isequal(labelsT,'BT')
        error(message('deep:dlarray:CtcTargetsNotBT'))
    end
    
    numObsT = size(T, finddim(T,'B'));
    maxSeqT = size(T, finddim(T,'T'));
end

if numObsT ~= numObsX
    error(message('deep:dlarray:LossBatchDimMismatch'))
end

validTMask = iValidateAndExtractTargetMask(TMask, numObsT, maxSeqT, 'TMask');

logicalTMask = logical(TMask);

% Validate every element of T is an integer between 1 and numLabels and
% does not contain blankIdx.
if any( (T < 1) & logicalTMask, 'all' ) || any( (T > numLabels) & logicalTMask, 'all' ) 
    error( message( 'deep:dlarray:CtcIndicesOutOfRange', numLabels ) )
end

if any( T == blankIdx & TMask, 'all' )
    error( message( 'deep:dlarray:CtcIndicesContainBlankIndex', blankIdx ) )
end

% Create dlarray of indices
T = dlarray(T, 'BT');
end

function extractedMask = iValidateAndExtractPredictionMask(mask, numObs, maxSeq, numLabels, maskVariable)

% Extract data if Mask is a dlarray
if isdlarray(mask)
    fmt = dims(mask);
    if ~isempty(fmt) && ~isequal(fmt, 'CBT')
        % Specified mask dimension label mismatches expected dimension
        % label
        error(message('deep:dlarray:InvalidFormattedXMaskArray'));
    end
    extractedMask = iGetExtractedData(mask);
else
    extractedMask = mask;
end

% Validate that Mask is of correct size: a three-dimensional array with
% size numLabels x numObservation x MaxSeq as the input X.
if ~isequal( size(extractedMask, 1:3), [numLabels numObs, maxSeq] ) || ndims(extractedMask) > 3
    error(message('deep:dlarray:InvalidSizeXMaskArray'));
end

% Verify that along the channel dimension, XMask has either all ones or all
% zeros and collapse the channel dimension.
reducedExtractedMask = min(extractedMask,[],1);
constantAlongChannelDim = all( reducedExtractedMask == max(extractedMask,[],1), 'all' );

if ~constantAlongChannelDim
    error(message('deep:dlarray:InvalidXMaskArrayConstantAlongC'))
end

% From Channel x Batch x Time, with singleton Channel dim, to Batch x Time
% array.
extractedMask = shiftdim(reducedExtractedMask, 1);

% Validate that Mask is valid (e.g. Mask = logical([1 1 0 0]) is OK, Mask =
% logical([1 1 0 1]) is invalid)
diffMask = extractedMask(:,2:end) - extractedMask(:,1:end-1);
if any( (diffMask ~= 0) & (diffMask ~= -1), 'all' )
    iThrowInvalidMask(maskVariable);
end

end

function extractedMask = iValidateAndExtractTargetMask(mask, numObs, maxSeq, maskVariable)

% Extract data if Mask is a dlarray
if isdlarray(mask)
    fmt = dims(mask);
    if ~isempty(fmt) && ~isequal(fmt, 'BT')
        % Specified mask dimension label mismatches expected dimension
        % label
        error(message('deep:dlarray:InvalidFormattedTMaskArray'));
    end
    extractedMask = iGetExtractedData(mask);
else
    extractedMask = mask;
end

% Validate that Mask is of correct size
if ~isequal( size(extractedMask), [numObs, maxSeq] )
    iThrowInvalidSizeMask(maskVariable);
end

% Validate that Mask is valid (e.g. Mask = logical([1 1 0 0]) is OK, Mask =
% logical([1 1 0 1]) is invalid)
diffMask = extractedMask(:,2:end) - extractedMask(:,1:end-1);
if any( (diffMask ~= 0) & (diffMask ~= -1), 'all' )
    iThrowInvalidMask(maskVariable);
end

end

function iValidateSequenceDims(XMask, TMask, repeats)
% Extract the sequence lengths
XLens = sum(XMask, 2);
TLens = sum(TMask, 2);

% Verify that XLens is greater than TLens for every observation
if any(TLens + repeats > XLens, 'all' )
    error(message('deep:dlarray:CtcSequenceMismatch'));
end
end

% Calculate number of repeats for each observation in the target sequence T
function repeats = iCalculateRepeats(T, TMask, numObs)

% Apply diff to find adjacent character which are the same
TData = extractdata(T);
repeats = [ ones(numObs, 1), diff(TData, 1, 2) ] == 0;

% Multiply by TMask to sum over only repeats which will affect calculation
repeats = sum( repeats .* TMask, 2 );

end

function extractedData = iGetExtractedData(data)
if isdlarray(data)
    extractedData = deep.internal.dlarray.toNumeric(data);
else
    extractedData = data;
end
end

function iThrowInvalidSizeMask(maskVariable)
if strcmp(maskVariable,"XMask")
    error(message('deep:dlarray:InvalidSizeXMaskArray'));
else
    error(message('deep:dlarray:InvalidSizeTMaskArray'));    
end
end

function iThrowInvalidMask(maskVariable)
if strcmp(maskVariable,"XMask")
    error(message('deep:dlarray:InvalidXMaskArray'));
else
    error(message('deep:dlarray:InvalidTMaskArray'));    
end
end