% function function [b, E] = iw(N, U, Hmag, No, Ex, gap, mode, b_target,cb)
%
% Calculates data rates of M users and corresponding bit distributions and Energy distributions
%    using iterative waterfilling.
%
% Inputs
% ------
% N: number of sub-channels
% U: number of users
% Hmag: squared channel transfer and crosstalk matrix (N x U x U matrix)
%        Hmag(n,i,j) is the crosstalk transfer function from loop i to j at the nth bin.
% No: noise power spectrum per tone (N x U)
% Ex: signal energy/SYMBOL
% gap: gap in dB
% mode: (U x 1 vector) each value is one of the followings
%       0 - rate adaptive
%       1 - fixed margin (power minimization)
%       2 - margin adaptive
% b_target: target bits on 1 DMT symbol for modes 1 and 2
% bitcap: maximum possible number of bits at each frequency bin
% cb =1 for cplx BB and =2 for real BB
%
% Outputs
% -------
% b: bit distribution (N x U matrix)
% E: energy distribution (N x U matrix)
%
% Remarks
% Iterate waterfiling for each user 10 times
% Youngjae(Sean) Kim - modified J. Cioffi, April 2024

function [b, E] = iw(N, U, Hmag, No, Ex, gap, mode, b_target,cb)


E = zeros(N,U); % mJ
b = zeros(N,U);
gaplin=10^(gap/10);

for count = 1:10
    for m = 1:U
        % Waterfilling for user m
        H = Hmag(:,m,m);% need consistency on real/complex H and No.
        No_all = No(:,m); %mW/Hz
        for n=1:U
            if n~=m
                No_all = No_all + Hmag(:,n,m).*E(:,n); % mW/Hz = mJ
            end
        end
        gn=H./No_all;
        switch mode(m)
        case 0
            [b(:,m), E(:,m)] = waterfill_gn(gn', Ex(m)/N, gap, cb);
        case 1
            [b(:,m), E(:,m)] = fmwaterfill_gn(gn', b_target(m)/N, gap, cb);            
        case 2
            [b(:,m), E(:,m)] = waterfill_em(gn', b_target(m)/N, gap, cb);
        end
    end
    % Remove the comment here, to see the bit distribution profile evolve
    %for m = 1:M
    %    stairs([1:N]*df, b(:,m)); hold on;
    %end
    %pause;
    %hold off;
end
for m =1:U
    No_all = No(:,m);
    for u=1:U
        if u~=m
            No_all = No_all + Hmag(:,u,m).*E(:,u) ; % mW/Hz = mJ
        end
     end
    b(:,m) = (1/cb)*(log2(1+1/gaplin*Hmag(:,m,m).*E(:,m)./No_all));
end
return