% [b,E,K,g]=dmt(Hds,Nds,E_total,Gap,VDSLmask)
% Computes bit distribution, and energy distribution for a dmt system.
% Also returns sub-channel SNR's and variable K.   
% The channel magnitude is specified in Hds (one-sided,dBm), where
% upstream and downstream bandwidth allocation must have been taken into 
% account. The noise psd is specified in Nds (in dBm).  The total energy 
% E_total, the effective Gap, and the psd mask (one-sided, in dBm) 
% also have to be specified. 
% The sub-channel center frequencies will be at: (2k-1)df/2 .  

%function [b,E,K,g]=dmt(Hds,Nds,E_total,Gap,PSDmask)

function [b, E, g, Marg] =dmt(Hds,Nds,E_total,Gap,PSDmask,const);

b_min=min(const);
b_max=max(const);

%g=(Hds.^2 ./ Nds)'; 		% sub-channel normalized SNR's
g=(Hds.^2 ./ Nds);

[g_sorted,index] = sort(g);

g_sorted=fliplr(g_sorted);      % flip sorted gains
index=fliplr(index);

Ntotal=length(Hds);      % total number of 1-dimension sub-channels

b=zeros(1,Ntotal);
E=zeros(1,Ntotal);

num_zero_g=length(find(g_sorted<10^-8));     
Nstar=Ntotal-num_zero_g;   % start from worst channel (with non-zero gain)

% water-filling
while(1)
    K = 1/Nstar*(E_total+Gap*sum(1./g_sorted(1:Nstar)));
    E_min=K-Gap/g_sorted(Nstar);
    if (E_min<0)
        Nstar=Nstar-1;
    else
        break;
    end;
end;

Es=K-Gap./g_sorted(1:Nstar);         % energies of subchannels 

count = 0;
E_extra = 0;
for k=1:Nstar
    if (Es(k) < PSDmask(index(k)))
      count = count + 1;
    else
      E_extra = E_extra + Es(k) - PSDmask(index(k));
    end;
end;
if (count == 0)		% count==0 means E_extra=0. avoid 0/0
    count = 1;
end;

for k=1:Nstar		% redistribute excess energy
    if (Es(k)+E_extra/count < PSDmask(index(k)))
      E(index(k)) = Es(k) + E_extra/count;
    else
      E(index(k)) = PSDmask(index(k));
    end;
end;

b = log2(1+E.*g/Gap);             % bit rates of subchannels

b=round(b);       
b(find(b>b_max))=b_max*ones(size(find(b>b_max)));  
b(find(b<b_min))=zeros(size(find(b<b_min)));

E=(2.^b-1).*Gap./g;

tmp=find(b>0);
Marg=zeros(size(b));

Marg(tmp) = E(tmp).*g(tmp)./(2.^b(tmp)-1)./Gap;