function [YR,YI,F,T] = dlstft(x,timeValue,NameValueArgs)
%DLSTFT Deep learning short-time Fourier transform.
%   [YR,YI] = 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 numeric array, an
%   unformatted dlarray, or a dlarray in 'CBT' format. The real and
%   imaginary parts of the STFT are returned separately in YR and YI. If X
%   is a formatted dlarray, YR and YI are dlarrays with 'SCBT' format. The
%   'S' dimension corresponds to frequency in the STFT.
%   
%   DLSTFT computes the short-time Fourier transform along the 'T'
%   dimension of X. X must have at least as many samples along the 'T'
%   dimension as the length of the window. To use DLSTFT when X has fewer
%   than 128 samples than the size of the 'T' dimension, specify a smaller
%   window.
%  
%   [YR,YI] = DLSTFT(X,Fs) specifies the sample rate of X in hertz as a
%   positive scalar.
%
%   [YR,YI] = DLSTFT(X,Ts) specifies the sample time, Ts, as a positive 
%   scalar duration. The sample rate in hertz is computed as the inverse
%   of Ts with Ts converted to seconds.
%   
%   [YR,YI] = DLSTFT(...,'DataFormat',FMT) specifies the data format, FMT,
%   if X is an unformatted dlarray or numeric array. FMT is any permutation
%   of 'CBT'. 'DataFormat' must be specified if X is an unformatted dlarray
%   or a numeric array. Do not specify 'DataFormat' if X is a formatted
%   dlarray. If you specify the data format, YR and YI are unformatted
%   dlarrays.
%   
%   [YR,YI] = DLSTFT(...,'Window',WINDOW) specifies the window used to
%   compute the STFT as a vector. DLSTFT divides each signal in X 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. The function then windows each segment 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.
% 
%   [YR,YI] = 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.
% 
%   [YR,YI] = 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 YR and YI corresponds to the DFT bins in the STFT.
% 
%   [YR,YI,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.
% 
%   [YR,YI,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:
%         % Create a 'CTB' formatted dlarray with 3 channels representing a
%         % batch of 5.
%         X = dlarray(rand([3 1000 5]),'CTB');
% 
%         % Compute the deep learning short-time Fourier transform
%         [YRDL,YIDL] = 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'.
%         [YRDL,YIDL,FDL,TDL] = dlstft(X,fs,'Window',win,...
%            'OverlapLength',noverlap,'FFTLength',nfft,'DataFormat','TCB');

% Copyright 2020-2021 The MathWorks, Inc.

arguments
    x dlarray {mustBeFloat,mustBeNonempty}
    timeValue = []
    NameValueArgs.Window {mustBeFloat,mustBeNonempty,mustBeFinite,...
        mustBeReal,mustBeVector}
    NameValueArgs.OverlapLength = []
    NameValueArgs.FFTLength = []
    NameValueArgs.DataFormat {mustBeText} = ''
    NameValueArgs.Weights {mustBeA(NameValueArgs.Weights,'cell'),mustBeNonempty}     
end

% Check outputs
nargoutchk(0,4);

%---------------------------------
% Validate input and further parse inputs
[x,opts,isInitialized] = dlstftParser(x,timeValue,NameValueArgs);

%---------------------------------
% Compute real and imaginary parts of STFT separately for dlarray. 
[YR,YI] = computeDLSTFT(x,opts,isInitialized);


% Reformat YR and YI as needed
[YR,YI] = reformatOutput(YR,YI,opts);

% Frequency Vector
if nargout > 2
    
    if opts.IsNormalizedFreq
        f = psdfreqvec('npts',opts.FFTLength,'Range','half');
    else
        f = psdfreqvec('npts',opts.FFTLength,'Range','half',...
            'Fs',opts.EffectiveFs);
    end
    
    % cast frequency vector to match YR 
    F = cast(f(1:opts.NumFreqSamples),'like',YR); 
    
   
end

% Time Vector
if nargout > 3
    coloffsets = (0:(opts.TimeAxisLength - 1)) * opts.HopSize;
    t = (coloffsets + (opts.WindowLength / 2)') / opts.EffectiveFs;
    
    % Scale time vectors in the case of normalized frequency
    if opts.IsNormalizedFreq
        t = t .* opts.EffectiveFs; % samples
    end

    T = t(1:opts.TimeAxisLength);
    if isnumeric(T) && ~isempty(opts.TimeUnits)
        T = duration(0,0,T,'Format',opts.TimeUnits);
    else
        T = cast(T,'like',YR);
    end
    
    % Ensure time output is a column vector
    T = T(:);
end

end

function [wR,wI] = initializeDFTBasisCore(nfft, win)
% Initialize windowed DFT basis 
nwin = length(win);
if nwin < nfft
    win(end+nfft-nwin,1) = 0;
end

nfreq = fix(nfft/2+1);
N = length(win);
n = 0:N-1;
idx = 1:nfreq;

wR = win.*cos(2*pi*(idx-1).*n'/N);
wI = win.*-sin(2*pi*(idx-1).*n'/N);

end

function [wR,wI] = shapeBasis(wR,wI,shape,replication)
% Reshape and zero-pad basis with given parameters
% Real part
wR = reshape(wR,shape);
wR = repmat(wR,replication);

% Imaginary part
wI = reshape(wI,shape);
wI = repmat(wI,replication);
end

function [wR,wI] = initializeDFTBasisCBT(nfft, win, nchannel)
% Initialize DFT basis for CBT format X. Convolution is done along T.
[wR,wI] = initializeDFTBasisCore(nfft, win);
[wR,wI] = shapeBasis(wR,wI,[size(wR,1),1,size(wR,2)],[1,1,1,nchannel]);
end

function [YR,YI] = computeDLSTFT(x,opts,isInitialized)
% Computes the short-time Fourier transform for 'CBT' or 'SSCB' format
% dlarray. YR and YI are unformatted dlarray.

if ~isInitialized
% Initialize DFT Basis 
    [wReal,wImag] = initializeDFTBasisCBT(opts.FFTLength, opts.Window, ...
        opts.NumChannel);
    
else
    wR = opts.Weights(:,:,1);
    wI = opts.Weights(:,:,2);
    [wReal,wImag] = shapeBasis(wR,wI,[size(wR,1),1,size(wR,2)],...
        [1,1,1,opts.NumChannel]);    
    
end

% Padding data
xpad = zeros([opts.NumChannel,opts.BatchSize,opts.PaddingSize],'like',x);
x = cat(3,x,xpad);

% Compute STFT for first channel
YR = dlconv(x,wReal,0,'Stride',opts.HopSize,'DataFormat','CBT',...
    'WeightsFormat','TCUU');
YI = dlconv(x,wImag,0,'Stride',opts.HopSize,'DataFormat','CBT',...
    'WeightsFormat','TCUU');

end

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

isInitialized = isfield(NameValueArgs,'Weights');

if isInitialized && isfield(NameValueArgs,'Window')
    error(message('signal_deep:dlstft:InvalidWeightWindow'));    
end
if isInitialized
    % Window size is size of weight matrix along the row dimension
    nwin = NameValueArgs.Weights{2};
    win = [];
elseif ~isInitialized && ~isfield(NameValueArgs,'Window')
    win = hann(128,'periodic');
    nwin = 128;
else
    win = 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
if ~isInitialized
    % 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

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 isInitialized && ~isempty(NameValueArgs.FFTLength)
    error(message('signal_deep:dlstft:InvalidFFTWeight'));
elseif isInitialized && 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);
if ~isInitialized
    npadding = nfft - nwin;
else
    npadding = 0;
end

% Set opts for output
opts = struct(...
    'DataFormat',NameValueArgs.DataFormat,...
    'DataLength',dataLength,...
    'PaddingSize',npadding,...
    '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);

if isInitialized 
    opts.Weights = NameValueArgs.Weights{1}; 
end

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


function [YR,YI] = reformatOutput(YR,YI,opts)

% Output is SCBT
YR = reshape(YR,[size(YR,1)/opts.NumChannel,opts.NumChannel,...
        opts.BatchSize,size(YR,3)]);
YI = reshape(YI,[size(YI,1)/opts.NumChannel,opts.NumChannel,...
        opts.BatchSize,size(YI,3)]);

if isempty(opts.DataFormat)
    YR = dlarray(YR,'SCBT');
    YI = dlarray(YI,'SCBT');
end



end



