function w=fresolve(x,K,w_GT)
% USAGE    : w=fresolve(x,K,[w_GT]);
% FUNCTION : Super-resolves the largest K peaks of the DFT of the signal x
% using the formula published in a "Tips & Tricks" column of the IEEE
% Signal Processing Magazine (2023). If the ground-truth frequencies, w_GT,
% are given as input, comparisons with the super-resolved frequencies are  
% shown as well, together with some visualization.
% If the DFT has less than K peaks, a warning occurs.
%
% INPUT    : * x, the samples of the signal
%            * K, the number of peaks to resolve
%            * w_GT, the ground-truth frequencies (optional), only needed
%              for comparison with the super-resolved frequencies and plots
% OUPUT    : * w, the (cyclic) frequency after super-resolution (in ]-pi,pi])
%
% DATE     : October 2023
% AUTHOR   : Thierry Blu, mailto:thierry.blu@m4x.org 
% REFERENCE: Guo, R. & Blu, T.,"Super-Resolving a Frequency Band", 
%            IEEE Signal Processing Magazine, 2023. To appear. 
%__________________________________________________________________________
% EXAMPLE OF USE
% % Generation of N samples of a signal made of K frequencies
% K=5;N=30;
% % First, choosing frequencies (reasonably separated from each other)
% sep=5;
% w_GT=sort(2*pi*rand(sep*K,1));
% w_GT=pi-w_GT(1:sep:end);
% % Second, choosing amplitudes (with reasonably close absolute values)
% a_GT=randn(K,2)*[1;1i];
% a_GT=a_GT./abs(a_GT).*(1+0.1*rand(K,1));
% % Third, building the samples of the signal
% x=a_GT.'*exp(1i*w_GT*(0:(N-1)));
%
% % Finally, super-resolving the K DFT peaks
% w=fresolve(x,K,w_GT);

x=x(:);
N=length(x);
n=0:(N-1);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%% Main calculation %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%__________________________________________________________________________
% Functions
arg=@(theta)pi-mod(pi-theta,2*pi);                  % Wrap frequencies in ]-pi,pi]
Fourier=@(w)exp(-i*w*n)*x;                          % DTFT coefficient at w
trick=@(w)w+2*atan(tan(pi/(2*N))*...                % Super-resolution formula
    (abs(Fourier(w+pi/N))-abs(Fourier(w-pi/N)))./...
    (abs(Fourier(w+pi/N))+abs(Fourier(w-pi/N))));   

%__________________________________________________________________________
% Peaks of the DFT of the signal x
y=abs(fft(x));
y=[y(N);y;y(1)];                                    % Periodic extension of the samples 
[ym,j]=findpeaks(y,'SortStr','descend','NPeaks',K); % Peaks found exclude first and last sample
w_DFT=arg(2*pi*mod(j-2,N)/N);                       % Frequencies of the DFT peaks

% Has findpeaks found the requested K peaks, or less?
if length(j)<K
    K_SR=length(j);
    peaks=' peaks';if K_SR==1,peaks=' peak';end
    warning(['Only ' num2str(K_SR) peaks ' found!'])
end

%__________________________________________________________________________
% Super-resolution by applying the trick (main idea 
% of the Signal Processing Magazine paper in reference)
w=trick(w_DFT);                                     % First application of the trick
w=trick(w);                                         % Refinement
w=arg(w);                                           % frequencies are in ]-pi,pi]

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%% Plots and comparisons %%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
if exist('w_GT')
%__________________________________________________________________________
% Plot of the DTFT, the DFT, and the super-resolved/ground-truth frequencies
    y=y(2:end-1);
    y_DTFT=abs(fft(x,10*N));
    [w_DTFT,k_DTFT]=sort(arg(2*pi*(0:10*N-1)/(10*N)));
    [wk,k]=sort(arg(2*pi*n/N));
    plot(w_DTFT,y_DTFT(k_DTFT),':k',wk,y(k),'b','linewidth',1),hold on
    plot(w,abs(Fourier(w)),'or',w_GT,abs(Fourier(w_GT)),'*k','markersize',10),hold off
    axis([-pi,pi,0,max(y_DTFT)*(1+1/3)])
    xlabel('Cyclic frequency \omega')
    ylabel('Fourier amplitude')
    title({'Discrete-Time Fourier Transform (DTFT)' 'and super-resolved DFT frequency peaks'})
    legend('DTFT','DFT','Super-resolved frequencies','Ground-truth frequencies')
    set(gca,'fontsize',14)
%__________________________________________________________________________
% Comparison with the ground-truth
    % Pairing the resolved frequencies with the GT by
    %   1. sorting the super-resolved and the ground-truth frequencies
    %   together (w_all)
    %   2. identifying iteratively the closest SR-GT frequencies and removing
    %   them from the list of frequencies to pair

    K_SR=length(w);K_GT=length(w_GT);
    [w_all,j]=sort([w;w_GT]);    
    j_GT=[];j_SR=[];stop=0;
    while ~stop
        idx_trans=find(...                          % Indices for which a frequency
            diff(double([j(end);j]>K_SR))~=0);      % from w and from w_GT are consecutive
        err=abs(diff([w_all(end)-2*pi;w_all]));     % Error between consecutive frequencies
        [~,k]=min(err(idx_trans));                  % Index of the minimal error when the consecutive
        k=idx_trans(k);                             % frequencies are from w and w_GT
        if j(k)>K_SR
            j_GT=[j_GT;j(k)-K_SR];
            j_SR=[j_SR;j(mod(k-2,length(j))+1)];
        else
            j_GT=[j_GT;j(mod(k-2,length(j))+1)-K_SR];
            j_SR=[j_SR;j(k)];
        end
        w_all=w_all(setdiff(1:end,...               %
            [mod(k-2,length(j))+1,k]));             % Remove paired frequencies from w_all
        j=j(setdiff(1:end,...                       % 
            [mod(k-2,length(j))+1,k]));             % and their indices in w and w_GT
        stop=(max(j)<=K_SR|min(j)>K_SR);            % Only remaining frequencies are either SR, or GT 
    end
    if max(j)>K_SR
        w_SR=[w(j_SR);NaN(K_GT-length(j_SR),1)];
        w_GT=[w_GT(j_GT);w_GT(j-K_SR)];
    else
        w_SR=[w(j_SR);w(j)];
        w_GT=[w_GT(j_GT);NaN(K_SR-length(j_GT),1)];
    end
    err=arg(w_SR-w_GT)/(2*pi/N);                    % 2*pi/N normalized error between the SR and the GT frequencies
    
    disp(' ')
    matlab_version=version('-release');
    if str2num(matlab_version(1:4))<=2019           % variable names are more restricted in older Matlab versions
        disp(table(w_SR,w_GT,err,...
            'variablenames',{'w_SR' 'w_GT' 'error_x_N/2\pi'}))
    else
        disp(table(w_SR,w_GT,err,...
            'variablenames',{'w_SR' 'w_GT' 'error x N/2\pi'}))
    end
    disp(' ')
    disp('w_SR = super-resolved frequency')
    disp('w_GT = ground-truth frequency')
    disp(['DFT bin size = ' num2str(2*pi/N) '  (2\pi/N)'])
end
