function varargout = welch(x,funcName,varargin)
%WELCH Welch spectral estimation method on the GPU.
%   [Pxx,F] = WELCH(X,'pwelch',WINDOW,NOVERLAP,NFFT,Fs,SPECTRUMTYPE,TRACE,'psd')
%   [Pxx,F] = WELCH(X,'pwelch',WINDOW,NOVERLAP,NFFT,Fs,SPECTRUMTYPE,TRACE,'ms')
%   [Pxx,F] = WELCH(X,'pwelch',WINDOW,NOVERLAP,NFFT,Fs,SPECTRUMTYPE,TRACE,'power')
%   [Pxy,F] = WELCH({X,Y},'cpsd',WINDOW,NOVERLAP,NFFT,Fs,SPECTRUMTYPE)
%   [Txy,F] = WELCH({X,Y},'tfe',WINDOW,NOVERLAP,NFFT,Fs,SPECTRUMTYPE)
%   [Txy,F] = WELCH({X,Y},'tfeh2',WINDOW,NOVERLAP,NFFT,Fs,SPECTRUMTYPE)
%   [Cxy,F] = WELCH({X,Y},'mscohere',WINDOW,NOVERLAP,NFFT,Fs,SPECTRUMTYPE)
%   [...,F] = WELCH({X,Y},'mscohere',WINDOW,NOVERLAP,NFFT,Fs,...,'mimo')
%   [Pxx,F,Pxxc] = WELCH(...)
%   [Pxx,F,Pxxc] = WELCH(...,'ConfidenceLevel',P)
%
%   Inputs:
%      see "help pwelch" for complete description of all input arguments.
%      ESTTYPE - is a string specifying the type of estimate to return, the
%                choices are: psd, cpsd, tfe, tfeh2, and mscohere.
%
%   Outputs:
%      Depends on the input string ESTTYPE:
%      Pxx - Power Spectral Density (PSD) estimate, or
%      MS  - Mean-square spectrum, or
%      Pxy - Cross Power Spectral Density (CPSD) estimate, or
%      Txy - Transfer Function Estimate (TFE), or
%      Cxy - Magnitude Squared Coherence.
%      F   - frequency vector, in Hz if Fs is specified, otherwise it has
%            units of rad/sample
%
%   See also WELCH, GPUARRAY, GPUARRAY/PWELCH.

%   Copyright 2019-2023 The MathWorks, Inc.

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

if strcmp(funcName,'pwelch')
    % Look for psd, power, and ms window compensation flags for pwelch.
    [esttype, args] = signal.internal.psdesttype({'psd','power','ms'},'psd',varargin);
    if strcmpi(esttype,'ms')
        ME = MException('signal:welch:InvalidMsLegacyOption', message('signal:welch:InvalidMsLegacyOption','gpuArray'));
        throwAsCaller(ME);
    end
    % Look for legacy freqrange options.
    if any(strcmpi(args,'whole')) || any(strcmpi(args,'half'))
        ME = MException('signal:welch:InvalidFreqrangeLegacyOption', message('signal:welch:InvalidFreqrangeLegacyOption','gpuArray'));
        throwAsCaller(ME);
    end
else
    esttype = funcName;
    args = varargin;
end

% Parse input arguments.
[x,~,~,y,~,win,winName,winParam,noverlap,k1,L,options] = ...
    signal.internal.spectral.welchparse(x,esttype,args{:});
% Make sure that the default window argument is sent to the GPU
win = gpuArray(win);
% Cast to enforce precision rules
% Data type of nfft and Fs are carried over from welchparse to properly
% compute frequency vector. If nfft is interpreted as frequency vector (f
% or w), then its data type has impact on precision of computation.
noverlap = signal.internal.sigcasttofloat(noverlap,'double','WELCH',...
    'NOVERLAP','allownumeric');
k = double(k1);

% Cast to enforce precision rules.
% Both x and window are gpuArray, y can be empty (one input signal) or a
% gpuArray (two input signals). x, y, win are cast in welchparse. Data
% types of nfft and Fs are properly set in welchparse, but the impact of
% the data type of the vector nfft (i.e., frequency vector) is not
% propagated up to this point.
if isUnderlyingType(x,'single') || isUnderlyingType(options.nfft,'single')
    x = single(x);
    y = single(y);
    win = single(win);
    options.nfft = single(options.nfft);
    options.Fs = single(options.Fs);
end

% Frequency vector was specified, return and plot two-sided PSD
freqVectorSpecified = false;
nrow = 1;
if length(options.nfft) > 1
    freqVectorSpecified = true;
    [~,nrow] = size(options.nfft);
end
% Record if input x is a vector for backwards compatibility
isVectorX = isvector(x);

% Compute the periodogram power spectrum of each segment and average always
% compute the whole power spectrum, we force Fs = 1 to get a PS not a PSD

% Cast to enforce precision rules
xProto = x([]);

% First check x and y must be vectors and matrices
validateattributes(x,{'numeric'},{'2d'},'X');
nChX = size(x,2);
if ~isempty(y)
    validateattributes(y,{'numeric'},{'2d'},'Y');
    nChY = size(y,2);
end

method = "plus";

switch funcName
    case 'pwelch'
        % Place x into columns according to the window size and overlap
        x = parallel.internal.gpu.extractWindows(x,length(win),noverlap);
        % Combining method
        if options.maxhold
            method = "max";
            isAvgMethod = false;
            [Pxx,w,units] = localComputeSpectraGPU(x,[],win,...
                options,esttype,k,isAvgMethod,freqVectorSpecified);
        elseif options.minhold
            method = "min";
            isAvgMethod = false;
            [Pxx,w,units] = localComputeSpectraGPU(x,[],win,...
                options,esttype,k,isAvgMethod,freqVectorSpecified);
        else
            isAvgMethod = true;
            [Pxx,w,units] = localComputeSpectraGPU(x,[],win,...
                options,esttype,k,isAvgMethod,freqVectorSpecified);
        end
        
    case 'cpsd'
        [Pxx,w,units] = computeCrossSpectraWithSubsetsOfChannels({x,y},xProto,options.MIMO,...
                win,noverlap,options,esttype,k,freqVectorSpecified);
        
    case 'tfe'
        if options.MIMO
            % Pxx
            [Pxx,w,units] = computeCrossSpectraWithSubsetsOfChannels(x,xProto,options.MIMO,...
                win,noverlap,options,esttype,k,freqVectorSpecified);

            % Pyx
            Pyx = computeCrossSpectraWithSubsetsOfChannels({y,x},xProto,options.MIMO,...
                win,noverlap,options,esttype,k,freqVectorSpecified);
        else
            % Place x and y into columns according to the window size and
            % overlap.
            x = parallel.internal.gpu.extractWindows(x,length(win),noverlap);
            y = parallel.internal.gpu.extractWindows(y,length(win),noverlap);

            % Pxx
            [Pxx,w,units] = localComputeSpectraGPU(x,[],win,options,esttype,k,[],freqVectorSpecified);

            % Pyx
            Pyx = localComputeSpectraGPU(y,x,win,options,esttype,k,[],freqVectorSpecified);
        end
        
        if options.MIMO && nChX > 1
            Pxx = localComputeMIMO(Pxx,Pyx,[],nChX,nChY,'tfe');
        else
            Pxx = Pyx./Pxx;
        end
        
    case 'tfeh2'
        if options.MIMO
            % Pyy
            [Pyy,w,units] = computeCrossSpectraWithSubsetsOfChannels(y,xProto,options.MIMO,...
                win,noverlap,options,esttype,k,freqVectorSpecified);

            % Pxy
            Pxy = computeCrossSpectraWithSubsetsOfChannels({x,y},xProto,options.MIMO,...
                win,noverlap,options,esttype,k,freqVectorSpecified);
        else
            % Place x and y into columns according to the window size and
            % overlap.
            x = parallel.internal.gpu.extractWindows(x,length(win),noverlap);
            y = parallel.internal.gpu.extractWindows(y,length(win),noverlap);

            % Pyy
            [Pyy,w,units] = localComputeSpectraGPU(y,[],win,options,esttype,k,[],freqVectorSpecified);

            % Pxy
            Pxy = localComputeSpectraGPU(x,y,win,options,esttype,k,[],freqVectorSpecified);
        end
        
        if options.MIMO && nChX > 1
            Pxx = localComputeMIMO([],Pxy,Pyy,nChX,nChY,'tfeh2');
        else
            Pxx = Pyy./Pxy;
        end
        
    case 'mscohere'
        % Note: (Sxy1+Sxy2)/(Sxx1+Sxx2) != (Sxy1/Sxy2) + (Sxx1/Sxx2)
        % ie, we can't push the computation of Cxy into computeperiodogram

        % Pxx
        [Pxx,w,units] = computeCrossSpectraWithSubsetsOfChannels(x,xProto,options.MIMO,...
                win,noverlap,options,esttype,k,freqVectorSpecified);
        
        % Pyy - Welch Periodogram for default combination and MIMO.
        Pyy = computeCrossSpectraWithSubsetsOfChannels(y,xProto,false,...
                win,noverlap,options,esttype,k,freqVectorSpecified);
        
        % Pxy
        Pxy = computeCrossSpectraWithSubsetsOfChannels({x,y},xProto,options.MIMO,...
                win,noverlap,options,esttype,k,freqVectorSpecified);
        
        if options.MIMO && nChX > 1
            Pxx = localComputeMIMO(Pxx,Pxy,Pyy,nChX,nChY,'mscohere');
        else
            Pxx = (abs(Pxy).^2)./(Pxx.*Pyy); % Cxy
        end
        
    otherwise
        Pxx = zeros(0,'like',x);
        w = cast([],'like',real(x));
        method = "";
        
end

% Compute confidence intervals if needed
if ~isnan(options.conflevel)
    if any(strcmpi(esttype,{'ms','power','psd'})) && method ~= "plus"
        % Always compute the confidence interval around the average
        % spectrum
        isAvgMethod = true;
        [avgPxx,w] = localComputeSpectraGPU(x,[],win,...
                options,esttype,k,isAvgMethod,freqVectorSpecified);
    else
        avgPxx = Pxx;
    end
    % Cast confidence level to the correct precision 
    % This casting is done to make sure that computation of confidence
    % interval is performed in the correct precision. Since data type of Fs
    % also indicates whether computation is performed in single precision
    % or double precision, confidence level is cast to data type of Fs.
    cl = cast(options.conflevel,'like',options.Fs);
    Pxxc = signal.internal.spectral.confInterval(cl, avgPxx, isreal(x), w, options.Fs, k);
elseif nargout>2
    if any(strcmpi(esttype,{'ms','power','psd'})) && method ~= "plus"
        % Always compute the confidence interval around the average
        % spectrum
        isAvgMethod = true;
        [avgPxx,w] = localComputeSpectraGPU(x,[],win,...
                options,esttype,k,isAvgMethod,freqVectorSpecified);
    else
        avgPxx = Pxx;
    end
    cl = cast(0.95,'like',options.Fs);
    Pxxc = signal.internal.spectral.confInterval(cl, avgPxx, isreal(x), w, options.Fs, k);
else
    Pxxc = cast([],'like',real(x));
end

if nargout==0
    signal.internal.spectral.plotWelch(Pxx,w,Pxxc,esttype,noverlap,L,winName,winParam,units,options);
else
    if options.centerdc
        [Pxx, w, Pxxc] = signal.internal.spectral.psdcenterdc(Pxx, w, Pxxc, options, esttype);
    end
    
    % If the input is a vector and a row frequency vector was specified,
    % return output as a row vector for backwards compatibility.
    if nrow > 1 && isVectorX
        Pxx = Pxx.';
        w = w.';
    end
    
    % Reshape output if MIMO
    if options.MIMO
        switch lower(esttype)
            case {'tfe','tfeh2'}
                Pxx = reshape(Pxx,size(Pxx,1),nChY,nChX);
            case 'cpsd'
                Pxx = reshape(Pxx,size(Pxx,1),nChX,nChY);
        end
    end
    
    % For mscohere, ensure output is real
    if lower(esttype) == "mscohere"
        Pxx = real(Pxx);
    end
    
    if nargout < 2
        varargout = {Pxx,w}; % Pxx=PSD, MEANSQUARE, CPSD, or TFE
    else
        varargout = {Pxx,w,Pxxc};
    end
end

end

function [Pxx,w,units] = localComputeSpectraGPU(x,y,win,options,esttype,k,isAvgMethod,freqVectorSpecified)

if ~isempty(y)
    [Sxxk,w] = computeperiodogram({x,y},win,options.nfft,esttype,options.Fs);
    Sxx1 = sum(Sxxk, 2);
else
    [Sxxk,w] = computeperiodogram(x,win,options.nfft,esttype,options.Fs);
    if options.maxhold && ~isAvgMethod
        % cmethod1 = @(seg,nextseg) max(seg,real(k*nextseg));
        Sxx1 = max(real(k.*Sxxk), [], 2);
        Sxx1 = squeeze(Sxx1);
    elseif options.minhold && ~isAvgMethod
        % cmethod2 = @(seg,nextseg) min(seg,real(k*nextseg));
        Sxx1 = min(real(k.*Sxxk), [], 2);
        Sxx1 = squeeze(Sxx1);
    else
        Sxx1 = sum(Sxxk, 2);
    end
end

Sxx1 = Sxx1./k; % Average the sum of the periodograms

% Generate the freq vector directly in Hz to avoid roundoff errors due to
% conversions later.
if ~freqVectorSpecified
    w = psdfreqvec('npts',options.nfft, 'Fs',options.Fs);
end
sxx1Proto = Sxx1([]);
w = cast(w, 'like', real(sxx1Proto));

% Compute the 1-sided or 2-sided PSD [Power/freq] or mean-square [Power].
% Also, corresponding freq vector and freq units.
Sxx1 = squeeze(Sxx1);
[Pxx,w,units] = computepsd(Sxx1,w,options.range,options.nfft,options.Fs,esttype);
end

function [numChan, chindx, chindy] = localChannelIdx(x,y,MIMO)
% Find the number of cross-spectra of the columns of x and y to compute,
% numChan, and the indices into the inputs x and y to compute each
% cross-spectrum. If MIMO is specified, every combination of the columns of
% x and y is needed. Otherwise, compute cross-spectra of column-wise pairs
% of column vectors in x and y.
Nx = size(x,2);
Ny = size(y,2);
if MIMO
    numChan = Nx*Ny;
    chindx = repmat(1:Nx,1,Ny);
    chindy1 = repmat(1:Ny,Nx,1);
    chindy = chindy1(:)';
else
    % Here, the number of channels in x and y is the same, or they are
    % different with one of them being equal to one. This is, cases of
    % multichannel signals with single-channel signals.
    numChan = max(Nx,Ny);
    chindx = 1:Nx;
    chindy = 1:Ny;
    if Nx ~= Ny
        if Nx == 1
            chindx = repmat(chindx,1,Ny);
        elseif Ny == 1
            chindy = repmat(chindy,1,Nx);
        end
    end
end
end

function [Pxy,w,units] = computeCrossSpectraWithSubsetsOfChannels(x,outProto,isMIMO,...
    win,noverlap,options,esttype,k,freqVectorSpecified)
% Compute cross-spectra dividing the multichannel signal into smaller
% subsets of channels.

% Maximum number of samples in the subset of channels of each input to
% process at a time. If a single channel has more samples, it will be
% processed at once, it's not further divided into smaller segments.
maxNumSamples = 2e6;

if iscell(x)
    % Cross-spectra between two signals
    y = x{2};
    x = x{1};
    numChannelsInSubsetX = max(1, floor(maxNumSamples/size(x,1)));
    numChannelsInSubsetY = max(1, floor(maxNumSamples/size(y,1)));
    numChannelsInSubset = min(numChannelsInSubsetX, numChannelsInSubsetY);
else
    % Auto cross-spectra of a signal
    y = x;
    numChannelsInSubset = max(1, floor(maxNumSamples/size(x,1)));
end

[numChannelsXY,chindx,chindy] = localChannelIdx(x,y,isMIMO);
numSubsetsXY = ceil(numChannelsXY/numChannelsInSubset);
Pxy = zeros(0,'like',outProto);
for ii = 1:numSubsetsXY
    % Extract subset of channels to process from chindx1 and
    % chindx2.
    subsetIdx = (ii-1)*numChannelsInSubset + 1: min(ii*numChannelsInSubset, numChannelsXY);
    subsetX = x(:, chindx(subsetIdx));
    subsetY = y(:, chindy(subsetIdx));

    % Place x into columns according to the window size and
    % overlap.
    subsetX = parallel.internal.gpu.extractWindows(subsetX,length(win),noverlap);
    subsetY = parallel.internal.gpu.extractWindows(subsetY,length(win),noverlap);

    if ii ~= 1
        P = localComputeSpectraGPU(subsetX,subsetY,win,...
            options,esttype,k,[],freqVectorSpecified);
    else
        % Request w and units only once
        [P,w,units] = localComputeSpectraGPU(subsetX,subsetY,win,...
            options,esttype,k,[],freqVectorSpecified);
    end
    Pxy = [Pxy, P]; %#ok<AGROW>
end
end

function Pout = localComputeMIMO(Pxx,Pxy,Pyy,nIn,nOut,esttype)
% Compute the transfer function for a MIMO system. At each frequency, the
% transfer function H is given by H(f) = MPxy(f)/MPxx(f), where MPxy and
% MPxx are matrices reshaped from the columns of Pxy and Pxx. Note that Pxy
% is a misnomer here, since it is formed from the DFT of y times the
% complex conjugate of the DFT of x.

switch esttype
    case 'tfe'
        % Create the corresponding MIMO matrices (nIn-by-nIn for Pxx, and
        % nOut-by-nIn for Pxy) and use pagefun to compute the division that
        % defines the MIMO transfer function.
        Pxx = permute(Pxx, [3 2 1]);
        Pxx = reshape(Pxx, nIn, nIn, []);
        Pxy = permute(Pxy, [3 2 1]);
        Pxy = reshape(Pxy, nOut, nIn, []);
        Pout = pagefun(@mrdivide, Pxy, Pxx);
        Pout = reshape(Pout, nIn*nOut, []).';
    case 'tfeh2'
        % Create the corresponding nOut-by-nOut MIMO matrix from each PSD
        % matrix and use pagefun to compute the division that defines the
        % MIMO transfer function.
        Pyy = permute(Pyy, [3 2 1]);
        Pyy = reshape(Pyy, nOut, nOut, []);
        Pxy = permute(Pxy, [3 2 1]);
        Pxy = reshape(Pxy, nOut, nOut, []);
        Pout = pagefun(@mrdivide, Pyy, Pxy);
        Pout = reshape(Pout, nOut*nOut, []).';
    case 'mscohere'
        pxyProto = cast([],"like",Pxy);
        realPxyProto = real(pxyProto);
        Pout = zeros(size(Pxx,1),nOut,'like',realPxyProto); %Cxy
        Pxx = reshape(Pxx,[],nIn,nIn);
        Pxx = permute(Pxx,[2,3,1]);
        
        for iOut = 1:nOut
            Pyx0 = Pxy(:, (iOut-1)*nIn+1:(iOut-1)*nIn+nIn);
            Pxy0 = Pyx0.';
            Pyx0 = conj(Pyx0);
            Pyx0 = permute(Pyx0, [3, 2, 1]);
            Pxy0 = permute(Pxy0, [1, 3, 2]);
            Pyy0 = Pyy(:,iOut);
            Pyy0 = permute(Pyy0, [3, 2, 1]);
            
            PyxPxx = pagefun(@mrdivide, Pyx0, Pxx);
            PxyPyy = pagefun(@mrdivide, Pxy0, Pyy0);
            PoutTemp = real(pagefun(@mtimes, PyxPxx, PxyPyy));
            Pout(:,iOut) = squeeze(PoutTemp);
        end
end
end