%% Test OMP and BLOOMP on spectral estimation
%% y_k = \sum_{j=1}^s x_j*e^{-2*\pi*i*k*supp(j)} + z_k    
%% supp(j) \in [0,1)
%% k = -fc : fc
%% Rayleigh length RL = 1/(2*fc)
clear all; 
close all; 
display('===========================================')

noise    = 5;   % noise level ||z||_2/||\sum_{j=1}^s x_j*e^{-2*\pi*i*k*supp(j)}||_2 = noise/100
range    = 10;  %dyanmic range = xmax/xmin
fprintf('noise = %6.0f  dynamic range = %6.0f\n',noise,range)
        
%% sample range
fc   = 50;
RL   = 1/(2*fc);   %Rayleigh length

%% exact frequency 
%%%%%%%%%%%% 15 equilly spaced frequencies
%sep  = sepn * RL;       %support separation
%sparsity  = 15;
%supp      = (0.1 + 0.1*rand + (1:sparsity)*sep).';
%%%%%%%%%%%% frequencies separated between sepn1*RL and sepn2*RL
sepn1= 2;
sepn2= 3;
sepn = sepn1;
sep1 = sepn1*RL;
sep2 = sepn2*RL;
supp = sep1;
while supp(end) < 1-2*sep2
    move = sep1+rand*(sep2-sep1);
    supp = [supp ; supp(end)+move];
end
sparsity = length(supp);

%% sample range
k         = -fc:1:fc;
n         = 2*fc+1;

        
%% real-valued frequencies with random signs
amp  = rand(sparsity,1);
amp  = (amp-min(amp))/(max(amp)-min(amp));      %[0,1]
amp  = 1 + amp*(range-1);
sig  = (binornd(1,1/2,sparsity,1)-0.5)*2;
x    = amp.*sig;%.*exp(2*pi*1i*rand(sparsity,1));
xabs = abs(x);



%% generate measurement
F = exp(-1i*2*pi*k'*supp'); % Fourier matrix
y = F*x;
        
%% add noise
sigma    = noise/100*norm(y)/sqrt(2*n);
z        = normrnd(0,sigma,n,1) + 1i*normrnd(0,sigma,n,1);
y        = y + z;
if noise > 0                       % ||z|| <= ep
    ep = sqrt(2*n)*sigma;
else
    ep = 10^(-12);
end
fprintf('noise to signal ratio = %6.2fp \n',100*norm(z)/norm(y))
        


%% discretize the problem on a fine grid
N = 2*fc+1;            % the number of rows of sensing matrix A
F = 20;                % refinement factor / super-resolution factor
M = N * F;             % the number of columns of sensing matrix A
        
%% construct the sensing matrix A
A = exp(-1i*2*pi*(-fc:fc)'*(0:M-1)/M);  



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% OMP
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
[xrec , supprec , r ] = OMP( A , y , sparsity , ep);    
% xrec:    recovered 
% supprec: recovered indices of nonzero frequencies on the fine grid
% r = b - A * xrec: residual
if size(supprec,1)~=size(supp,1)    % make supprec be a column vector
    supprec = supprec.';
end
supprec = (supprec-1)/M;               % recover frequencies on [0,1)
dist1  = H_dist(supp,supprec)*(2*fc);  % Hausdorff distance
res1   = norm(r) / norm(y);            % residual


% plot
figure
stem(supp*(2*fc),x,'MarkerSize',5,'Color','r','LineWidth',2);
hold on
stem(supprec*(2*fc),xrec(supprec*M+1),'MarkerSize',5,'Color','b','Marker','x','LineWidth',2);
legend('exact','recovered')
title('OMP')
xlabel(['dist = ' num2str(dist1) 'RL     residual = ' num2str(res1)])
axis tight

fprintf('OMP:          dist = %6.2f RL  res = %6.2f\n',dist1,res1)


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% BLOOMP
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% radius of excluded band and radius of locally optimized neighborhood
if sepn >= 3
    band = 2*F-1;
    neighbor = F-1;
else
    band = floor(sepn/2 * F - 1);      % radius of local optimization = 1 * rayleigh length
    neighbor = floor(sepn/2 * F - 1);      % radius of excluded band = 2 * rayleigh length
end       

[xrec , supprec , r ] = BLOOMP(A , y , sparsity , band , neighbor);
% xrec:    recovered 
% supprec: recovered indices of nonzero frequencies on the fine grid
% r = b - A * xrec: residual
if size(supprec,1)~=size(supp,1)    % make supprec be a column vector
    supprec = supprec.';
end
supprec = (supprec-1)/M;               % recover frequencies on [0,1)
dist2  = H_dist(supp,supprec)*(2*fc);  % Hausdorff distance
res2   = norm(r) / norm(y);            % residual

% plot
figure
stem(supp*(2*fc),x,'MarkerSize',5,'Color','r','LineWidth',2);
hold on
stem(supprec*(2*fc),xrec(supprec*M+1),'MarkerSize',5,'Color','b','Marker','x','LineWidth',2);
legend('exact','recovered')
title('BLOOMP')
xlabel(['dist = ' num2str(dist2) 'RL     residual = ' num2str(res2)])
axis tight



fprintf('BLOOMP:       dist = %6.2f RL  res = %6.2f\n',dist2,res2)




