%% SU-MIMO channel estimation using codebook
% This code estimates the AoA, AoD, and complex path gain of the
% unquantized received data X.
% Bartlett beamformin is used to estimate AoAs
% Beamforming with designed codebook (eqn/section number) is used to
% estimate AoDs
% Least squares is used to estimate complex path gains

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


%%%% Input parameters  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% X - unquantized received signal at the BS (N_r \times N) where N is the
% nuber of snapshots
% d_r : inter-element spacing at the BS
% d_t : inter-element spacing at the UE 
% L : number of paths
% S : Pilot matrix
% res: search resolution of beamformer to estimate AoA
% M : number of grid points used to estimate AoD

%%%%%%%% Output parameters %%%%%%%%%%%%%%%%%%%%%%%%%%%%
% aoa_est : Estimated AoAs
% aod_est : Estimated AoDs
% alpha_est : Estimated path gains



function [aoa_est aod_est alpha_est] = SU_MIMO_ch_est_bf_codebook(X,d_r,L,S,res,M, d_t)


u_grid = linspace(-1,1,M); % search grid for AoD
[N_RX N] = size(X);
[N_TX N] =size(S);

% X1  = X*inv(S);
X1 = X*S';
R_doa = (1/N).*X1*X1'; % estimated data-covariance matrix at BS



%% Angle estimation

% AoA estimation
aoa_search = -90:res:90; % search grid for AoAs
cost_aoa_a = zeros(length(aoa_search),1);

for k = 1:length(aoa_search)
    a_doa = gen_a(N_RX, d_r, aoa_search(k));      
    cost_aoa_a(k) = abs(a_doa'*R_doa*a_doa); % Bartlett beamformer spectrum
end

[val pos] = local_max(cost_aoa_a,L); % finding peaks of the bartlett spectrum
aoa_est = sort(aoa_search(pos)); % estimated AoAs

%% AoD estimation
% we use proposed codebook to estimate the AoDs

aod_est_dc_ind = zeros(1,L);
K = 2; %  Beams per stage
Nstage = log2(M);

aod_check_a = zeros(K,1);

for iter=1:L
    
    d_index = 1; 
    w_mrc = gen_a(N_RX, d_r, aoa_est(iter));
    
    for s = 1:Nstage 
        
        W = get_code_TX(N_TX,d_t,M,K,s);
        pos_req(1:K) = 2*(d_index-1)+(1:K); 

        for j=1:K
            f_tx = W(:,pos_req(j));       
            aod_check_a(j) = abs((w_mrc'*X1*f_tx))^2; % beamformer spectrum
        end

        [val c_index] = max(aod_check_a);
        d_index = 2*(d_index-1) + c_index; % Update beamformer to be used in next stage

    end
    aod_est_dc_ind(iter) = (d_index); % estimated AoD
end

aod_est = asind(sort(u_grid(aod_est_dc_ind)));


%% Path gain estimation

B_aux = sqrt(1/L).*(kron(transpose(S),eye(N_RX)))*krp(conj(gen_a(N_TX, d_t, aod_est)),gen_a(N_RX,d_r,aoa_est));
alpha_est = B_aux\reshape(X, N*N_RX,1);



end

