function [S,F,T] = dlstft(x,timeValue,NameValueArgs)
%DLSTFT Deep learning short-time Fourier transform.
%   Y = DLSTFT(X) returns the one-sided short-time Fourier transform (STFT)
%   of X. The STFT is computed over the frequency range [0,pi]
%   radians/sample using a default 128-sample Hann window and an equal
%   number of discrete Fourier transform (DFT) bins. The number of
%   overlapped samples between windows is 96. X is a real-valued numeric
%   array, gpuArray, unformatted dlarray, or a dlarray in 'CBT' format. The
%   underlying data type of X is double or single. The STFT of X is in
%   general complex valued. If X is a formatted dlarray, Y is a dlarray with
%   'SCBT' format. The 'S' dimension corresponds to frequency in the STFT.
% 
%   DLSTFT computes the short-time Fourier transform of X along its 'T'
%   dimension. Ensure X has enough samples in the 'T' dimension to match
%   the window length. If X has fewer than 128 samples in the 'T' dimension,
%   specify a smaller window.
% 
%   Y = DLSTFT(X,Fs) specifies the sample rate of X, Fs, in hertz as a
%   positive scalar.
%
%   Y = DLSTFT(X,Ts) specifies the sample time, Ts, as a positive
%   duration scalar. The sample rate in hertz is computed as the inverse of
%   Ts with Ts converted to seconds.
%
%   Y = DLSTFT(...,DataFormat=FMT) specifies the data format, FMT, if X
%   is an unformatted dlarray, numeric array, or gpuArray. FMT is any
%   permutation of 'CBT'. 'DataFormat' must be specified if X is an
%   unformatted dlarray, a numeric array, or a gpuArray. Do not specify
%   'DataFormat' if X is a formatted dlarray. If you specify the data
%   format, Y is an unformatted dlarray.
%
%   Y = DLSTFT(...,Window=WINDOW) specifies the window used to compute
%   the STFT. WINDOW must be a real-valued numeric, gpuArray, or dlarray
%   vector. DLSTFT divides each signal in X along the 'T' dimension into
%   segments of the same length as WINDOW with an overlap between adjoining
%   segments set to the largest integer less than or equal to 75% of the
%   window length. DLSTFT multiplies each segment elementwise with the
%   vector specified in WINDOW. The default is a Hann window of length 128.
%   X must have at least as many samples along the 'T' dimension as the
%   length of the window.
%
%   Y = DLSTFT(...,OverlapLength=NOVERLAP) specifies an integer
%   number of samples of overlap between adjoining segments. NOVERLAP must
%   be smaller than the length of the window. If NOVERLAP is not specified,
%   the overlap is set to the largest integer less than or equal to 75% of
%   the window length.
%
%   Y = DLSTFT(...,FFTLength=NFFT) specifies the integer number of
%   frequency points used to calculate the DFT. NFFT must be greater than
%   or equal to the window length. NFFT defaults to the length of WINDOW.
%   The STFT is computed over the frequency range [0,pi] or [0,Fs/2] when
%   NFFT is even and [0,pi) or [0,Fs/2) when NFFT is odd. The 'S' dimension
%   of Y corresponds to the DFT bins in the one-sided STFT.
%
%   [Y,F] = dlstft(...) returns the frequencies at which the STFT is
%   evaluated. If the input contains a sample rate or a sample time, F has
%   units of hertz. Otherwise, F has units of radians/sample. F is an
%   unformatted dlarray.
%
%   [Y,F,T] = dlstft(...) returns the times at which the STFT is
%   evaluated. If a sample rate is provided, T is an unformatted dlarray of
%   time values in seconds. If a sample time is provided, T is a duration
%   array with the same time format as the input. If no time information is
%   provided, the output is an unformatted dlarray of sample numbers.
%
%     % EXAMPLE 1:
%     %     Compute the deep learning short-time Fourier transform of a
%     %     'CTB' formatted dlarray with 3 channels and 1000 samples
%     %     representing a batch of 5.
%     X = dlarray(rand([3 1000 5]),"CTB");
%     YDL = dlstft(X);
%
%     % EXAMPLE 2:
%     %     Create an unformatted dlarray containing a batch of two
%     %     single-channel signals. Each signal in the batch contains a
%     %     quadratic chirp.
%     fs = 8000;
%     t = 0:1/fs:2-1/fs;
%     X = dlarray(zeros([16000 1 2]));
%     X(:,:,1) = chirp(t,250,1,500,"q");
%     X(:,:,2) = chirp(t,500,1,1000,"q");
%     %     Set the FFT length, the window, and the overlap length.
%     nfft = 600;
%     win = hann(512);
%     noverlap = 128;
%     %     Compute the deep learning short-time Fourier transform.
%     %     Specify the format of the input dlarray as 'TCB'.
%     [YDL,FDL,TDL] = dlstft(X,fs,Window=win,OverlapLength=noverlap,...
%       FFTLength=nfft,DataFormat="TCB");
%
%   See also STFTLAYER, STFT

% Copyright 2020-2023 The MathWorks, Inc.

arguments
    x dlarray {mustBeFloat,mustBeNonempty,mustBeReal}
    timeValue = []
    NameValueArgs.Window dlarray {mustBeFloat,mustBeNonempty,mustBeFinite,...
        mustBeReal,mustBeVector}
    NameValueArgs.OverlapLength = []
    NameValueArgs.FFTLength = []
    NameValueArgs.DataFormat {mustBeTextScalar} = ''
end

% Check outputs
nargoutchk(0,3);

%---------------------------------
% Validate input and further parse inputs
[x,opts] = dlstftParser(x,timeValue,NameValueArgs);
win = opts.Window;
noverlap = opts.OverlapLength;
fs = opts.EffectiveFs;
nfft = opts.FFTLength;
% timeunits = opts.TimeUnits;
addformat = isempty(opts.DataFormat);

if nargout < 2 % When no need T output
    S = signalwavelet.deep.internal.stft.computeDLSTFT(x,win,noverlap,nfft,addformat,fs);
else
    [S,T] = signalwavelet.deep.internal.stft.computeDLSTFT(x,win,noverlap,nfft,addformat,fs);
end

if nargout > 1
    if opts.IsNormalizedFreq
        F = psdfreqvec('npts',nfft,'Range','half');
    else
        F = psdfreqvec('npts',nfft,'Fs',fs,'Range','half');
    end
    F = cast(F,'like',real(S));
end

if nargout > 2
    if isnumeric(T) && ~isempty(opts.TimeUnits)
        T = duration(0,0,T,'Format',opts.TimeUnits);
    else
        T = cast(T,'like',real(S));
    end
    % Ensure time output is a column vector
    T = T(:);
end

end

function [x,opts] = dlstftParser(x,timeValue,NameValueArgs)
%DLSTFTPARSER is a function for parsing value-only inputs, flags, and
%   name-value pairs for the DLSTFT functions.

if ~isfield(NameValueArgs,'Window')
    win = dlarray(hann(128,'periodic'));
    nwin = 128;
else
    win = stripdims(NameValueArgs.Window);
    nwin = length(win);
end

%---------------------------------
% Set defaults
timeUnits = '';
defaultOverlapPercent = 0.75;
timeMode = 'sp';
effectiveFs = 1;
isNormalizedFreq = true;

%---------------------------------
% Validate input and compute parameters
x = deep.internal.dlarray.validateDataFormatArg(x, NameValueArgs.DataFormat);

if ~strcmpi(dims(x),'CBT')
    error(message('signal_deep:dlstft:InvalidDataFormat'));
end

dataLength = size(x,3);
numChannel = size(x,1);
batchSize = size(x,2);
x = stripdims(x);

% The following is not needed when dlstft is called by stftLayer Determine
% time type
if ~isempty(timeValue)
    if isduration(timeValue)
        timeUnits = timeValue.Format;
        if isscalar(timeValue)
            timeValue = seconds(timeValue);
            timeMode = 'ts';
        else
            error(message('signal:stft:InvalidValueOnlyInput'));
        end
    else
        if isscalar(timeValue)
            timeMode = 'fs';
        else
            error(message('signal:stft:InvalidValueOnlyInput'));
        end
    end
end
% Validate the time value
switch timeMode
    case 'fs' % Sample rate provided
        validateattributes(timeValue, {'numeric'},{'scalar','real',...
            'finite','positive'},'DLSTFT','sample rate')
        effectiveFs = double(timeValue);
        isNormalizedFreq = false;
    case 'ts' % Duration provided
        validateattributes(timeValue, {'numeric'},{'scalar','real',...
            'finite','positive'},'DLSTFT','sample time')
        effectiveFs = double(1/timeValue);
        isNormalizedFreq = false;
end

validateattributes(nwin,{'numeric'},...
    {'scalar','integer','nonnegative','real','nonnan','nonempty',...
    'finite','>',1},'dlstft','WindowLength');

% Check that the number of samples is greater than the window length
if nwin > dataLength
    error(message('signal:stft:InvalidWindowLength',nwin));
end

% Assign other n-v pairs dependent on window
noverlap = ifEmptyUseDefault(NameValueArgs.OverlapLength,...
    floor(nwin.*defaultOverlapPercent));
if  isempty(NameValueArgs.FFTLength)
    nfft = nwin;
else
    nfft = ifEmptyUseDefault(NameValueArgs.FFTLength,nwin);
end

% OverlapLength
validateattributes(noverlap,{'numeric'},...
    {'scalar','integer','nonnegative','real','nonnan',...
    'finite','nonempty','>=',0,'<',nwin},...
    'dlstft','OverlapLength');
noverlap = double(noverlap);

% FFTLength
validateattributes(nfft,{'numeric'},...
    {'scalar','integer','nonnegative','real','nonnan',...
    'finite','nonempty','>=',nwin},'dlstft','FFTLength');
nfft = double(nfft);

% Compute useful parameters for given inputs
nfreq = fix(nfft/2) +1;
hopSize = nwin - noverlap;
ntime = fix((dataLength - noverlap) / hopSize);

% Set opts for output
opts = struct(...
    'DataFormat',NameValueArgs.DataFormat,...
    'DataLength',dataLength,...
    'TimeAxisLength',ntime,...
    'TimeMode',timeMode,...
    'TimeUnits',timeUnits,...
    'EffectiveFs',effectiveFs,...
    'IsNormalizedFreq',isNormalizedFreq,...
    'Window',win,...
    'WindowLength',nwin,...
    'OverlapLength',noverlap,...
    'HopSize',hopSize,....
    'FFTLength',nfft,...
    'NumFreqSamples',nfreq,...
    'NumChannel',numChannel,....
    'BatchSize',batchSize);
end


function output = ifEmptyUseDefault(input,default)
% Check input. If it is empty, use provided default value.
if isempty(input)
    output = default;
else
    output = input;
end
end
