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

% EE379C 2001-2002 Spring,
% Written by Seong Taek Chung
%
% P is the pulse response
% SNRmfb is the SNRmfb in dB
% b_bar is the normalized bit rate
% 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

% 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=(2^(Ntot*b_bar*2)/prod(gn_sorted(1:Nstar)))^(1/Nstar)*gap;
    %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

en_bar=en;
bn_bar=bn;

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

%calculate margin
margin=10*log10(Ntot*Ex_bar/sum(en));

%>> [gn,en_bar,bn_bar,Nstar,b_bar_check,margin]=MAwaterfill([1 0.9],10,1,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 =
%
%  Columns 1 through 7 
%
%    0.2000    0.4369    0.4782    0.4868    0.4782    0.4369    0.2000
%
%  Column 8 
%
%         0
%
%
%bn_bar =
%
%  Columns 1 through 7 
%
%    0.3361    1.2123    1.5964    1.7103    1.5964    1.2123    0.3361
%
%  Column 8 
%
%         0
%
% Nstar =
%
%     7
%
%
%b_bar_check =
%
%   1.0000
%
%
%margin =
%
%    4.6903
%