function [W, U, p] = waterFilling(H, P, sigma2)
% W: precoder; U: combiner. 
    [U,Lambda,W] = svd(H); 
    
    [~, Nt] = size(H); 
    lambda2 = diag(Lambda).^2;
    M = max(lambda2); 
    eps = 1e-7 * M;      % do waterfilling only with quality >= eps. 
    
    % sel = (lambda2 >= eps); 
    sel = true(Nt, 1); 
    lambda2(lambda2 <= eps) = eps; 

    sel_lambda2 = lambda2(sel); 

    % disp(sel_lambda2.');
    Nsel = sum(sel); 

    [wminLevel, order] = sort(sigma2./sel_lambda2, 'ascend'); 

    accumulativeMaxPower = zeros(Nsel, 1);

    for idx = 2:Nsel
        accumulativeMaxPower(idx) = accumulativeMaxPower(idx-1) + (idx-1) * (wminLevel(idx)-wminLevel(idx-1)); 
    end

    t = 1;
    while t <= Nsel && accumulativeMaxPower(t) <= P
        t = t+1;
    end

    mu = (P-accumulativeMaxPower(t-1))/(t-1) + wminLevel(t-1); 
    q = mu - wminLevel; 
    q(q<0) = 0; 
    
    p = zeros(Nt, 1); 
    indices = 1:Nt; 
    indices = indices(sel); 

    for ii = 1:Nsel
        p(indices(order(ii))) = q(ii); 
    end
    
end
