clc; 
clear all
close all


N = 512; % the number of the antennas at the BS

c = 3e8; %the speed of light
fc = 60e9; % the frequency
B = 1e8;
M = 128;
lambda = c/fc; %the wave length
d = 0.5*lambda; % the antenna spacing lambda/2


f = zeros(1, M);
topk = 1;
for m= 1:M
    f(m)=fc+B/(M)*(m-1-(M-1)/2);
end
Aperture = N*d
Rayleigh_dis = 2 * (N*d)^2/lambda
delta = 0.25;


SNR_dB = [5];  
SNR_linear = 10.^(SNR_dB/10.);


%%% 生成普通近场码本
rho_min = 20;
rho_max = Rayleigh_dis*0.375;
sector = 1;
uniform = 1;
eta = 1.2;
grid_r = 16;
[near_codebook,po_label] = generate_near_field_codebook2(N, d, lambda, eta,...
    rho_min, rho_max,sector,N,uniform,grid_r);
near_codebook = near_codebook';

overhead_max = size(near_codebook,1);


%% 远场分层码字 分3层
theta_max = pi/2;
theta_min = -pi/2;
sin_theta1_min = sin(theta_min);
sin_theta1_max = sin(theta_max);
theta_num = [4,4,32];%[32, 16, 8];


%% load近场分层码字    
load('near_codebook_W1.mat')
load('near_codebook_W2.mat')
load('recordW1.mat')
load('recordW2.mat') 


r_min = rho_min;
r_max = 100;

%% 近场宽带码字
rho_min_wide = 5;
delta = 1.2;

theta_c = 0;
theta_max_wide = pi/3;
theta_min_wide = -pi/3;
[w_delay_near, sin_theta_near, alpha_near] = delay_polar_beam(N, theta_c, theta_max_wide, theta_min_wide, B, fc, M, d, rho_min_wide, delta);
S = size(alpha_near, 1)


sample = 2000;

sum_rate1 = zeros([sample,size(near_codebook,1)]);
sum_rate2 = zeros([sample,size(near_codebook,1)]);
sum_rate3 = zeros([sample,size(near_codebook,1)]);
sum_rate4 = zeros([sample,size(near_codebook,1)]);
sum_rate_far_hierarchy = zeros([sample,size(near_codebook,1)]);
sum_rate_near = zeros([sample,size(near_codebook,1)]);

sum_rate1_wide = zeros([sample,size(near_codebook,1)]);
sum_rate2_wide = zeros([sample,size(near_codebook,1)]);
sum_rate3_wide = zeros([sample,size(near_codebook,1)]);
sum_rate4_wide = zeros([sample,size(near_codebook,1)]);
sum_rate_far_hierarchy_wide = zeros([sample,size(near_codebook,1)]);
sum_rate_near_wide = zeros([sample,size(near_codebook,1)]);

far_codebook = exp(1i*2*pi*[-(N-1)/2:1:(N-1)/2]'*d*[-(N-1)/2:1:(N-1)/2]*(2/(N/1))/lambda);

rate_far_ex = zeros(sample,length(SNR_dB));
rate_near_ex = zeros(sample,length(SNR_dB));
rate_near_hi = zeros(sample,length(SNR_dB));
rate_far_hi = zeros(sample,length(SNR_dB));
rate_rainbow = zeros(sample,length(SNR_dB));

rate_opt = zeros(sample,length(SNR_dB));
overhead = zeros(sample,length(SNR_dB));

Distance = [r_min+5:10:75]

for t = 1:sample
    t
for s = 1:length(Distance)
%     s
     dis = Distance(s);
     
    [h, hc] = near_field_channel(N, d, fc, B, M, dis, dis,pi/2);
    H=hc;
    SNR = SNR_linear;
    
    %% far-field beam training 
    array_gain = 0;
    for i =1:length(far_codebook)
        if array_gain<=abs(far_codebook(i,:)*H)^2
            array_gain=max(array_gain,abs(far_codebook(i,:)*H)^2);
            array_gain_wide = abs(far_codebook(i,:)*h).^2;
            sum_rate1_wide(t,i) = sum(log2(1+SNR*abs(far_codebook(i,:)*h).^2));
        else
            sum_rate1_wide(t,i) = sum(log2(1+SNR*array_gain_wide));
        end
            sum_rate1(t,i) =  log2(1 + SNR * array_gain);
    end
    sum_rate1(t,length(far_codebook)+1:overhead_max) = log2(1 + SNR * array_gain);
    sum_rate1_wide(t,length(far_codebook)+1:overhead_max) = sum_rate1_wide(t,length(far_codebook));

    rate_far_ex(t,s) = sum_rate1_wide(t,end);
    
    %% near-field beam training 
    array_gain = 0;
    
    i_max_=0;
    for i =1:size(near_codebook,1)
         if array_gain<=abs(near_codebook(i,:)*H)^2
            array_gain=max(array_gain,abs(near_codebook(i,:)*H)^2);
            array_gain_wide = abs(near_codebook(i,:)*h).^2;
            sum_rate2_wide(t,i) = sum(log2(1+SNR*abs(near_codebook(i,:)*h).^2));
        else
            sum_rate2_wide(t,i) = sum(log2(1+SNR*array_gain_wide));
        end
         sum_rate2(t,i) =  log2(1 + SNR * array_gain);
    end
    

    rate_near_ex(t,s) = sum_rate2_wide(t,end);
    %% hierachy beamforming   
    
    array_gain = 0;array_gain_1 = 0;array_gain_2 = 0;array_gain_3 = 0;
    %%%% level one
    i_max_1 = 0;
    overhead1 = size(near_codebook_W1,1);
    for i =1:size(near_codebook_W1,1)
        if array_gain_1<=abs(near_codebook_W1(i,:)*H)^2
            array_gain_1=max(array_gain_1,abs(near_codebook_W1(i,:)*H)^2);
            array_gain_wide = abs(near_codebook_W1(i,:)*h).^2;
            i_max_1 = i;
            sum_rate3_wide(t,i) = sum(log2(1+SNR*abs(near_codebook_W1(i,:)*h).^2));
        else
            sum_rate3_wide(t,i) = sum(log2(1+SNR*array_gain_wide));
        end
        sum_rate3(t,i) =  log2(1 + SNR * array_gain_1);
    end
    po1 = recordW1(:,i_max_1);
    
    order1_themin = find(recordW2(1,:)>=po1(1));
    order1_themax = find(recordW2(2,:)<=po1(2));
    order1_rmin = find(recordW2(3,:)>=po1(3));
    order1_rmax = find(recordW2(4,:)<=po1(4));
    
    order1 = intersect(order1_themin,order1_themax);
    order1 = intersect(order1,order1_rmin);
    order1 = intersect(order1,order1_rmax);

    %%%% level two
    near_codebook2 = near_codebook_W2(order1,:);
    overhead2 = size(near_codebook2,1);
    record2 = recordW2(:,order1);
    i_max_2 = 0;
    for i =1:size(near_codebook2,1)
        if array_gain_2<=abs(near_codebook2(i,:)*H)^2
            array_gain_2=max(array_gain_2,abs(near_codebook2(i,:)*H)^2);
            array_gain_wide = abs(near_codebook2(i,:)*h).^2;
            i_max_2 = i;
            sum_rate3_wide(t,overhead1+i) = sum(log2(1+SNR*abs(near_codebook2(i,:)*h).^2));
        else
            sum_rate3_wide(t,overhead1+i) = sum(log2(1+SNR*array_gain_wide));
        end
        sum_rate3(t,overhead1+i) =  log2(1 + SNR * array_gain_2);
    end
    po2 = record2(:,i_max_2);
   
   
    order3_the = find((po_label(1,:)>=po2(1))&(po_label(1,:)<=po2(2)));
    order3_r = find((po_label(2,:)>=po2(3))&(po_label(2,:)<=po2(4)));
    order3 = intersect(order3_the,order3_r);
% 

    %%% level four
    near_codebook4 = near_codebook(order3,:);
    overhead3 = size(near_codebook4,1);
    record4 = po_label(:,order3);
    i_max = 0;
    for i =1:size(near_codebook4,1)
        if array_gain<=abs(near_codebook4(i,:)*H)^2
            array_gain=max(array_gain,abs(near_codebook4(i,:)*H)^2);
            i_max = i;
            array_gain_wide = abs(near_codebook4(i,:)*h).^2;
            sum_rate3_wide(t,overhead1+overhead2+i) = sum(log2(1+SNR*abs(near_codebook4(i,:)*h).^2));
        else
            sum_rate3_wide(t,overhead1+overhead2+i) = sum(log2(1+SNR*array_gain_wide));
        end
        sum_rate3(t,overhead1+overhead2+i) =  log2(1 + SNR * array_gain);
    end
    sum_rate3(t,overhead1+overhead2+overhead3+1:overhead_max) = log2(1 + SNR * array_gain);
    sum_rate3_wide(t,overhead1+overhead2+overhead3+1:overhead_max) = sum_rate3_wide(t,overhead1+overhead2+overhead3);

    overhead(t,s) = overhead1+overhead2+overhead3;
    
    rate_near_hi(t,s) = sum_rate3_wide(t,end);
        %% 远场分级
        
    [rate_far_hierarchy_iter,rate_far_hierarchy_iter_wide,~] = training_far_hierarchy(N, H.',h,sin_theta1_min,...
        sin_theta1_max,theta_num,SNR, 10*log10(SNR), 1, 2, M);
    sum_rate_far_hierarchy(t,1:sum(theta_num)) = rate_far_hierarchy_iter;
    temp = max(rate_far_hierarchy_iter);
    sum_rate_far_hierarchy(t,sum(theta_num)+1:overhead_max) = temp;
    sum_rate_far_hierarchy_wide(t,1:sum(theta_num))=rate_far_hierarchy_iter_wide;
    sum_rate_far_hierarchy_wide(t,sum(theta_num)+1:overhead_max)=max(rate_far_hierarchy_iter_wide);
   
    rate_far_hi(t,s) = max(rate_far_hierarchy_iter_wide);
    
    
        %% 近场彩虹
    [rate_near_iter,rate_near_iter_wide] = training_near_rainbow(N, B, fc, f, M, d, S, h.', w_delay_near,...
        theta_c, sin_theta_near, alpha_near, SNR, 10*log10(SNR), 2, 1, topk);
    sum_rate_near(t,1:S) = rate_near_iter;
    sum_rate_near_wide(t,1:S) = rate_near_iter_wide;
    sum_rate_near(t, S+1:overhead_max) = rate_near_iter(end);
    sum_rate_near_wide(t,S+1:overhead_max) = rate_near_iter_wide(end);

    rate_rainbow(t,s) = rate_near_iter_wide(end);
    
    %% perfect CSI based beamforming
    wc_opt = 1/sqrt(N)*exp(1j*angle(H));
    array_gain = abs(wc_opt.*H).^2;   
    rate_opt(t,s)=sum(log2(1+SNR*sum(abs(H)).^2));
end
end

mean(overhead)

far_exhau = mean(rate_far_ex/M);
near_exhau = mean(rate_near_ex/M);
near_hiera = mean(rate_near_hi/M);
far_hiera = mean(rate_far_hi/M);
near_wide = mean(rate_rainbow/M);


figure('color',[1,1,1]); hold on; box on; grid on;
ha=gca;
plot(Distance,far_exhau,'m^-', 'Linewidth', 1.2)
hold on
plot(Distance,far_hiera,'->','Color',[0.18039 0.5451 0.34118],'Linewidth',1.2)
hold on
plot(Distance,near_exhau,'b*-','Linewidth',1.2)
hold on
plot(Distance,near_wide,'-s','Color',[1.0 0.5 0],'Linewidth',1.2)
hold on
plot(Distance,near_hiera,'rd-','Linewidth',1.2)

legend('Far-field exhaustive search beam training','Far-field hierarchical beam training',...
    'Near-field exhaustive search beam training','Time delay based near-field beam training',...
    'Proposed near-field 2D hierarchical beam training')


xlabel('Distance (m)')
ylabel('Achievable Rate (bit/s/Hz)')


ylim([4.0,10.5])
