% function [Rxx, Rwcn, bmax] = bcmax(iRxx, H, Lyu)
%
%  Uses cvx_wcnoise.m and rate-adaptive waterfill.m (Lagrange
%  Multiplier based)
%   Arguments:
%       - iRxx: initial input autocorrelation array, size is Lx x Lx x N. 
%           Only the sum of traces matters, so can initialize to any valid 
%           autocorrelation matrix Rxx to run wcnoise.
%           needs to include factor N/(N+nu) if nu ~= 0
%       - H: channel response, size is Ly x Lx x N, w/o sqrt(N)
%       normalization
%       - Lyu: number of antennas at each user
%             can create variable-u by just using dummy zero rows in H for
%             all output receivers that have less than max (=Lyu input)
%   Outputs:
%       - Rxx: optimized input autocorrelation 
%       - Rwcn: optimized worst-case noise autocorrelation, with white local noise
%               SO IF H is noise-whitened for Rnn, then actual noise is
%               Rwcn^(1/2)*Rnn*Rwcn^(*/2)
%       - b: maximum sum rate/real-dimension - user must mult by 2 for
%            complex case

function [Rxx, Rwcn, bmax] = bcmax(iRxx, H, Lyu)
    
    [Ly,Lx,N] = size(H);
    total_en = trace(sum(iRxx, 3));
    bmax = 0;
    ib=zeros(1,N);
    Rwcn = zeros(Ly,Ly,N);
    for n=1:N % worst-case noise independent over tones for BC
    [Rwcn(:,:,n), ib(n)] = cvx_wcnoise(iRxx(:,:,n), H(:,:,n), Lyu);
    end
    
    while (abs(sum(ib) - bmax) > 1e-5) %tolerance
        % uncomment the following two lines to see how the loop progresses
        %bmax
        %bmin = sum(ib)
     % Vector Coding Gains for each tone
        M=zeros(Lx,Lx,N);
        gains = zeros(Lx,N);
        for n=1:N 
            [V,D,~]=svd(Rwcn(:,:,n));
            sqD = sqrt(D).*(D>1e-6);
            invsqRwcn = V*pinv(sqD)*V';
            Htil = invsqRwcn*H(:,:,n);
            [~, g, M(:,:,n)] = svd(Htil);
            g=diag(g);
            gains(1:length(g),n)=g.^2;
        end
     % Water-filling step
        En = waterfill(total_en, reshape(gains',N*Lx,1), 1);
        bvec=0.5*log2(ones(N*Lx,1)+En.*reshape(gains',N*Lx,1))'; 
        bmax = real(sum(bvec));
        En=reshape(En,N,Lx)';
     % update worst-case-noise step
        for n=1:N
            newiRxx = M(:,:,n)*diag(En(:,n))*M(:,:,n)';
            iRxx(:,:,n) = (iRxx(:,:,n)+newiRxx)/2;
            %iRxx(:,:,n) = M(:,:,n)*diag(En(:,n))*M(:,:,n)';
           [Rwcn(:,:,n), ib(:,n)] = cvx_wcnoise(iRxx(:,:,n), H(:,:,n), Lyu);
        ib(:,n)=real(ib(:,n));
        end
    end
    Rxx = iRxx;
end