function [Xpad, mask] = padsequences(X, paddingDimension, options)
%PADSEQUENCES Pad sequence data to the same length
%
%   [XPAD,MASK] = PADSEQUENCES(X,DIM) concatenates the sequences in X and
%   pads them in the dimension specified by DIM. X must be a cell vector
%   containing numeric or categorical arrays. DIM must be a scalar integer.
%   Except the dimension specified by DIM, the size of the arrays in X must
%   be the same in all dimensions.
%
%   Padding is added at the end of each sequence to match the size of the
%   longest sequence in the input X. The default value of padding is 0 for
%   numeric arrays and <undefined> for categorical arrays.
%
%   The output XPAD is an array with N+1 dimensions where N is the number
%   of dimensions of the arrays in X. The last dimension of XPAD has the
%   same size as the number of sequences in input X. The output XPAD is
%   returned with the same underlying type as the arrays in input X.
%
%   The output MASK is a logical array with the same size as XPAD. The
%   position of values of true or 1 in MASK correspond to the positions of
%   original sequence data in XPAD and values of false or 0 correspond to
%   padded values.
%
%   [___] = PADSEQUENCES(___,'PARAM',VAL) specifies optional parameter
%   name/value pairs:
%
%       'Length'        - Strategy to determine the length of the output
%                         sequences. Supported values are:
%                           - 'longest'. Pad the input sequences to the
%                           length of the longest sequence.
%                           - 'shortest'. Truncate the input sequences
%                           to the length of the shortest sequence.
%                           - Positive integer. Pad or truncate the
%                           input sequences to the specified size.
%                          The default value is 'longest'.
%
%       'Direction'     - Direction of padding or truncation. Supported
%                         values are:
%                           - 'right'. Add padding or truncate at the end
%                           of the original sequences.
%                           - 'left'. Add padding or truncate at the
%                           beginning of the original sequences.
%                           - 'both'.  Center the original sequences and
%                           pad or truncate evenly at both sides. Apply
%                           half the amount of padding or truncation to
%                           both ends of the sequence.
%                         The default value is 'right'.
%
%       'PaddingValue'  - Value to use for padding. Supported values are:
%                           - 'auto'. Infer padding value from input
%                           sequences. Padding value is 0 for numeric
%                           sequences, and <undefined> for categorical
%                           sequences.
%                           - 'symmetric'. Pad sequences with mirror
%                           reflections of themselves.
%                           - Scalar value. Pad sequence with the specified
%                           value. Padding value must be numeric or 
%                           categorical.
%                         The default value is 'auto'.
%
%       'UniformOutput' - Logical scalar value indicating whether to return
%                         the sequences in a uniform array or in a cell
%                         array. The default value is true.
%
%   Example 1
%       % Load sequence data
%       S = japaneseVowelsTrainData;
%       timeDimension = 2;
%
%       % Pad the first 10 sequences and retrieve the mask
%       [xPadded,mask] = padsequences(S(1:10),timeDimension);
%
%   Example 2
%       % Load sequence data
%       S = japaneseVowelsTrainData;
%       timeDimension = 2;
%
%       % Pad the first 10 sequences to a length of 22 using a padding
%       % value of -1. Add padding in front of each sequence.
%       sequenceLength = 22;
%       paddingValue = -1;
%       [xPaddedFixed,mask] = padsequences(S(1:10),timeDimension,...
%           "Direction","left",...
%           "Length",sequenceLength,...
%           "PaddingValue",paddingValue);
%
%   Example 3
%       % Use padsequences to transform a datastore containing sequences.
%       % Pad or truncate each sequence such that all sequences have a size
%       % of 17. Add padding on both sides of each sequence, using the
%       % symmetric reflection of the original sequence as the padding
%       % value.
%
%       % Load sequence data
%       S = japaneseVowelsTrainData;
%
%       % Create an arrayDatastore
%       ds = arrayDatastore(S,"OutputType","same","ReadSize",4);
%
%       % Transform the datastore
%       tds = transform(ds, @(x) padsequences(x,2,...
%               "Direction","both",...
%               "Length",17,...
%               "PaddingValue","symmetric",...
%               "UniformOutput",false));
%
%       % Read the transformed datastore to obtain 4 padded sequences
%       paddedSequences = read(tds);
%
%
% See also minibatchqueue, sequenceInputLayer, dlarray, lstm, gru

%   Copyright 2020-2022 The MathWorks, Inc.

arguments
    X cell {mustBeVector}
    paddingDimension (1,1) double {mustBeInteger,mustBePositive}
    options.Direction (1,1) deep.internal.data.padsequences.DirectionEnum = "right"
    options.Length (1,:) {deep.internal.dlarray.mustBeNonDlarray} = iGetLongestSupportedValue()
    options.PaddingValue (1,:) = "auto"
    options.UniformOutput (1,1) {iAssertBinary} = true
end

% Convert arguments
sequenceLength = iValidateLength(options.Length);
uniformOutput = options.UniformOutput;
direction = options.Direction;
exampleInput = iGetExampleFromX(X);

% Validate the padding value and find the padding mode
[paddingValue, paddingMode] = iValidatePaddingValueAndGetPaddingMode(options.PaddingValue);

% Find the output sequence length and ensure the input is supported
[targetLength, maxNDimsInput] = iParseInputSequencesAndFindTargetLength(X, exampleInput, sequenceLength, paddingDimension);
numObservations = numel(X);

% Initialize mask if asked for
computeMask = nargout>1;
mask = [];

if uniformOutput
    % Initialize the output assuming all the sequences have the same number of
    % dimensions. Indexing will error later if this is not the case
    [Xpad, nDimsOutput] = iGenerateOutput(targetLength, paddingDimension, paddingValue, paddingMode, exampleInput, numObservations);
    if computeMask
        mask = false(size(Xpad));
    end
else
    nDimsOutput = max(maxNDimsInput,paddingDimension);
    Xpad = cell(size(X));
    if computeMask
        mask = cell(size(Xpad));
    end
end

% Initialize indexing
[inputIndexing, outputIndexing] = iPreallocateIndexing(maxNDimsInput, paddingDimension, nDimsOutput);

if uniformOutput
    % With UniformOutput true, capture indexing error and suggest to use
    % UniformOutput false instead
    try
        [Xpad, mask] = iPadSequencesUniform(X, Xpad, targetLength, direction, ...
            paddingDimension, paddingMode, numObservations, inputIndexing, outputIndexing, computeMask, mask);
    catch me
        subassignErrors = ["MATLAB:subsassigndimmismatch" ...
            "parallel:array:SubsasgnTooManyRHSDims"];
        if any(me.identifier == subassignErrors)
            % Throw a better error if the indexing failed because of
            % inconsistent dimensions
            me = MException(message("deep:padsequences:InconsistentDims"));
        end
        throw(me)
    end
else
    [Xpad, mask] = iPadSequencesNonUniform(X, Xpad, targetLength, direction, ...
        paddingDimension, paddingValue, paddingMode, numObservations, inputIndexing, outputIndexing, computeMask, mask);
end
end

% Input validation helpers

function sequenceLength = iValidateLength(sequenceLength)
% Validate Length is positive integer scalar, 'longest' or 'shortest'
if isnumeric(sequenceLength)
    validateattributes(sequenceLength, {'numeric'}, {'integer', 'scalar', 'positive'}, 'padsequences', 'Length')
elseif ischar(sequenceLength) || isstring(sequenceLength)
    longestValue = iGetLongestSupportedValue();
    shortestValue = iGetShortestSupportedValue();
    n = strlength(sequenceLength);
    if strncmp(sequenceLength, longestValue,n)
        sequenceLength = longestValue;
    elseif strncmp(sequenceLength, shortestValue,n)
        sequenceLength = shortestValue;
    else
        % Incorrect string. Use validatestring to throw a nice error
        sequenceLength = validatestring(sequenceLength,{longestValue,shortestValue});
    end
else
    error(message("deep:padsequences:InvalidOutputSize"))
end
end

function longestValue = iGetLongestSupportedValue()
longestValue = 'longest';
end

function shortestValue = iGetShortestSupportedValue()
shortestValue = 'shortest';
end

function exampleX = iGetExampleFromX(X)
% Grab the first sequence, to use as example for allocating the output
if ~isempty(X)
    exampleX = X{1};
else
    % No sequence provided
    exampleX = [];
end
end


function [paddingValue, paddingMode] = iValidatePaddingValueAndGetPaddingMode(paddingValue)
% Validate the input paddingValue and find the padding mode to use

% Validate PaddingValue is scalar, 'auto' or 'symmetric'
if ischar(paddingValue) || isstring(paddingValue)
    try
        paddingMode = deep.internal.data.padsequences.PaddingValueEnum(paddingValue);
    catch me
        if me.identifier == "MATLAB:class:CannotConvert"
            % Incorrect string. Use validatestring to throw a nice error
            iValidatePaddingValueString(paddingValue);
        else
            throw(me)
        end
    end
    if paddingMode == deep.internal.data.padsequences.PaddingValueEnum.custom
        % User selected 'custom', this is not supported
        iValidatePaddingValueString(paddingValue);
    end
    % Set the padding value to empty as it will be evaluated
    % later
    paddingValue = [];
else
    % Custom padding was specified
    paddingMode = deep.internal.data.padsequences.PaddingValueEnum.custom;
    validateattributes(paddingValue, {'numeric','categorical'}, {'scalar'}, 'padsequences', 'PaddingValue');
end
end

function paddingValue = iValidatePaddingValueString(paddingValue)
% Validate that the string padding value is one of the supported values
paddingValue = validatestring(paddingValue,{'auto','symmetric'});
end

% Padding helpers

function [inputIndexing, outputIndexing] = iPreallocateIndexing(maxNDimsInput, paddingDimension, nDimsOutput)
% Pre-allocate arrays for indexing
inputIndexing = iInitializeSubsrefIndices(max(maxNDimsInput,paddingDimension));
outputIndexing = iInitializeSubsrefIndices(max(nDimsOutput,paddingDimension));
end

function indices = iInitializeSubsrefIndices(numDims)
% Return a cell array with values ':'
indices = cell(1,numDims);
indices(:) = {':'};
end

function [targetLength, maxNDimsInput] = iParseInputSequencesAndFindTargetLength(input, exampleInput, outputSize, paddingDimension)
% Perform a single pass on the input sequences. During this pass,
% verification on the sequences are done. We assert that the input data is
% supported and the size and number of dimensions of the sequences is
% computed.


% Remove the stack trace
try
    % Assert the data type is supported and find the length of each
    % sequence
    [sequenceSizes, nDims] = iParseSequenceAssertDataTypeAndGetMetaData(input, exampleInput, paddingDimension);
    if isnumeric(outputSize)
        targetLength = outputSize;
    else
        if ~isempty(input)
            % Length specified as "longest" or "shortest"
            if strcmp(outputSize, iGetLongestSupportedValue())
                targetLength = max(sequenceSizes);
            elseif strcmp(outputSize, iGetShortestSupportedValue())
                targetLength = min(sequenceSizes);
            end
        else
            % No sequence provided
            targetLength = 0;
        end
    end
catch me
    throwAsCaller(me)
end

maxNDimsInput = max(nDims);
end

function [sz, nDims] = iParseSequenceAssertDataTypeAndGetMetaData(x, exampleInput, paddingDimension)
% Parse the sequences to verify the type and compute the sequence size and
% number of dimensions
numSeq = numel(x);
sz = zeros(numSeq,1);
nDims = zeros(numSeq,1);

% Collect info on the example sequence
exampleInputClass = class(exampleInput);
if iscategorical(exampleInput)
    exampleInputCategories = categories(exampleInput);
else
    exampleInputCategories = [];
end

% Parse the sequences
for i = 1:numSeq
    [sz(i), nDims(i)] = iAssertDataTypeAndGetMetaData(x{i}, exampleInputClass, exampleInputCategories, paddingDimension);
end
end

function [sz, nDims] = iAssertDataTypeAndGetMetaData(sequence, exampleInputClass, exampleInputCategories, paddingDimension)
isCategoricalSequence = iscategorical(sequence);
% Only numeric and categorical sequences are supported
if ~isnumeric(sequence) && ~isCategoricalSequence
    % The input array is not supported
    error(message("deep:padsequences:InvalidInputType"))
end
if ~isa(sequence, exampleInputClass)
    % The sequences do not have the same data type
    error(message("deep:padsequences:MismatchingType"))
end
if isCategoricalSequence && ~isequal(categories(sequence), exampleInputCategories)
    % The categorical sequences do not have the same categories
    error(message("deep:padsequences:MismatchingCategories")) 
end
% Compute nDims and size
nDims = ndims(sequence);
sz = size(sequence, paddingDimension);
end

function [array, nDimsArray] = iGenerateOutput(targetLength, paddingDim, paddingValue, paddingMode, input, numObservations)
% Allocate the output array using the padding value and the selected
% padding mode

% Compute the output array size
outputArraySize = iFindPaddedOutputSize(input, targetLength, numObservations, paddingDim);

% Find the output padding value
outputPaddingValue = iFindOutputPaddingValue(paddingValue, paddingMode, input);

% Create the padded array depending on the input type
if isnumeric(input)
    array = ones(outputArraySize, 'like', input).*cast(outputPaddingValue,'like',input);
else
    array = repmat(outputPaddingValue, outputArraySize);
end

if nargout > 1
    % Return the number of dimension using numel to preserve singleton
    % dimensions
    nDimsArray = numel(outputArraySize);
else
    nDimsArray = [];
end
end

function outputSize = iFindPaddedOutputSize(input, targetLength, numObservations, paddingDimension)
% Compute the output size from the input dimensions

inputSize = size(input);
nDimsInput = numel(inputSize);

% Replace empty dimensions with 1 to ensure an empty sequence get filled
% with the padding value
inputSize(inputSize==0) = 1;

% Padding dimension is not present in the original array, add enough
% singleton dimensions
if paddingDimension > nDimsInput
    inputSize = [inputSize ones(1,paddingDimension-nDimsInput)];
end
inputSize(paddingDimension) = targetLength;

outputSize = [inputSize, numObservations];
end

function outputPaddingValue = iFindOutputPaddingValue(paddingValue, paddingMode, input)
% Find the output padding value to use depending on the padding mode that
% was selected

useCustomPaddingValue = paddingMode == deep.internal.data.padsequences.PaddingValueEnum.custom;

if useCustomPaddingValue
    % With custom padding, attempt to cast the padding value to the output
    % type
    if iscategorical(input)
        castFcn = @(x) categorical(x);
    else
        castFcn = @(x) cast(x,'like',input);
    end
    try
        outputPaddingValue = castFcn(paddingValue);
    catch c
        me = MException(message("deep:padsequences:InvalidPaddingValue"));
        me = addCause(me,c);
        throw(me)
    end
else
    % With auto and other string modes, use 0 or missing as padding value
    if iscategorical(input)
        % Create a categorical with value "missing" and the original
        % classes. 
        outputPaddingValue = [input([]) missing];
    else
        % The value will be cast to the target precision when indexing
        % the output
        outputPaddingValue = 0;
    end
end
end

function [Xpad, mask] = iPadSequencesUniform(X, Xpad, targetLength, direction, ...
    paddingDim, paddingMode, numObservations, inputIndexing, outputIndexing, computeMask, mask)
% Padding function for a non-cell array output
for i = 1:numObservations
    currentSequence = X{i};
    outputIndexing{end} = i;
    
    % Fill the pre-allocated input with expected data and padding
    [Xpad, maskIndexing] = iPadSequence(currentSequence, Xpad, paddingDim, targetLength, direction, ...
        paddingMode, inputIndexing, outputIndexing);
    
    % Generate the mask using the output indexing
    if computeMask
        mask(maskIndexing{:}) = true;
    end
end
end

function [Xpad, mask] = iPadSequencesNonUniform(X, Xpad, targetLength, direction, ...
    paddingDim, paddingValue, paddingMode, numObservations, inputIndexing, outputIndexing, computeMask, mask)
% Padding function for a cell array output
for i = 1:numObservations
    currentSequence = X{i};
    
    % Allocate the output sequence on the fly to allow variable size
    % sequences
    padSeq = iGenerateOutput(targetLength, paddingDim, paddingValue, paddingMode, currentSequence, 1);
    
    % Fill the pre-allocated input with expected data and padding
    [padSeq, maskIndexing] = iPadSequence(currentSequence, padSeq, paddingDim, targetLength, direction, ...
        paddingMode, inputIndexing, outputIndexing);
    
    % Put the padded sequence in the output cell array
    Xpad{i} = padSeq;
    
    % Generate the mask using the output indexing
    if computeMask
        padMask = false(size(padSeq));
        padMask(maskIndexing{:}) = true;
        mask{i} = padMask;
    end
end
end

function [paddedSequence, maskIndexing] = iPadSequence(currentSequence, paddedSequence, dim, ...
    targetLength, direction, paddingMode, inputIndexing, outputIndexing)
% padSequence  Compute the padding for a single sequence. Add the
% currentSequence in the paddedSequence according to the other options. If
% padding mode is not uniform, generate also the padding from the
% currentSequence.

% Ensure an empty sequence has a zero length
if isempty(currentSequence)
    currentSequenceLength = 0;
else
    currentSequenceLength = size(currentSequence,dim);
end

% Offset with the targetLength determines if cropping or padding
offset = targetLength - currentSequenceLength;

% Get the indices for this sequence
if offset < 0
    % Need to crop the sequence
    maskIndices = 1:targetLength;
    outputIndexing{dim} = ':';
    switch direction
        case deep.internal.data.padsequences.DirectionEnum.right
            inputIndexing{dim} = maskIndices;
        case deep.internal.data.padsequences.DirectionEnum.left
            % Left padding, offset the indices to start from the right-most
            % side of the sequence
            inputIndexing{dim} = maskIndices - offset;
        case deep.internal.data.padsequences.DirectionEnum.both
            % Edge padding, offset the indices by half the offset
            inputIndexing{dim} = maskIndices - floor(offset/2);
    end
elseif offset > 0
    % Need to pad the sequence
    maskIndices = 1:currentSequenceLength;
    switch direction
        case deep.internal.data.padsequences.DirectionEnum.right
            % No-op
        case deep.internal.data.padsequences.DirectionEnum.left
            % Left padding, offset the indices to start from the right-most
            % side of the sequence
            maskIndices = maskIndices + offset;
        case deep.internal.data.padsequences.DirectionEnum.both
            % Edge padding, offset the indices by half the offset
            maskIndices = maskIndices + floor(offset/2);
    end
    % Get the padding indexing depending on the padding mode
    switch paddingMode
        case deep.internal.data.padsequences.PaddingValueEnum.symmetric
            [inputIdx, outputIdx] = iGetSymmetricPaddingIndices(currentSequenceLength, offset, direction);
        otherwise
            % Uniform padding
            [inputIdx, outputIdx] = iGetUniformPaddingIndices(maskIndices);
    end
    inputIndexing{dim} = inputIdx;
    outputIndexing{dim} = outputIdx;
else
    % The sequence is already of size targetLength
    maskIndices = 1:targetLength;
    inputIndexing{dim} = ':';
    outputIndexing{dim} = ':';
end

% Get the mask indexing
maskIndexing = outputIndexing;
maskIndexing{dim} = maskIndices;

% Return the data for this sequence
currentSequenceData = currentSequence(inputIndexing{:});

paddedSequence(outputIndexing{:}) = currentSequenceData;
end

function [inputIdx, outputIdx] = iGetSymmetricPaddingIndices(currentSequenceLength, offset, direction)

if currentSequenceLength == 0
    % If the sequence is empty, avoid indexing
    inputIdx = [];
    outputIdx = [];
    return
end

% Generate the symmetric indices for this sequence
symmetricMask = uint32([1:currentSequenceLength currentSequenceLength:-1:1]);

% Depending on the direction, generate the indices we'll use
switch direction
    case deep.internal.data.padsequences.DirectionEnum.left
        idx = -offset:currentSequenceLength-1;
        
    case deep.internal.data.padsequences.DirectionEnum.right
        idx = 0:currentSequenceLength+offset-1;
        
    case deep.internal.data.padsequences.DirectionEnum.both
        idx = -floor(offset/2):currentSequenceLength+ceil(offset/2)-1;
end
% Index the symmetric indices to obtain an indexing valid for
% the sequence
inputIdx = symmetricMask(mod(idx, 2*currentSequenceLength)+1);

% With symmetric padding, we generate as many indices as the size of the
% output. This will repeat values of the input, we need to index the entire
% padded sequence.
outputIdx = ':';
end

function [inputIdx, outputIdx] = iGetUniformPaddingIndices(maskIndices)
% With Uniform output, we need to take the entire original sequence
inputIdx = ':';

% With uniform padding, we just need to index the original sequence, since
% the output is initialized to the padding value.
outputIdx = maskIndices;
end

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