function varargout = buffer(x, n, varargin)
%BUFFER Buffer a signal vector into a matrix of data frames on the GPU.
%	Y = BUFFER(X,N)
%	Y = BUFFER(X,N,P)
%	Y = BUFFER(X,N,P,OPT)
%	Y = BUFFER(X, ...)
%	[Y,Z] = BUFFER(X, ...)
%	[Y,Z,OPT] = BUFFER(X, ...)
%
%   See also: BUFFER, GPUARRAY.

%   Copyright 2020-2023 The MathWorks, Inc.

narginchk(2, 4);
nargoutchk(0, 3);

if ~isgpuarray(x)
    % Only extra arguments are on the GPU, gather extra arguments and
    % dispatch in-memory buffer.
    n = gather(n);
    if nargin > 2
        [varargin{:}] = gather(varargin{:});
    end
    varargout = cell(1, nargout);
    [varargout{:}] = buffer(x, n, varargin{:});
    return;
end

[x, n, p, opt, isOptNodelay] = iParseAndValidateInputs(x, n, varargin{:});

% Here we only handle column vectors, keep track if input x was a row
% vector.
isRowX = isrow(x);
if isRowX
    x = x.';
end

if p >= 0
    [varargout{1:nargout}] = iBufferOverlap(x, n, p, opt, isOptNodelay, isRowX);
else
    [varargout{1:nargout}] = iBufferUnderlap(x, n, p, opt);
end

function [x, n, p, opt, isOptNodelay] = iParseAndValidateInputs(x, n, varargin)
% Parse and validate inputs for buffer.

% Validate input signal x: non-sparse, numeric/logical, vector. If it's on
% the GPU, it must be numeric or logical. It can't have a different
% datatype.
if issparse(x)
    error(message("signal:buffer:InputMustBeNonSparse"));
end

if ~isvector(x)
    error(message("signal:buffer:InputMustBeVector"));
end

% Validate buffer size n: numeric, real, scalar, integer-valued, positive
if ~isUnderlyingType(n, "numeric") || ~isreal(n) || ~isscalar(n) || n ~= floor(n) || n < 0
    error(message("signal:buffer:InvalidBufferSize"));
end
n = double(n);

% Validate buffer overlap p: numeric, real, scalar, nonnegative,
% interger-valued p must be smaller than n
if nargin > 2
    p = varargin{1};
    if ~isUnderlyingType(p, "numeric") || ~isreal(p) || ~isscalar(p) || p ~= floor(p)
        error(message("signal:buffer:InvalidBufferOverlap"));
    end
    if p > n-1
        error(message("signal:buffer:PMustBeSmallerThanN"));
    end
else
    p = 0;
end
p = double(p);

% Validate optional argument OPT:
isOptNodelay = false;
if nargin > 3
    opt = varargin{2};
    if p < 0
        % Underlap (p<0): must be empty or double real integer-valued
        % scalar in the range [0 -p].
        if numel(opt)>1 || ~isUnderlyingType(opt, "numeric") || ~isreal(opt)
            error(message("signal:buffer:InvalidOptForUnderlap"));
        end
        if ~isempty(opt)
            if opt < 0 || opt > -p || opt ~= floor(opt)
                error(message("signal:buffer:InvalidSkipOffsetForUnderlap"));
            end
        end
        opt = double(opt);
    elseif p >= 0
        % Overlap (p>0): must be empty, 'nodelay' or vector of length p of
        % the same datatype as x (no complexity check). If it's 'nodelay',
        % length(x)>p.
        % No-overlap (p==0): must be empty or 'nodelay'. If it's 'nodelay',
        % length(x)>p.
        if matlab.internal.datatypes.isScalarText(opt, false)
            if startsWith("nodelay", opt)
                if length(x) < p
                    error(message("signal:buffer:InvalidInputLengthForOverlapNodelay"));
                end
                isOptNodelay = true;
            else
                error(message("signal:buffer:InvalidOptForOverlap"));
            end
        else
            if p == 0
                if ~isempty(opt)
                    error(message("signal:buffer:InvalidOptForNoOverlap"));
                end
            end
            
            if ~isempty(opt) && (~isvector(opt) || length(opt) ~= p)
                error(message("signal:buffer:InvalidInitialConditions"));
            end
            
            if ~isUnderlyingType(x, underlyingType(opt))
                error(message("signal:buffer:InvalidInitialConditionsDatatype"));
            end
        end
    end
else
    opt = [];
end

function varargout = iBufferOverlap(x, n, p, opt, isOptNodelay, isRowX)
% Perform buffering with no overlap between frames or overlapping frames.

if nargout < 2
    % Y contains all the frames including any incomplete frame at the end.
    % Pre-padding might be considered depending on the input arguments.
    if p > 0 && ~isOptNodelay && ~isempty(opt)
        % Overlap case when the user specifies the samples to use for
        % pre-padding. In this case, opt is a vector of p samples.
        prePad = reshape(opt, [], 1);
        x = [prePad; x];
        % We've done pre-padding, treat the entire input as in the
        % 'nodelay' case by setting the last input to true.
        varargout{1} = parallel.internal.gpu.extractIncompleteWindowsWithPadding(x, n, p, true);
    else
        % Default syntax with isOptNodelay set to false, pre-padding with
        % zeros is added. If the user has specified 'nodelay', no
        % pre-padding is added.
        varargout{1} = parallel.internal.gpu.extractIncompleteWindowsWithPadding(x, n, p, isOptNodelay);
    end
else
    % Y only contains full frames.
    lengthX = length(x);
    if (p > 0 && ~isOptNodelay && ~isempty(opt))
        % opt is a vector of p samples
        prePad = reshape(opt, [], 1);
        x = [prePad; x];
        % We've done pre-padding, treat the entire input as in the
        % 'nodelay' case setting the last input to true.
        varargout{1} = parallel.internal.gpu.extractWindows(x, n, p);
    elseif isOptNodelay
        varargout{1} = parallel.internal.gpu.extractWindows(x, n, p);
    else
        % Default case, add prepadding and return full frames only.
        varargout{1} = parallel.internal.gpu.extractWindowsWithPrePadding(x, n, p);
    end
    
    % Z contains unused samples form x. Extract discarded samples from x.
    if isOptNodelay
        L = max(lengthX - p, 0); % Valid for p>=0
    else
        L = lengthX;
    end
    numFullFrames = fix(L/(n-p));
    totalSamplesUsed = numFullFrames*(n-p) + p*isOptNodelay;
    numDiscardedSamples = lengthX - totalSamplesUsed;
    varargout{2} = tail(x, numDiscardedSamples);
    % If X was a row, return the discarded samples as a row vector too.
    if isRowX
        varargout{2} = varargout{2}.';
    end
    
    if nargout == 3
        % OPT contains the last p samples that would be required for
        % the next frame.
        if p == 0
            varargout{3} = gpuArray([]);
        else
            if isOptNodelay || isempty(opt)
                % Do not consider prepadding samples
                startOptOffset = p - 1;
            else
                % User has provided opt as a vector of initial conditions
                totalSamplesUsed = totalSamplesUsed + length(opt);
                startOptOffset = p - 1;
            end
            idxOpt = colon(totalSamplesUsed - startOptOffset, totalSamplesUsed);
            if any(idxOpt < 1)
                % There are not enough samples to fill in the first frame,
                % opt is set to zeros as the default behaviour.
                % Create a prototype of the input datatype.
                xProto = x([]);
                varargout{3} = zeros(length(idxOpt), 1, 'like', xProto);
            else
                varargout{3} = x(idxOpt);
            end
        end
    end
end

function varargout = iBufferUnderlap(x, n, p, opt)
% Perform buffering with underlap (p < 0) to skip samples between frames.

% If opt is specified, it's an integer-valued scalar specifying how many
% samples to skip at the beginning of the input.
if ~isempty(opt)
    x = tail(x, max(0, length(x)-opt));
end
L = length(x);

if nargout < 2
    % Y contains all the frames including any incomplete frame at the end
    % with padded zeros. For all underlap cases, opt can't be set to
    % 'nodelay'.
    varargout{1} = parallel.internal.gpu.extractIncompleteWindowsWithPadding(x, n, p, false);
else
    % Y only contains full frames and Z contains unused samples from x.
    % Extract windows only returns full frames
    varargout{1} = parallel.internal.gpu.extractWindows(x, n, p);
    % Now extract discarded samples from x.
    numFullFrames = round(L/(n-p));
    if numFullFrames == 0 && L>n
        % We have enough samples to fill the frame of length n but not
        % enough to skip all the requested samples in p.
        numFullFrames = 1;
    end
    totalSamplesUsed = numFullFrames*(n-p) + p;
    % Take into account the following -p samples to skip, they don't form
    % part of z.
    numSkipSamples = L - totalSamplesUsed + p;
    varargout{2} = tail(x, max(0, numSkipSamples));
    % For underlap, always return the discarded samples as a row vector.
    varargout{2} = varargout{2}.';
    
    if nargout > 2
        % OPT returns the number of samples that must be skipped for the
        % next frame.
        if isempty(varargout{2})
            numUnusedSamples = L - totalSamplesUsed;
            varargout{3} = gpuArray(-p - numUnusedSamples);
        else
            % When there are samples in z (second output argument) that
            % weren't used for the first output, opt is equal to 0. There
            % is no need to skip samples for the next frame.
            varargout{3} = gpuArray(0);
        end
    end
end