% function [bn, en, Nstar] = waterfill_em(gn, bsum_bar, gap, cb)
%
% This function minimizes total energy for dimensions in gn.  
% THIS FUNCTION ACCEPTS bit sum divied by number of dimensions.
%
% Basically, issues of guard period are not addressed because the
% subchannel gains are provided as an input
%
% INPUT
% gn is the channel gain (a row vector).
% bsum_bar is the target sum rate in bits/real-dimension
% gap is the gap in dB
% cb = 1 for complex bb and cb=2 for real bb (if not given, assumes cb=2)
%
% OUTPUT
% en is the energy in the nth subchannel; for complex baseband divide by 2
%    to obtain energy/real-dimension
% bn is the bit in the nth subchannel; for complex baseband multiply by 2
%    to obtain bits/real dimension
% Nstar is the number of energized dimensions (complex dim if cb =1)
%
% dB into normal scale
function [bn , en , Nstar] = waterfill_em(gn, bsum_bar, gap , cb)
gap = 10^(gap/10);
if nargin < 4
    cb=2;
end

[col, Ntot] = size(gn);

if col ~= 1
    error = 1
    return;
end

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

%%%%%%%%%%%%%%%%%%%%%%%
% 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));
%----------
Nstar= Ntot - num_zero_gn;   
 	% Number of used channels, 
 	% start from N - (number of zero gain subchannels)
b=(3-cb)*Ntot*bsum_bar;
Klog2tilde=cb*b-sum(log2(gn_sorted(1:Nstar)));

while(1) 
 	Klog2=(1/Nstar)*Klog2tilde + log2(gap);
 	En_min=2^(Klog2)-gap/gn_sorted(Nstar);	% En_min occurs in the worst channel
 	if (En_min<0)
            Klog2tilde=Klog2tilde+log2(gn_sorted(Nstar));
    		Nstar=Nstar-1;  % If negative En, continue with less channels
 	else 
    		break;       % If all En positive, done.
 	end
end
Kma=2^Klog2;

En = Kma - gap./gn_sorted(1:Nstar); % Calculate En
Bn =(1/cb) * log2(2^Klog2 * 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