% This matlab code computes multiuser rate performance as in Fig 7.
% Yuan-Pei Lin, ypl@mail.nctu.edu.tw

clearvars;
close all;

t=32;  %number of transmit antennas
r=8; %number of receive antennas
M=1;
%Nrf=2*M;
maxK=16; %number of users
subK=[1 4:4:maxK];

cl=3;
ray=10;
path=cl*ray;

ULA=0;
    
   sigma_ta=7.5*ones(1,cl)/180*pi;% standard deviation of phi_t
   sigma_te=7.5*ones(1,cl)/180*pi;% standard deviation of phi_t

   sigma_ra=7.5*ones(1,cl)/180*pi;
   sigma_re=7.5*ones(1,cl)/180*pi;

    if ceil(sqrt(t))==sqrt(t)
        
       Wt=sqrt(t);
       heightt=Wt;
    else
        Wt=sqrt(2*t);
        heightt=Wt/2;
    end   
   

    if ceil(sqrt(r))==sqrt(r)
        
       Wr=sqrt(r);
       heightr=Wr;
    else
        Wr=sqrt(2*r);
        heightr=Wr/2;
    end   




ga=-log(5*10^(-5))/1.5;

dt=1/2; %normalized transmit antenna spacing
dr=1/2; %normalized receive antenna spacing
kr=(0:r-1)';
kt=(0:t-1)';


N_theta=8;N_phi=N_theta;

N_theta2=N_theta*2;N_phi2=N_phi/2;
N_theta2less=N_theta;N_phi2less=N_phi/4;
N_theta2less2=N_theta/2;N_phi2less2=N_phi/4;

N_theta1=N_theta; N_phi1=N_theta; 
N_theta1less=N_theta/2; N_phi1less=N_theta/2; 
N_theta1less2=N_theta/2; N_phi1less2=N_theta/4;
snr_range=(-10:3:5);
pt=10.^(snr_range/10);

ptK=10.^(5/10);

phi_itr=100;mu_itr=10;
itr=phi_itr*mu_itr;


r0=zeros(1,length(pt));r1=r0; r2=r0;r2ba=r0;r2less=r0;r2less2=r0;r1less2=r0;r1less=r0;
r0K=zeros(1,length(subK));r1K=r0K; r2K=r0K;r2baK=r0K;r2lessK=r0K;r2less2K=r0K;r1less2K=r0K;r1lessK=r0K;
mu_ta=zeros(cl,maxK);mu_te=zeros(cl,maxK);mu_ra=zeros(cl,maxK);mu_re=zeros(cl,maxK);
F0=zeros(t,maxK); F1=F0; F1less=F0;F1less2=F0;F2=F0; F2ba=F0; F2less=F0;F2less2=F0;
Geq0=zeros(maxK,t);Geq1=Geq0;Geq1less=Geq0; Geq1less2=Geq0; 
Geq2=Geq0; Geq2less=Geq0; Geq2less2=Geq0; Geq2ba=Geq0;
noise0=zeros(maxK,1);
noise1=noise0;
noise1less=noise0;
noise1less2=noise0;

noise2=noise0;
noise2less=noise0;
noise2less2=noise0;
noise2ba=noise0;

rng(1970);

for n_mu=1:mu_itr
    for n_k=1:maxK
    mu_ta(:,n_k)=2*pi*rand(cl,1); % mean of azimuth angle at tx
    mu_te(:,n_k)=2*pi*rand(cl,1); % mean of elevation angle at tx
    mu_ra(:,n_k)=2*pi*rand(cl,1); % mean of azimuth angle at rx
    mu_re(:,n_k)=2*pi*rand(cl,1); %mean of elevation angle at rx
    end

for n_itr=1:phi_itr
    
    for n_k=1:maxK
    
   [H, At, Ar]=mmwCh(ULA,dr,dt,t,r,Wt,heightt,Wr,heightr, path, cl,ray,mu_ta(:,n_k), mu_te(:,n_k), mu_ra(:,n_k), mu_re(:,n_k), sigma_ta,sigma_te,sigma_ra,sigma_re);

   
   [U,La,V]=svd(H);

   Ph1=zeros(t,M);Ph1less=zeros(t,M);Ph1less2=zeros(t,M);
   for i2=1:M
       Ph1(:,i2)=V(:,i2)-Ph1*pinv(Ph1)*V(:,i2);
       s=max(abs(Ph1(:,i2)))/2; % for scaling f0, extra 2 s.t. max(abs(f0))=2
       Ph1(:,i2)=Ph1(:,i2)/s;
       Ph1less(:,i2)=V(:,i2)-Ph1less*pinv(Ph1less)*V(:,i2);
       s=max(abs(Ph1less(:,i2)))/2;
       Ph1less(:,i2)=Ph1less(:,i2)/s;
       Ph1less2(:,i2)=V(:,i2)-Ph1less2*pinv(Ph1less2)*V(:,i2);
       s=max(abs(Ph1less2(:,i2)))/2; % for scaling f0, extra 2 s.t. max(abs(f0))=2
       Ph1less2(:,i2)=Ph1less2(:,i2)/s;
        for i1=1:t 
            %[~,i0]=min(abs(tab1-Ph1(i1,i2)));  Ph1(i1,i2)=tab1(i0);
            Ph1(i1,i2)=q1(Ph1(i1,i2),N_theta1,N_phi1);
            Ph1less(i1,i2)=q1(Ph1less(i1,i2),N_theta1less,N_phi1less);
            Ph1less2(i1,i2)=q1(Ph1less2(i1,i2),N_theta1less2,N_phi1less2);
        end
   end
   
   
   Ph2=zeros(t,M);
   Ph2ba=zeros(t,M);
   Ph2less=zeros(t,M);
   Ph2less2=zeros(t,M);
   % quantize the subsequent vectors with orthogonal projection
   for i2=1:M
       Ph2(:,i2)=V(:,i2)-Ph2*pinv(Ph2)*V(:,i2);
       Ph2ba(:,i2)=V(:,i2)-Ph2ba*pinv(Ph2ba)*V(:,i2);
       Ph2less(:,i2)=V(:,i2)-Ph2less*pinv(Ph2less)*V(:,i2);
       Ph2less2(:,i2)=V(:,i2)-Ph2less2*pinv(Ph2less2)*V(:,i2);
       Ph2(:,i2)=Ph2(:,i2)/max(abs(Ph2(:,i2)))*2/1.0000001;
       Ph2ba(:,i2)=Ph2ba(:,i2)/max(abs(Ph2ba(:,i2)))*2/1.0000001;
       Ph2less(:,i2)=Ph2less(:,i2)/max(abs(Ph2less(:,i2)))*2/1.0000001;
       Ph2less2(:,i2)=Ph2less2(:,i2)/max(abs(Ph2less2(:,i2)))*2/1.0000001;
        for i1=1:t 
 
            Ph2(i1,i2)=q2(Ph2(i1,i2),N_theta,N_phi);
            Ph2ba(i1,i2)=q2(Ph2ba(i1,i2),N_theta2,N_phi2); 
            Ph2less(i1,i2)=q2(Ph2less(i1,i2),N_theta2less,N_phi2less);
            Ph2less2(i1,i2)=q2(Ph2less2(i1,i2),N_theta2less2,N_phi2less2);
        end
   end

   
   F0(:,n_k)=V(:,1);
   F1(:,n_k)=Ph1;
   F1less(:,n_k)=Ph1less;
   F1less2(:,n_k)=Ph1less2;
   F2(:,n_k)=Ph2;
   F2ba(:,n_k)=Ph2ba;
   F2less(:,n_k)=Ph2less;
   F2less2(:,n_k)=Ph2less2;

   G0=U(:,1:M)';
   G1zf=G0;
   G2zf=G0;
   G2bazf=G0;
   G1zfless=G0;
   G1zfless2=G0;
   G2zfless=G0;
   G2zfless2=G0;
   
   
   Phg1=zeros(r,M);
   G1zf=G1zf';
   Phg1less=zeros(r,M);
   G1zfless=G1zfless';
   Phg1less2=zeros(r,M);
   G1zfless2=G1zfless2';
   for i2=1:M
       Phg1(:,i2)=G1zf(:,i2)-Phg1*pinv(Phg1)*G1zf(:,i2);
       s=max(abs(Phg1(:,i2)))/2; % for scaling f0, extra 2 s.t. max(abs(f0))=2
       Phg1(:,i2)=Phg1(:,i2)/s;
       Phg1less(:,i2)=G1zfless(:,i2)-Phg1less*pinv(Phg1less)*G1zfless(:,i2);
       s=max(abs(Phg1less(:,i2)))/2; % for scaling f0, extra 2 s.t. max(abs(f0))=2
       Phg1less(:,i2)=Phg1less(:,i2)/s;
 
       Phg1less2(:,i2)=G1zfless2(:,i2)-Phg1less2*pinv(Phg1less2)*G1zfless2(:,i2);
       s=max(abs(Phg1less2(:,i2)))/2; % for scaling f0, extra 2 s.t. max(abs(f0))=2
       Phg1less2(:,i2)=Phg1less2(:,i2)/s;
        for i1=1:r
 
            Phg1(i1,i2)=q1(Phg1(i1,i2),N_theta,N_phi);
            Phg1less(i1,i2)=q1(Phg1less(i1,i2),N_theta1less,N_phi1less);
            Phg1less2(i1,i2)=q1(Phg1less2(i1,i2),N_theta1less2,N_phi1less2);
        end
   end     
   Phg1=Phg1';
   Phg1less=Phg1less';
   Phg1less2=Phg1less2';

   
   Phg2=zeros(r,M);
   G2zf=G2zf';
   for i2=1:M
       Phg2(:,i2)=G2zf(:,i2)-Phg2*pinv(Phg2)*G2zf(:,i2);
       s=max(abs(Phg2(:,i2)))/2; % for scaling f0, extra 2 s.t. max(abs(f0))=2
       Phg2(:,i2)=Phg2(:,i2)/s;
        for i1=1:r 
            Phg2(i1,i2)=q2uniform(Phg2(i1,i2),N_theta,N_phi);
        end
   end     
   Phg2=Phg2';
   G2zf=G2zf';
   
   Phg2ba=zeros(r,M);
   G2bazf=G2bazf';
   for i2=1:M
       Phg2ba(:,i2)=G2bazf(:,i2)-Phg2ba*pinv(Phg2ba)*G2bazf(:,i2);
       Phg2ba(:,i2)=Phg2ba(:,i2)/max(abs(Phg2ba(:,i2)))*2/1.000001;
        for i1=1:r 
            Phg2ba(i1,i2)=q2(Phg2ba(i1,i2),N_theta2,N_phi2);
        end
   end     
   Phg2ba=Phg2ba';
   
   Phg2less=zeros(r,M);
   G2zfless=G2zfless';
   for i2=1:M
       Phg2less(:,i2)=G2zfless(:,i2)-Phg2less*pinv(Phg2less)*G2zfless(:,i2);
       Phg2less(:,i2)=Phg2less(:,i2)/max(abs(Phg2less(:,i2)))*2/1.000001;
        for i1=1:r 
            Phg2less(i1,i2)=q2(Phg2less(i1,i2),N_theta2less,N_phi2less);
        end
   end     
   Phg2less=Phg2less';
   
   Phg2less2=zeros(r,M);
   G2zfless2=G2zfless2';
   for i2=1:M
       Phg2less2(:,i2)=G2zfless2(:,i2)-Phg2less2*pinv(Phg2less2)*G2zfless(:,i2);
       Phg2less2(:,i2)=Phg2less2(:,i2)/max(abs(Phg2less2(:,i2)))*2/1.000001;
        for i1=1:r 
            Phg2less2(i1,i2)=q2(Phg2less2(i1,i2),N_theta2less2,N_phi2less2);
        end
   end     
   Phg2less2=Phg2less2';
   
 
   
   Geq0(n_k,:)=G0*H;
   Geq1(n_k,:)=Phg1*H;
   Geq2(n_k,:)=Phg2*H;
   Geq1less(n_k,:)=Phg1less*H;
   Geq1less2(n_k,:)=Phg1less2*H;
   Geq2ba(n_k,:)=Phg2ba*H;
   Geq2less(n_k,:)=Phg2less*H;
   Geq2less2(n_k,:)=Phg2less2*H;

   
   
   
   noise0(n_k)=norm(G0)^2;   
   noise1(n_k)=norm(Phg1)^2;
    noise1less(n_k)=norm(Phg1less)^2;
    noise2less(n_k)=norm(Phg2less)^2;
        noise1less2(n_k)=norm(Phg1less2)^2;
    noise2less2(n_k)=norm(Phg2less2)^2;
   noise2(n_k)=norm(Phg2)^2; 
   noise2ba(n_k)=norm(Phg2)^2;  
    end
    
   for nn=1:length(subK)
       
       n_subK=subK(nn);
   Heq0=Geq0(1:n_subK,:)*F0(:,1:n_subK);
   Heq1=Geq1(1:n_subK,:)*F1(:,1:n_subK);
   Heq2=Geq2(1:n_subK,:)*F2(:,1:n_subK);
   Heq2ba=Geq2ba(1:n_subK,:)*F2ba(:,1:n_subK);
   Heq1less=Geq1less(1:n_subK,:)*F1less(:,1:n_subK);
   Heq2less=Geq2less(1:n_subK,:)*F2less(:,1:n_subK);
      Heq1less2=Geq1less2(1:n_subK,:)*F1less2(:,1:n_subK);
   Heq2less2=Geq2less2(1:n_subK,:)*F2less2(:,1:n_subK);
   
   Fbb0=(pinv(Heq0'))';
   Fbb1=(pinv(Heq1'))';
   Fbb2=(pinv(Heq2'))';   
   Fbb2ba=(pinv(Heq2ba'))';
   Fbb1less=(pinv(Heq1less'))';
   Fbb2less=(pinv(Heq2less'))';
      Fbb1less2=(pinv(Heq1less2'))';
   Fbb2less2=(pinv(Heq2less2'))';
%    display('check if Heq2ba*Fbb2ba=I');
%    Heq2ba*Fbb2ba
   
   for n_k=1:n_subK
       Fbb0(:,n_k)=Fbb0(:,n_k)/norm(F0(:,1:n_subK)*Fbb0(:,n_k));
       Fbb1(:,n_k)=Fbb1(:,n_k)/norm(F1(:,1:n_subK)*Fbb1(:,n_k));
       Fbb1less(:,n_k)=Fbb1less(:,n_k)/norm(F1less(:,1:n_subK)*Fbb1less(:,n_k));
       Fbb2less(:,n_k)=Fbb2less(:,n_k)/norm(F2less(:,1:n_subK)*Fbb2less(:,n_k));
       Fbb1less2(:,n_k)=Fbb1less2(:,n_k)/norm(F1less2(:,1:n_subK)*Fbb1less2(:,n_k));
       Fbb2less2(:,n_k)=Fbb2less2(:,n_k)/norm(F2less2(:,1:n_subK)*Fbb2less2(:,n_k));
       Fbb2ba(:,n_k)=Fbb2ba(:,n_k)/norm(F2ba(:,1:n_subK)*Fbb2ba(:,n_k));
   end
 
   
 
   
   snr0=abs(diag(Heq0*Fbb0)).^2./noise0(1:n_subK);
   snr1=abs(diag(Heq1*Fbb1)).^2./noise1(1:n_subK);
   snr1less=abs(diag(Heq1less*Fbb1less)).^2./noise1less(1:n_subK);
   snr1less2=abs(diag(Heq1less2*Fbb1less2)).^2./noise1less2(1:n_subK);
   snr2=abs(diag(Heq2*Fbb2)).^2./noise2(1:n_subK);
   snr2ba=abs(diag(Heq2ba*Fbb2ba)).^2./noise2ba(1:n_subK);  
   snr2less=abs(diag(Heq2less*Fbb2less)).^2./noise2less(1:n_subK);
   snr2less2=abs(diag(Heq2less2*Fbb2less2)).^2./noise2less2(1:n_subK);

   
   r0K(nn)=r0K(nn)+sum(log2(1+snr0*ptK/ga),1);
   r1K(nn)=r1K(nn)+sum(log2(1+snr1*ptK/ga),1);
   r1lessK(nn)=r1lessK(nn)+sum(log2(1+snr1less*ptK/ga),1);
   r2lessK(nn)=r2lessK(nn)+sum(log2(1+snr2less*ptK/ga),1);
   r1less2K(nn)=r1less2K(nn)+sum(log2(1+snr1less2*ptK/ga),1);
   r2less2K(nn)=r2less2K(nn)+sum(log2(1+snr2less2*ptK/ga),1);
   r2K(nn)=r2K(nn)+sum(log2(1+snr2*ptK/ga),1);
   r2baK(nn)=r2baK(nn)+sum(log2(1+snr2ba*ptK/ga),1);


   end

  
   

    r0=r0+sum(log2(1+snr0*pt/ga),1);
    r1=r1+sum(log2(1+snr1*pt/ga),1);
    r1less=r1less+sum(log2(1+snr1less*pt/ga),1);
    r2less=r2less+sum(log2(1+snr2less*pt/ga),1);
    r1less2=r1less2+sum(log2(1+snr1less2*pt/ga),1);
    r2less2=r2less2+sum(log2(1+snr2less2*pt/ga),1);
    r2=r2+sum(log2(1+snr2*pt/ga),1);
    r2ba=r2ba+sum(log2(1+snr2ba*pt/ga),1);


end

end


r0=r0/itr;
r1=r1/itr;
r2=r2/itr;
r2ba=r2ba/itr;
r1less=r1less/itr;
r1less2=r1less2/itr;
r2less=r2less/itr;
r2less2=r2less2/itr;

r0K=r0K/itr;
r1K=r1K/itr;
r2K=r2K/itr;
r2baK=r2baK/itr;
r1lessK=r1lessK/itr;
r1less2K=r1less2K/itr;
r2lessK=r2lessK/itr;
r2less2K=r2less2K/itr;

r0K=r0K./subK;
r1K=r1K./subK;
r2K=r2K./subK;
r2baK=r2baK./subK;
r1lessK=r1lessK./subK;
r1less2K=r1less2K./subK;
r2lessK=r2lessK./subK;
r2less2K=r2less2K./subK;


%% average rate/user vs SNR
r0=r0/maxK;
r1=r1/maxK;

r2=r2/maxK;
r2ba=r2ba/maxK;
 r1less=r1less/maxK;
 r1less2=r1less2/maxK;
 r2less=r2less/maxK;
 r2less2=r2less2/maxK;
figure

plot(snr_range, r0,'k--');hold
plot(snr_range,r2ba,'r-o');
plot(snr_range, r1,'b--+');
plot(snr_range,r2less,'r--s');
plot(snr_range, r1less,'b:*');

legend('unconstrained','S2, b=2','S1, b=2','S2, b=1.5','S1, b=1.5');
ylabel('Average rate');xlabel('P_t/\sigma_n^2');


grid on
%% average rate/user vs number of users

figure

plot(subK, r0K,'k--');hold
plot(subK,r2baK,'r-o');


plot(subK, r1K,'b:+');
plot(subK,r2lessK,'r-s');
plot(subK, r1lessK,'b:*');
plot(subK,r2less2K,'r->');
plot(subK, r1less2K,'b:<');
legend('unconstrained','S2, b=3','S1, b=3','S2, b=2','S1, b=2','S2, b=1.5','S1, b=1.5');

ylabel('Average rate');xlabel('number of users K');
set(gca,'xTick',subK)
grid on

