function [Ri, Di] = reverse_water_filling( variances, Rtot, mode)
%REVERSE_WATER_FILLING   ... determines the MSE-optimal bit allocation
%
%[Ri, Di] = reverse_water_filling(variances, Rtot); % or:
%[Ri, Di] = reverse_water_filling(variances, Rtot, 'rate');
% given sum rate Rtot
%
%[Ri, Di] = reverse_water_filling(variances, Dtot, 'dist');
% given sum distortion Dtot
%
% (c) by Christian Feldbauer, 2009-2010

if nargin==2
 mode = 'rate'; % default
elseif nargin == 3
 if strcmp( lower(mode), 'dist')
   mode = 'dist';
   Dtot = Rtot;
 elseif strcmp( lower(mode), 'rate')
   mode = 'rate';
 else
   warning('invalid mode! assume "rate" given.')
   mode = 'rate';
 end
else
 error('at least 2, at most 3 arguments needed.') 
end


variances = variances(:)';

[variances, sortind] = sort(variances, 'descend');
M = length(variances); % total number of variables


for kk = M:-1:1 % number of variables to be encoded
 switch mode
 case 'rate'
  d_kk = geomean(variances(1:kk)) * 2^(-2*Rtot/kk); % dist. per encoded variable
 case 'dist'
  d_kk = (Dtot-sum(variances((kk+1):end)))/kk;
 end 
 % note, d_kk is the Lagrange multiplier...
 if sum( d_kk > variances(1:kk) ) == 0
  break;
 end
end
% kk and d_kk has been found, the rest is simple:
Di = [repmat(d_kk,1,kk) variances((kk+1):end)];
Di(sortind) = Di;
Ri = max( 0, 1/2*log2(variances ./ d_kk));
Ri(sortind) = Ri;

