% Author: Zukang Shen and Ian C. Wong  
% Copyright (C) 2004   Ian C. Wong
% 
% This program is free software; you can redistribute it and/or
% modify it under the terms of the GNU General Public License
% as published by the Free Software Foundation; either version 2
% of the License, or (at your option) any later version.
% 
% This program is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.
% 
% You should have received a copy of the GNU General Public License
% along with this program; if not, write to the Free Software
% Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
% 
% You may reach the author at wongic@mail.utexas.edu.
% Or visit his website at www.ece.utexas.edu/~iwong

function [capacity] = waterfilling(Ptot,eigenVec)
eigenVec=eigenVec+1e-10;
eigenVec=eigenVec(:);
sortedEigVec=sort(eigenVec);

EigVecTemp=sortedEigVec;


single = 0;
while sum((EigVecTemp-EigVecTemp(1)*ones(length(EigVecTemp),1))./(EigVecTemp*EigVecTemp(1)))>Ptot
    EigVecTemp=EigVecTemp(2:length(EigVecTemp));
    if length(EigVecTemp) < 2,
        single = 1;
        break;
    end;        
end;
if single == 1,
    capacity = log2(Ptot*EigVecTemp);
else,
	P1=(Ptot-sum((EigVecTemp-EigVecTemp(1)*ones(length(EigVecTemp),1))./(EigVecTemp*EigVecTemp(1))))/length(EigVecTemp);
	mu=P1+1/EigVecTemp(1);
	capacity=sum(log2(mu*EigVecTemp));
end;