%% BAZZI Downlink training sequence %%

clearvars; rng('shuffle');

% SNR
SNRdB=10; 
SNR = 10^(SNRdB/10);

% Number of Users
K = 10;

% Number of BS antenna
M=20;

%Angular Spread

AS=10;

% Wavelength 
Lambda=0.5; 

% Choose a different AoAs
AoA_all=[-57.5,57.5, -23, 23.5, -7.5,7.5, -41.5, 41.5,-45,45]; 

Ck_OR=zeros(M,M,K);

  for k = 1:K
    AoA=AoA_all(k); 
    C_k=funOR(M,AS,Lambda,AoA) ;  
    Ck_OR(:,:,k)=C_k;
  end
  
% Pilot length     
  T=10; 

 
%% Generate a Random Gaussian   
A=[randn(M,M)+1i*randn(M,M)]/sqrt(2); 

% Select a Unitary 
[U_p,~,~]=svd(A); % Take the SVD 

% Initial Unitary
%S_t1=U_p(:,1:T);

% Initial DFT
 S_dft=sqrt(1/M)*dftmtx(M);
 S_t1=S_dft(:,1:T);

second=1;
first=1;

tr_Z=1000;

threshold=10e-5;

kkk=1;

while tr_Z<=threshold || kkk<1000
    
Sopt=S_t1;
    

% alph_pl=1e-1;
% w_k=1./alph_pl;

alph_pl=1;
w_k=1;
% Initial gamma
gamma=1;

Div_MSE=zeros(M,T,K);
    for kk=1:K
       Ck=alph_pl*Ck_OR(:,:,kk); 
       Div_MSE(:,:,kk)=(-1*w_k*(Ck*Ck*S_t1/(S_t1'*Ck*S_t1+(1/SNR)*eye(T)))+(Ck*S_t1)*(inv(S_t1'*Ck*S_t1+(1/SNR)*eye(T)))*S_t1'*w_k*Ck*Ck*S_t1/(S_t1'*Ck*S_t1+(1/SNR)*eye(T))) ;  
    end
 
    Div_MSE=sum(Div_MSE,3);

    Z=-1*(eye(M)-S_t1*S_t1')*2*Div_MSE;

    tr_Z=trace(Z'*Z);
   
   aaaa(kkk)=tr_Z;
    

   WSMSE_S_t1=zeros(T,T);
   for kk=1:K
       Ck=alph_pl*Ck_OR(:,:,kk); 
       WSMSE_S_t1=WSMSE_S_t1+(S_t1'*Ck*Ck*w_k*S_t1/(S_t1'*Ck*S_t1+(1/SNR)*eye(T)));
   end
   tr_WSMSE_S_t1=real(trace(WSMSE_S_t1));

   WSMSE_S=trace(sum(Ck_OR,3))-tr_WSMSE_S_t1;


% Consider new S
S_t2=(S_t1+2*gamma*Z);

WSMSE_S_t2=zeros(T,T);

   for kk=1:K
       Ck=alph_pl*Ck_OR(:,:,kk);     
       WSMSE_S_t2=WSMSE_S_t2+(S_t2'*Ck*Ck*w_k*S_t2/(S_t2'*Ck*S_t2+(1/SNR)*eye(T)));
    end

tr_WSMSE_S_t2=real(trace((WSMSE_S_t2)));

% Calculate WSMSE(S+2gammaZ)
WSMSE_S_new=trace(sum(Ck_OR,3))-tr_WSMSE_S_t2;

% Calculate the difference between the old WSMSE and new WSMSE
Diff_new= WSMSE_S-WSMSE_S_new;

%% first Direction 

if Diff_new >= gamma*tr_Z

    first=first+1;
    
while Diff_new<=10e-4
    
gamma=gamma*2;

S_t2=(S_t1+2*gamma*Z);

WSMSE_S_t2=zeros(T,T);
    for kk=1:K
       Ck=alph_pl*Ck_OR(:,:,kk); 
       WSMSE_S_t2=WSMSE_S_t2+(S_t2'*Ck*Ck*w_k*S_t2/(S_t2'*Ck*S_t2+(1/SNR)*eye(T)));
    end

tr_WSMSE_S_t2=trace((WSMSE_S_t2));

WSMSE_S_new=trace(sum(Ck_OR,3))-tr_WSMSE_S_t2;

Diff_new= WSMSE_S-WSMSE_S_new;

end


end

%% Second Direction 

if Diff_new< gamma*tr_Z
    
    second=second+1;

 
while Diff_new<=10e-4
    
gamma=gamma/2;

S_t2=(S_t1+gamma*Z);

WSMSE_S_t2=zeros(T,T);
    for kk=1:K
       Ck=alph_pl*Ck_OR(:,:,kk); 
       WSMSE_S_t2=WSMSE_S_t2+(S_t2'*Ck*Ck*w_k*S_t2/(S_t2'*Ck*S_t2+(1/SNR)*eye(T)));
    end

tr_WSMSE_S_t2=real(trace((WSMSE_S_t2)));

WSMSE_S_new=trace(sum(Ck_OR,3))-tr_WSMSE_S_t2;

Diff_new= WSMSE_S-WSMSE_S_new;

end

end

S_t1=S_t1+gamma*Z;

[Q,~]=qr(S_t1);
S_t1=Q(:,1:T);

kkk=kkk+1;
end

  
  Stop
  
  
  
  
  
  

