function [gn,en_bar,bn_bar,Nstar,b_bar]=waterfill(P,SNRmfb,Ex_bar,Ntot,gap)

% EE379C 2001-2002 Spring,
% EE379C 1999-2000 Spring, 
% EE379C 1997-1998 Spring, 
% Written by Atul Salvekar, Edited by Wonjong Rhee
%
% function [gn,en_bar,bn_bar,Nstar,b_bar]=waterfill(P,SNRmfb,Ex_bar,N,gap)
%
% P is the pulse response
% SNRmfb is the SNRmfb in dB
% Ex_bar is the normalized energy
% Ntot is the total number of real/complex subchannels, Ntot>2
% gap is the gap in dB
%
% gn is channel gain
% en_bar is the energy/dim in the nth subchannel
% bn_bar is the bit/dim in the nth subchannel
% Nstar is the number of subchannel used
% b_bar is the bit rate

% dB into normal scale
Noise_var=Ex_bar*(norm(P)^2)/(10^(SNRmfb/10));
gap=10^(gap/10);

% initialization
en=zeros(1,Ntot);
bn=zeros(1,Ntot);
gn=zeros(1,Ntot);
Hn = zeros(1,Ntot);

% subchannel center frequencies
f=-1/2+1/Ntot:1/Ntot:1/2;

% find Hn vector
for i=1:length(P)
	Hn=Hn+P(i)*exp(j*2*pi*f*(i-1)); 
        % This value will be different depending if P represents 
        % P(1) + P(2)*D^-1 + ....  or P(1) + P(2)*D^+1....,
        % but we'll get same gn, thus same waterfilling result.
        % (Note that both have the same magnitude response!)
end

% find gn vector
gn=abs(Hn).^2/Noise_var;
%plot(gn)

%%%%%%%%%%%%%%%%%%%%%%%
% Now do waterfilling %
%%%%%%%%%%%%%%%%%%%%%%%

%sort
[gn_sorted, Index]=sort(gn);  % sort gain, and get Index

gn_sorted = fliplr(gn_sorted);% flip left/right to get the largest 
                              % gain in leftside
Index = fliplr(Index);        % also flip index  


num_zero_gn = length(find(gn_sorted == 0)); %number of zero gain subchannels
Nstar=Ntot - num_zero_gn;    
 	% Number of used channels, 
 	% start from Ntot - (number of zero gain subchannels)

while(1) 
 	K=1/Nstar*(Ntot*Ex_bar+gap*sum(1./gn_sorted(1:Nstar))); 
 	En_min=K-gap/gn_sorted(Nstar);	% En_min occurs in the worst channel
 	if (En_min<0)		
    		Nstar=Nstar-1;  % If negative En, continue with less channels
 	else 
    		break;       % If all En positive, done.
 	end
end

En=K-gap./gn_sorted(1:Nstar); 		% Calculate En
Bn=.5*log2(K*gn_sorted(1:Nstar)/gap); 	% Calculate bn

bn(Index(1:Nstar))=Bn;		% return values in original index
en(Index(1:Nstar))=En;		% return values in original index

middle = Ntot/2;		% Since channel is even, need to display 
                                % only half of result
en_bar=en(middle:Ntot);
bn_bar=bn(middle:Ntot);

% calculate b_bar
b_bar=1/Ntot*(sum(bn));



% check if you get the following.
%>> [gn,en_bar,bn_bar,Nstar,b_bar] = waterfill([1 .9], 10,1,8,0)
%gn =
%  Columns 1 through 7 
%    2.9680   10.0000   17.0320   19.9448   17.0320   10.0000    2.9680
%  Column 8 
%    0.0552

%en_bar =
%    1.2415    1.2329    1.1916    0.9547         0

%bn_bar =
%    2.3436    2.2297    1.8456    0.9693         0

%Nstar =
%     7

%b_bar =
%    1.5541