%% MU-MIMO channel estimation illustration
% Example code that illustrate the proposed channel estimation algorithm
% This code perform Monte Carlo simulation to generate data used to obtain
% Fig 8. 

% Data generation with a large number of monte-carlo runs can take
% excessive amount of time. Kindly use the "Fig8x.m" to generate the curve from the provided data.

% For any queries, please contact R.S. Prasobh Sankar (prasobhsankar1@gmail.com)

clear all
close all
clc


N_r = 128;
K = 8; % number of UEs
L_k = 1; % number of paths per UE
no_sim = 10; % number of Monte Carlo runs


fade = 0;   % parameter for truncated Gaussian fading (if needed)
% fade - nature of fade. 
%       0 is pure Los (unit modulus path gains)
%       1 is Rayleigh fading
%       2 is arbitrary fading

dist_min = 0.5;
pd = makedist('Normal',0,sqrt(0.5));
trunc_pd = truncate(pd,dist_min,inf);
% Example: 
%     pd = makedist('Normal',0,sqrt(0.5));  trunc_pd = truncate(pd,dist_min,inf);
%     Results in real and imaginary parts of path gains (say, \alpha) to follow a
%     truncated Gaussian distribution with 0 <= |Re(\alpha)|,|Im(\alpha)| <= dist_min 


genie_aided_correlation = 0; 
% choose genie_aided_correlation = 1 to use the true channel correlation to
% compute the Bussgang decomposition. However, this is not practical in
% real life. 

OSR = 4; % spatial oversampling ratio at BS
delta = 0.5/OSR; % inter-element spacing at BS
N = K;  % Number of snapshots

aoa_grid = -70:1:70; % search grid for AoA
aspread_aoa = 45; % AoAs will be drawn from -aspread_aoa to aspread_aoa
spacing_aoa = 20; % minimum spacing between AoAs of same UE

aoa_grid_app = -aspread_aoa:1:aspread_aoa;
L_app = length(aoa_grid_app);



E_s = 1; % Signal power
SNR_dB = -12:4:20; 
SNR = 10.^(0.1.*SNR_dB);

tic

U = tril(ones(N_r));
V = U - eye(N_r);
U_inv = inv(U);





% Channel estimation errors
Channel_MSE_uq_lmmse = zeros(1,length(SNR));
Channel_MSE_uq_angular = zeros(1,length(SNR));
Channel_MSE_sd_Bussgang = zeros(1,length(SNR));
Channel_MSE_sd = zeros(1,length(SNR));
Channel_MSE_ob_blmmse = zeros(1,length(SNR));


for iter = 1:length(SNR)
    
     
    channel_err_uq_lmmse = 0;
    channel_err_uq_angular = 0;
    channel_err_sd_Bussgang = 0;
    channel_err_sd = 0;
    channel_err_ob_blmmse = 0;
      
    ch_err_dr = 0;
    
    % selection of voltage level (b)
    b = 3*sqrt((SNR(iter)+1)/2);

    U = tril(ones(N_r));
    U_inv = inv(U);
    R_n = eye(N_r) + ((2*(b^2))/(3)).*U_inv*U_inv'; % Noise covariance
    R_n_pw = inv(sqrtm(R_n)); % pre-whitening matrix

    
    for loop_inner = 1:no_sim
        
        
       % Get channel matrix
         [H_true doa_true alpha_true] = get_MUMIMO_channel(K,L_k,...
          N_r,delta,aspread_aoa, spacing_aoa, fade, trunc_pd);
       
        A_RX_temp = sqrt(1/L_app).*gen_a(N_r, delta, aoa_grid_app); 
      
        
        ch_err_dr = ch_err_dr + norm(H_true,'fro')^2;
        
        S = sqrt((SNR(iter))/K).*dftmtx(K); % pilot matrix
        N_noise = sqrt(1/2).*(randn(N_r,N) + (1j).*randn(N_r,N));
        
        X = H_true*S + N_noise; % unquantized received signal
        Y_q = sqrt(1/2).*sgn(X);  %One-bit 
        
        %% prameters for Bussgang decomposition based channel estimation [Ref. no]
        
        if(genie_aided_correlation == 0)
            C_G = A_RX_temp*A_RX_temp';
            C_g = kron(eye(K),C_G);
        elseif(genie_aided_correlation == 1)
            C_g = zeros(K*N_r);        
            for k=1:K
                range_consider = (k-1)*N_r+1:k*N_r;
                A_req = sqrt(1/L_k).*gen_a(N_r, delta, doa_true(:,k));
                C_g(range_consider,range_consider) = A_req*A_req';
            end 
        end

        P = kron(transpose(S), eye(N_r));
        C_x = P*C_g*P' + eye(N*N_r);
        beta = 1.05;

        C_r = zeros(N*N_r,N*N_r);


        sig_r_sq = zeros(N*N_r,1);
        sig_q_sq = zeros(N*N_r,1);
        sig_x_sq = diag(C_x);
        sig_y_sq = zeros(N*N_r,1);

        for k=1:N*N_r

            if(mod(k,N_r)==1)
                sig_r_sq(k) = sig_x_sq(k);
            else
                sig_r_sq(k) = sig_x_sq(k) + sig_q_sq(k-1);
            end

            sig_y_sq(k) = (pi/2)*(beta^2)*sig_r_sq(k);
            sig_q_sq(k) = sig_y_sq(k) - sig_r_sq(k);

        end

        r_t = sig_r_sq(1:N_r);
        a_m = zeros(N_r,1);
        a_m(1:N_r) = (beta/2).*sqrt(pi.*r_t(1:N_r)); % proposed voltage level in [Number]
        
        U_inv = kron(eye(N),inv(U));

        C_ql = diag(sig_q_sq);

        C_y = C_x + U_inv*C_ql*U_inv'; % Eq (31)

        Yl = sigma_delta_ADC_b(X,0,delta,a_m); % output of SD ADC proposed in []
        yl = reshape(Yl, N*N_r,1);

        C_gy = C_g*P';
        h_hat = C_gy*inv(C_y)*yl;
        H_est_sd_Bussgang = reshape(h_hat, N_r, K); % Estimated channel []
        
        channel_err_sd_Bussgang = channel_err_sd_Bussgang + ((norm(H_true - H_est_sd_Bussgang,'fro'))^2);
        
        %%%%%%%%%% Unquantized channel estimation - LMMSE
        x = reshape(X, N*N_r, 1);   
            
        C_gx = C_g*P';
        h_hat_lmmse = C_gx*inv(C_x)*x;
        H_est_lmmse = reshape(h_hat_lmmse, N_r, K); % LMMSE channel estimate    
        channel_err_uq_lmmse = channel_err_uq_lmmse + ((norm(H_true - H_est_lmmse,'fro'))^2);
        
        %% Proposed methods 
        
         T_s = 1;
        % UQ channel estimation        
        [H_est_uq_angular  doa_est_uq_angular alpha_est_uq_angular] = MU_MIMO_est_UQ(H_true,S, T_s, delta, L_k, SNR(iter), aoa_grid );
        channel_err_uq_angular = channel_err_uq_angular + norm(H_true - H_est_uq_angular,'fro')^2;

        % SD channel estimation
        [H_est_sd  doa_est_sd alpha_est_sd]  = MU_MIMO_est_SD(H_true, S, T_s, delta, L_k, SNR(iter), aoa_grid, b, R_n_pw );
        channel_err_sd = channel_err_sd + norm(H_true - H_est_sd,'fro')^2; 
      
               
        %% One bit BLMMSE channel estimation
        
        H_est = one_bit_ch_estm_blmmse(X,S,K,C_g);
        channel_err_ob_blmmse = channel_err_ob_blmmse + ((norm(H_true - H_est,'fro'))^2);
    
        [iter loop_inner]
        
    end
    
    
    
    Channel_MSE_uq_lmmse(iter) = channel_err_uq_lmmse/ch_err_dr;
    Channel_MSE_uq_angular(iter) = channel_err_uq_angular/ch_err_dr;
    Channel_MSE_sd_Bussgang(iter) = channel_err_sd_Bussgang/ch_err_dr;
    Channel_MSE_sd(iter) = channel_err_sd/ch_err_dr;
    Channel_MSE_ob_blmmse(iter) = channel_err_ob_blmmse/ch_err_dr;
    
end

toc
%% Plotting Results
% 

figure
plot(SNR_dB , 10.*log10(Channel_MSE_uq_lmmse), 'b-o','LineWidth',1.2);
hold on
plot(SNR_dB , 10.*log10(Channel_MSE_uq_angular), 'g-d','LineWidth', 1.2);
hold on
plot(SNR_dB , 10.*log10(Channel_MSE_sd_Bussgang), 'r-o','LineWidth',1.2);
hold on
plot(SNR_dB ,  10.*log10(Channel_MSE_sd), 'm-d','LineWidth', 1.2);
hold on
plot(SNR_dB , 10.*log10(Channel_MSE_ob_blmmse), 'k-o','LineWidth',1.2);
grid on

legend('LMMSE','Analog-angular','SD-Bussgang', 'SD-proposed','BLMMSE');
title('Channel Estimation');
xlabel('SNR (dB)');
ylabel('MSE (dB)');

