%
% Waterfilling algorithm (from [Palomar and Fonollosa, Trans-SP2004]) to compute:
%
%           pi = (mu*ai - bi)^+
%           sum(pi) = Pt
%
%  By Daniel Perez Palomar (last revision: May 10, 2004).
%  Feel free to distribute this file as it is (without including any modifications).
%  Please, email any improvement or error to Daniel.P.Palomar@ieee.org or danielp@princeton.edu
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
% Details of the function:
% ------------------------
%
% It uses Algorithm 2 of: Palomar and Fonollosa, "Practical Algorithms for a Family of Waterfilling Solutions", 
% IEEE Trans. Signal Proc., 2004.
%
% INPUTS: Vectors ai and bi of same size.
%
% OUTPUT: Power allocation pi and waterlevel mu.
%
function [pi, mu] = WaterFilling(ai,bi,Pt)

L=length(ai);

%Reorder substreams with decreasing ai./bi or increasing bi./ai
[tmp,index]=sort(bi./ai);
ai=ai(index);  %ai_original(index)=ai;
bi=bi(index);  %bi_original(index)=bi;

%Loop to obtain the optimal waterlevel
ai(L+1)=1; bi(L+1)=+Inf; %define a(L+1)/b(L+1)=0 or b(L+1)/a(L+1)=+Inf
L_=L;
while L_>=1,
   mu=bi(L_)/ai(L_);
   if mu < bi(L_+1)/ai(L_+1)  &&  mu < ( Pt + sum(bi(1:L_)) )/( sum(ai(1:L_)) ), 
      break; %accept hypothesis
   else
      L_=L_-1; %reject hypothesis and form a new one
   end
end

%Compute the definite waterlevel and the power allocation
mu = ( Pt + sum(bi(1:L_)) ) / sum(ai(1:L_));
pi = max(0, mu*ai(1:L) - bi(1:L));
      
%Recover the original ordering
pi(index)=pi;