%% Simulate the AMP algorithm. 

K = 6;
N = 500;
M = 50;

% sparse sensing matrix.
A = zeros(M, N);

rng(0);

% Compressing matrix A: Ensure E[|Ax|^2] == 1, when E[|x|^2] == 1. 
for n = 1:N
    A(:,n) = (randn([M,1])+1j*randn([M,1]))/sqrt(2)/sqrt(M);    % Energy: 1/M each. 
end

% randomly generate a K-sparse vector x.

SNR_range = flip(5:2:20);
len_SNR_range = length(SNR_range);


lambda = 1.7;
N_shot = 50;
MSE = zeros(len_SNR_range, 3);
    
for snr_idx = 1:len_SNR_range
    MSE_arr = zeros(N_shot, 3);
    SNR = SNR_range(snr_idx);   % in dB.
    
    for idx = 1:N_shot
        x = zeros(N,1);
        s = randperm(N);
        s = s(1:K);

        b = false(N,1); 
        b(s) = true;        

        x(b) = (randn([K,1])+1j*randn([K,1]))/sqrt(2)/sqrt(K);
        sigma = 1/db2mag(SNR);
        
        y = A*x + sigma*(randn([M,1])+1j*randn([M,1]))/sqrt(2)/sqrt(M); 

        % Perform AMP.
        x_hat = camp(y, A, lambda);
        x_hat2 = camp2(y, A, lambda);
        x_hat3 = uamp(y, K, A, 200); 

        % Estimate the error. 
        MSE_arr(idx, 1) = norm(x-x_hat)^2;
        MSE_arr(idx, 2) = norm(x-x_hat2)^2;
        MSE_arr(idx, 3) = norm(x-x_hat3)^2; 
    end
    MSE(snr_idx, :) = mean(MSE_arr, 1);
    fprintf('SNR = %.2f dB complete with avg. MSE1 = %.2f dB, MSE2 = %.2f dB, MSE3 = %.2f dB\n', ...
    SNR_range(snr_idx), pow2db(MSE(snr_idx,1)), pow2db(MSE(snr_idx, 2)), pow2db(MSE(snr_idx, 3)));
end


%% Visualization.
set(0,'DefaultLineMarkerSize',6);
set(0,'DefaultTextFontSize',14);
set(0,'DefaultAxesFontSize',12);
set(0,'DefaultLineLineWidth',1.4);
set(0,'defaultfigurecolor','w');

figure('color',[1 1 1]); box on; 

plot(SNR_range, pow2db(MSE(:,1))); hold on;
plot(SNR_range, pow2db(MSE(:,2)));
plot(SNR_range, pow2db(MSE(:, 3))); 
grid on;  xlabel('SNR (dB)'); ylabel('MSE(x) (dB)'); 
title('AMP performance');
legend({'CAMP1', 'CAMP2', 'AMP-Git'});


