function [X] = dlistft(s, NameValueArgs)
%DLISTFT Deep learning inverse short-time Fourier transform.
%   X = DLISTFT(Y) returns the inverse short-time Fourier transform (ISTFT)
%   of the one-sided short-time Fourier transform, Y. The ISTFT is computed
%   using the default 128-sample Hann window and an equal number of
%   discrete Fourier transform (DFT) bins. The number of overlapped samples
%   between windows is 75% of the window length. Y can be a formatted
%   dlarray or an unformatted dlarray. If Y is a formatted dlarray, it must
%   have 'SCBT' or 'CBT' format. In both cases, the first dimension of Y
%   must be divisible by floor(NFFT/2)+1 where NFFT is the FFTLength. If Y
%   has 'CBT' format, DLISTFT unflattens Y to have shape 'SCBT'. If Y is
%   real-valued, DLISTFT assumes that Y is formed by concatenating the real
%   and imaginary parts of the short-time Fourier transform along the
%   channel dimension. If Y is unformatted, it must be compatible with the
%   'SCBT' or 'CBT' formats and you must specify the DataFormat name-value
%   argument. See the documentation for detailed information on the size
%   and format requirements of DLISTFT.
%
%   DLISTFT computes the inverse short-time Fourier transform along the 'S'
%   dimension of Y. For perfect reconstruction for non-modified spectra,
%   the window must be COLA compliant.
%
%   X = DLISTFT(Y,DataFormat=FMT) specifies the input data format, FMT,
%   when Y is an unformatted dlarray. You must specify DataFormat if Y is
%   an unformatted dlarray. FMT is any permutation of 'SCBT' or 'CBT'. Y
%   must be compatible with the specified FMT.  Do not specify DataFormat
%   if Y is a formatted dlarray. If you specify FMT, X is an unformatted
%   dlarray compatible with the 'CBT' format.
%
%   X = DLISTFT(...,Window=WINDOW) specifies the window used to compute the
%   ISTFT. For perfect time-domain reconstruction, the window for the ISTFT
%   computation must match the window for the STFT computation. Use the
%   function ISCOLA to check a window/overlap combination for constant
%   overlap-add (COLA) compliance. If the window is a dlarray, extract data
%   from the dlarray before using the ISCOLA function. COLA compliance is a
%   requirement for perfect reconstruction for non-modified spectra. This
%   argument defaults to a 128-sample periodic Hann window.
%
%   X = DLISTFT(...,OverlapLength=NOVERLAP) specifies an integer number of
%   samples of overlap between adjoining segments. The specified value must
%   be smaller than the length of the window. This argument defaults to the
%   largest integer less than or equal to 75% of the window length.
%
%   X = DLISTFT(...,FFTLength=NFFT) specifies the integer number of
%   frequency bins used to calculate the DFT. The specified value must be
%   greater than or equal to the window length. This argument defaults to
%   the window length. For perfect time-domain reconstruction, specify an
%   equal number of bins for the STFT and the ISTFT.
%
%   X = DLISTFT(...,Method=METHOD) specifies the overlap-add method to use
%   for reconstruction. Valid options are: 'ola' (Overlap-add) and 'wola'
%   (Weighted overlap-add). This argument defaults to 'wola'.
%
%   X = DLISTFT(...,ExpectedOutputSize=OUTPUTSIZE) specifies the expected
%   number of channels and samples output by DLISTFT as a two-element
%   vector. The first element of OUTPUTSIZE is the expected number of
%   channels and the second element is the expected number of time samples.
%   DLISTFT errors if the ISTFT does not match OUTPUTSIZE in the number of
%   channels and samples. If you do not specify ExpectedOutputSize, DLISTFT
%   does not check the output size of the ISTFT. See the documentation for
%   details on how the size of the ISTFT is determined.
%
%   % EXAMPLE 1:
%   %   Compute the deep learning short-time Fourier transform of a 'CTB'
%   %   formatted random dlarray with three channels and 1024 samples representing
%   %   a batch of 5, using the default window, FFT length and overlap
%   %   length values for both DLSTFT and DLISTFT.
%   x = dlarray(rand([3 1024 5]),'CTB');
%   ydl = dlstft(x);
%   %   Compute the deep learning inverse short-time Fourier transform of a
%   %   'SCBT' formatted dlarray.
%   X = dlistft(ydl);
%
%   % EXAMPLE 2:
%   %   Create an unformatted dlarray containing a batch of two
%   %   single-channel signals. Each signal in the batch contains a
%   %   quadratic chirp.
%   fs = 8192;
%   t = 0:1/fs:2-1/fs;
%   x = dlarray(zeros([16384 1 2]));
%   x(:,:,1) = chirp(t,250,1,500,'q');
%   x(:,:,2) = chirp(t,500,1,1000,'q');
%   nfft = 600;
%   win = hann(512);
%   noverlap = 384;
%   outsize = [1 16384];
%   %   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');
%   %   Compute the deep learning inverse short-time Fourier transform.
%   %   Specify the data format as 'SCBT' because ydl is an unformatted
%   %   dlarray. To obtain perfect reconstruction, use the same name-value
%   %   arguments for the ISTFT computation as for the STFT computation.
%   X = dlistft(ydl, FFTLength=nfft, DataFormat='SCBT',...
%       OverlapLength=noverlap, Window=win, ExpectedOutputSize=outsize);
%
%   See also DLSTFT, STFTLAYER, STFT, ISTFT

%   Copyright 2023 The MathWorks, Inc.

arguments
    s dlarray {mustBeFloat,mustBeNonempty}
    NameValueArgs.Window dlarray {mustBeFloat,mustBeNonempty,mustBeFinite,...
        mustBeReal,mustBeVector}
    NameValueArgs.OverlapLength = []
    NameValueArgs.FFTLength = []
    NameValueArgs.DataFormat {mustBeTextScalar} = ''
    NameValueArgs.Method {mustBeMember(NameValueArgs.Method,{'wola','ola'})} = 'wola'
    NameValueArgs.ExpectedOutputSize {mustBeFinite,mustBeReal,mustBeVector}
end

% Check outputs (1 outputs -> original signal)
nargoutchk(0,1);

% Validate input and parse inputs
[s,opts] = inputParser(s, NameValueArgs);

win = opts.Window;
nfft = opts.FFTLength;
hopsize = opts.HopSize;
method = opts.Method;

if strcmpi(dims(s),'CBT')
    X = signal.deep.internal.istftCBTInput(s,win,nfft,hopsize,method);
else
    X = signal.deep.internal.istftSCBTInput(s,win,nfft,hopsize,method);
end

% X must be unformatted if the input was unformatted.
if ~isempty(NameValueArgs.DataFormat)
    X= stripdims(X);
end
end

function [s,opts] = inputParser(s,NameValueArgs)
%inputParser is a function for parsing value-only inputs, flags, and
%name-value pairs for the DLISTFT functions.

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

nfft = ifEmptyUseDefault(NameValueArgs.FFTLength,nwin);

% Set defaults
defaultOverlapPercent = 0.75;

% If specified DataFormat must be 'CBT' or 'SCBT'
% Validate input and compute parameters
s = deep.internal.dlarray.validateDataFormatArg(s, NameValueArgs.DataFormat);
% dims(s) must be 'CBT' or 'SCBT'
if ~ismember(dims(s),{'CBT','SCBT'})
    error(message('signal_deep:dlstft:InvalidISTFTDataFormat',dims(s)));
end

% validate window length
if nwin < 2
    error(message('signal_deep:dlstft:InvalidWindow'));
end

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

% Method
method = NameValueArgs.Method;

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

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

% Compute useful parameters for given inputs
nfreq = fix(nfft/2) +1;
hopSize = nwin - noverlap;
% Validate the size of the STFT is compatible with dlistft
if startsWith(dims(s),'S')
    if size(s,1) ~= nfreq
        error((message('signal_deep:dlstft:InvalidFirstDimSCBT', ...
            num2str(size(s,1)), num2str(nfreq))));
    end
    %Only option left is CBT -- must be divisible by nfreq
elseif mod(size(s,1),nfreq)
    error((message('signal_deep:dlstft:InvalidFirstDimCBT', ...
        num2str(size(s,1)), num2str(nfreq))));
end
if isreal(s) && startsWith(dims(s),'S')
    if signalwavelet.internal.isodd(size(s,2))
        error((message('signal_deep:dlstft:InvalidChannelSize', ...
            num2str(size(s,2)))));
    end
elseif isreal(s) && startsWith(dims(s),'C')
    if mod(size(s,1)/nfreq/2,1)
        error((message('signal_deep:dlstft:InvalidRealCBTFirstDim', ...
            num2str(size(s,1)), num2str(nfreq))));
    end
end

% Calculate the validate the output size of the DLISTFT with the expected
% output size.
if isfield(NameValueArgs,'ExpectedOutputSize')
    validateattributes(NameValueArgs.ExpectedOutputSize,{'numeric'},...
        {'numel',2},'dlistft','ExpectedOutputSize');
    expectedChannels = NameValueArgs.ExpectedOutputSize(1);
    expectedTime = NameValueArgs.ExpectedOutputSize(2);
    [match,outChannels,siglen] = signal.deep.internal.checkSTFTSize(s,nfreq,...
        nwin,hopSize,expectedChannels,expectedTime);
    if ~match
        error(message('signal_deep:dlstft:InvalidInputSize', ...
            num2str(expectedChannels), num2str(expectedTime), num2str(outChannels),...
            num2str(siglen)));
    end
end

% Set opts for output
opts = struct(...
    'DataFormat',NameValueArgs.DataFormat,...
    'Window',win,...
    'WindowLength',nwin,...
    'OverlapLength',noverlap,...
    'HopSize',hopSize,....
    'FFTLength',nfft,...
    'Method', method);
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 mustBeMember(A,B)
validatestring(A,B,'dlistft');
end

