% function [gn,en_bar,bn_bar,Nstar,b_bar,SNRdmt]=DMTra(h,NoisePSD,Ex_bar,N,gap)
%
% CAUTION - the user must know if the input channel h is complex baseband
% or real baseband in interpretting the outputs, as per comments below
%
% INPUTS
% h - the complex or real baseband sampled (temporal) pulse response
% NoisePSD and Ex_bar - noise power and energy/dimension or cpx-sample. 
%    These two parameters need to be consistent in terms of # dimensions.    
% N is the DFT size, N>2, and is "Nbar" for complex channels and N
%     for real channels
% The sampling rate (on each of Inphase or Quad for complex) is (N+v)/T 
%    where 1/T is DMT symbol rate and nu = length(p) - 1
% gap is the gap in dB
%
% OUTPUTS
% gn is vector of channel gains or magnitude-squared |P|^2 / NoisePSD
% En - vector of energy allocation. 
%      Per real dimension for real channels; 
%      Per tone for complex channel.
%    For real channels, the upper image frequencies duplicate the lower 
%      frequencies.
%    For complex channels, each En output is per that tone (= cmplx dim)
% bn_bar is the vector of bit/dim for all subchannels 
%    For complex channels, double bn_bar for bits/tone. 
%    For real channels, the upper image frequencies duplicate lower 
%      frequencies, but bn_bar remains bits per real dimension.
% Nstar is the number of used input-size DFT (real or complex) dimensions 
%     (2*Nstar is number of real dimensions for complex chan)
% b_bar is the number of bits per (real) dimension 
%     (so 2*(N+v)*b_bar is total number of bits/symbol for a complex chan )
% SNRdmt is the equivalent DMT "geometric" SNR (complex or real) in dB
%
function [gn,en_bar,bn_bar,Nstar,b_bar,SNRdmt]=DMTra(h,NoisePSD,Ex_bar,N,gap)

Noise_var=NoisePSD;
gap=10^(gap/10);
nu=length(h)-1;

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

% subchannel center frequencies
f=0:1/N:1-1/N;

% find Hn vector
for i=1:length(h)
	Hn=Hn+h(i)*exp(j*2*pi*f*(i-1)); 
end

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

%%%%%%%%%%%%%%%%%%%%%%%
% Waterfilling for DMT %
%%%%%%%%%%%%%%%%%%%%%%%

%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=N - num_zero_gn;    
 	% Number of used channels, 
 	% start from N - (number of zero gain subchannels)

while(1) 
 	K=1/Nstar*(N*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=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

		% Since channel is even, need to display 
                                % only half of result
en_bar=en(1:N);
bn_bar=0.5*bn(1:N);

% calculate b_bar
b_bar=1/(N+nu)*(0.5*sum(bn));
SNRdmt=10*log10(gap*(2^(2*b_bar)-1));





