function varargout = pspectrum(x,varargin)
%PSPECTRUM Analyze signals in the frequency and time-frequency domains on the GPU
%   P = PSPECTRUM(X)
%   P = PSPECTRUM(X,Fs)
%   P = PSPECTRUM(X,Ts)
%   P = PSPECTRUM(X,Tv)
%   P = PSPECTRUM(...,TYPE)
%   P = PSPECTRUM(...,'TwoSided',VAL)
%   P = PSPECTRUM(...,'FrequencyLimits',FLIMS)
%   P = PSPECTRUM(...,'Leakage',LKG)
%   P = PSPECTRUM(...,'MinThreshold',THRESH)
%   P = PSPECTRUM(...,'FrequencyResolution',FRES)
%   P = PSPECTRUM(...,'TimeResolution',TRES)
%   P = PSPECTRUM(...,'OverlapPercent',OP)
%   P = PSPECTRUM(...,'NumPowerBins',NPB)
%   [P,F] = PSPECTRUM(...)
%   [P,F,T] = PSPECTRUM(...,'spectrogram')
%   PSPECTRUM(...)
%
%   Limitations:
%   1. Persistence spectrum is not supported.
%   2. Reassigned spectrum or spectrogram is not supported.
%
%   See also PSPECTRUM, GPUARRRAY, PERIODOGRAM, PWELCH, SPECTROGRAM

%   Copyright 2021-2023 The MathWorks, Inc.

narginchk(1, 18);
nargoutchk(0, 3);

if nargin > 1
    if isgpuarray(varargin{1}) && ~isscalar(varargin{1})
        % A time vector has been supplied as a gpuArray, keep it on the
        % GPU.
        if nargin > 2
            [varargin{2:end}] = gather(varargin{2:end});
        end
        % Ensure the signal is also on the GPU.
        x = gpuArray(x);
    else
        [varargin{:}] = gather(varargin{:});
    end
end

% Dispatch in-memory pspectrum if input signal is not on the GPU.
if ~isgpuarray(x)
    varargout = cell(1, nargout);
    [varargout{:}] = pspectrum(x, varargin{:});
    return;
end

opts = signal.internal.pspectrum.pspectrumParser(x, varargin{:});
if opts.Reassign
    % gpuArray limitation: reassign spectrum/spectrogram not supported.
    error(message("signal:pspectrum:ReassignedNotSupportedGPU"));
end

generatePlots = nargout==0;

if opts.Type == "power"
    TOrPwr = gpuArray([]);
    [P, F] = computeSpectrum(opts, generatePlots);
elseif opts.Type == "spectrogram"
    [P, F, TOrPwr] = computeSpectrogram(opts, generatePlots);
else
    % gpuArray limitation: persistence spectrum not supported.
    error(message("signal:pspectrum:PersistenceNotSupportedGPU"));
end

if opts.IsNormalizedFreq
    F = F*pi;
end

if opts.IsSingle
    P = single(P);
    F = single(F);
    if isnumeric(TOrPwr)
        TOrPwr = single(TOrPwr);
    end
end

outputs = {P, F, TOrPwr};
varargout = outputs(1:nargout);
end

%--------------------------------------------------------------------------
%--------------------------------------------------------------------------
function [P, F] = computeSpectrum(opts, plotFlag)
% Compute power spectrum

if opts.IsNormalizedFreq
    timeFactor = 0.5;
else
    timeFactor = 1;
end
Fs = opts.EffectiveFs;
f1 = opts.FrequencyLimits(1);
f2 = opts.FrequencyLimits(2);
% Correct t1 and t2 by half a sample so that (t2-t1)*Fs equals the exact
% signal length.
t1 = opts.t1*timeFactor - 0.5/Fs;
t2 = opts.t2*timeFactor + 0.5/Fs;
kbeta = signal.internal.pspectrum.convertLeakageToBeta(opts.Leakage);

if isempty(opts.FrequencyResolution)
    % Auto mode
    Npoints = 4096;
else
    % Set Npoints to achieve desired resolution
    fspan = 4*Fs;
    Npoints = round((1 + fspan/opts.FrequencyResolution(1)));
end

zoomSpectrumObj = signal.internal.gpu.pspectrum.ZoomSpectrum(...
        t1, t2,...       % init and end times of input signal
        f1, f2,...       % frequency band
        kbeta,...        % Kaiser beta
        Npoints,...      % Number of frequency points
        4*Fs); ...       % Maximum sample rate

zoomSpectrumObj.specifyTimeVector(Fs, t1, t2);

convertToDB = false;
if opts.TwoSided
    P = zoomSpectrumObj.fetchTwoSidedSpectrum(opts.Data, convertToDB);
else
    P = zoomSpectrumObj.fetchOneSidedSpectrum(opts.Data, convertToDB);
end
F = zoomSpectrumObj.fetchFrequencyVector();


if opts.MinThreshold > 0
    P(P<opts.MinThreshold) = 0;
end

if plotFlag
    if opts.IsSingle
        P = single(P);
        F = single(F);
    end
    FRES = zoomSpectrumObj.getActualResolutionBandwidth();
    signal.internal.pspectrum.displaySpectrum(F, P, opts.IsNormalizedFreq, FRES);
end
end

%--------------------------------------------------------------------------
%--------------------------------------------------------------------------
function [P, F, T] = computeSpectrogram(opts, plotFlag)
% Compute spectrogram

Fs = opts.EffectiveFs;
f1 = opts.FrequencyLimits(1);
f2 = opts.FrequencyLimits(2);
if opts.IsNormalizedFreq
    timeFactor = 0.5;
else
    timeFactor = 1;
end

kbeta = signal.internal.pspectrum.convertLeakageToBeta(opts.Leakage);
ENBW = signal.internal.pspectrum.getENBWEstimate(kbeta);
Npoints = 1024;

timeResolutionAuto = false;
timeResolution = -1;

if isempty(opts.FrequencyResolution) && isempty(opts.TimeResolution)
    timeResolutionAuto = true;
    timeResolution = -1; % don't-care value
elseif ~isempty(opts.FrequencyResolution)
    timeResolutionAuto = false;
    % Need to divide by time factor because when working in normalized
    % frequency we want timeResolution in samples. Dividing by time factor
    % is the same as multiplying by Fs=2. In the non-normalized case
    % timeResolution must be in seconds so in that case we do not need to
    % multiply by Fs.
    timeResolution = (ENBW/opts.FrequencyResolution)/timeFactor;
    fspan = 4*opts.EffectiveFs;
    Npoints = round((1 + fspan/opts.FrequencyResolution));
elseif ~isempty(opts.TimeResolution)
    timeResolutionAuto = false;
    timeResolution = opts.TimeResolution;
end

% Do not convert to dBs for spectrogram.
convertToDB = false;

% t1 and t2 do not require correction by half a sample. Internal
% spectrogram code re-centers times when calling the spectrum estimator. 
spectrogramObj = signal.internal.gpu.pspectrum.Spectrogram(...
    timeResolutionAuto,...  % auto resolution if true
    timeResolution,...      % time res value - irrelevant if auto
    opts.OverlapPercent,... % overlap percentage
    opts.DataLength,...     % entire signal length
    timeFactor,...          % 0.5 if normalized freq, 1 otherwise
    Fs,...                  % sample rate
    opts.t1, opts.t2,...    % init and end times of input signal
    f1, f2,...              % frequency band
    kbeta,...               % Kaiser beta
    Npoints,...             % Number of frequency points
    4*Fs,...                % Maximum sample rate
    convertToDB,...         % Convert powers to dB if true
    opts.TwoSided);         % Two sided spectrum if true

% Call getSpectrogramMatrix directly and only with input data, here we
% always have an entire sequence to process. No need to include the time
% vector, opts.t1 and opts.t2 will always match with opts.TimeVector(1) and
% opts.TimeVector(end).
P = spectrogramObj.getSpectrogramMatrix(opts.Data);
F = spectrogramObj.fetchFrequencyVector();
T = spectrogramObj.fetchTimeVector()/timeFactor;
FRES = spectrogramObj.getTargetResolutionBandwidth();
TRES = spectrogramObj.getTimeResolution();

if opts.MinThreshold > 0
    P(P<opts.MinThreshold) = 0;
end

if ~isempty(opts.InitialDate)
    % Set times to datetime format if time information is in datetimes
    T = seconds(gather(T)) + opts.InitialDate;
end

if plotFlag
    if opts.IsSingle
        P = single(P);
        F = single(F);
        if isnumeric(T)
            T = single(T);
        end
    end
    signal.internal.pspectrum.displaySpectrogram(T, F, P, ...
        opts.IsNormalizedFreq, opts.MinThreshold, FRES, TRES);
else
    if isnumeric(T) && ~isempty(opts.TimeUnits)
        % Set times to duration format if time information is in durations
        T = duration(0, 0, gather(T), "Format", opts.TimeUnits);
    end
end
end
