function g = czt(x, varargin)
%CZT  Chirp z-transform on the GPU
%   G = CZT(X, K, W, A)

%   Copyright 2019-2023 The MathWorks, Inc.

narginchk(1,4);
if nargin > 1
    [varargin{:}] = gather(varargin{:});
end

if ~isgpuarray(x)
    % Dispatch to CPU if X not gpuArray
    g = czt(x, varargin{:});
    return;
end

% Parsing Inputs
[x,k,w,a,sz,transpose,thinArray] = inputParser(x,varargin{:});
m = sz(1);

%------- Length for power-of-two fft.
nfft = 2^nextpow2(m+k-1);

%------- Premultiply data.
kk = gpuArray.colon((-m+1),max(k-1,m-1)).';
nn = gpuArray.colon(0,(m-1)).';

ww = arrayfun(@(w,kk) w^((kk^2)/2),w,kk);
aa = arrayfun(@(ww,a,nn) ww*(a^-nn),ww(m:2*m-1),a,nn);

y = x .* aa(:,1);

%------- Fast convolution via FFT.

fy = fft(y,nfft);

fv = fft(1./ww(1:k-1+m),nfft);   % <----- Chirp filter.

fy = fy .* fv(:,1);
gg  = ifft(fy);

%------- Final multiply.
g = gg(m:m+k-1,:,:) .* ww(m:m+k-1,1);

% Transpose if input signal (X) on each page wasn't a column vector
if transpose == true
    g = permute(g,[2,1,3:length(sz)]);
end

if thinArray
    g = reshape(g,1,1,[]);
end

end

function [x,k,w,a,sz,transpose,thinArray] = inputParser(x,varargin)
% Input can be a vector, matrix or a 3-D array
validateattributes(x,{'double','single'},{'3d'},'czt','X',1)
sz = size(x);

if length(sz) == 3 && sz(1) == 1 && sz(2) == 1
    % Input is 1x1xK - to be treated as a vector of length K
    x = reshape(x,[],1);
    sz = size(x);
    thinArray = true;
else
    thinArray = false;
end

% Signal (X) is converted to a column vector on each page in InputParser so a transpose
% back would be required if it is a row vector
if sz(1) == 1
    transpose = true;
else
    transpose = false;
end

if sz(1) == 1
    x = permute(x,[2,1,3:length(sz)]);
    sz = size(x);
end

if nargin < 2
    k = max([sz(1),sz(2)]);
else
    validateattributes(varargin{1},{'numeric'},{'scalar','finite'},'czt','M',2)
    k = varargin{1};
end
% Cast transform length so default complex ratio W is computed in desired
% precision. Input signal X is of type floating point, single or double,
% which is in agreement with data type of W.
k = cast(k,underlyingType(x));

if nargin < 3
    w = exp(-1i.* 2.*pi./k);
else
    validateattributes(varargin{2},{'double','single'},{'scalar'},'czt','W',3)
    w = varargin{2};
end

if nargin < 4
    a = cast(1,'like',k+w);
else
    validateattributes(varargin{3},{'double','single'},{'scalar'},'czt','A',4)
    a = varargin{3};
end

% Signal (X), W and A determines type of output (G)
isSingle = isaUnderlying(x,'single') || isa(w,'single') || isa(a,'single');

% Cast all inputs to correct type
if isSingle
    x = single(x);
    k = single(k);
    w = single(w);
    a = single(a);
end

end
