clear all; close all;clc;
Ex_bar=1;
Noise_var=0.181;
Ntot=8;
tail=rand(1,100)-1;
H_D=[1 tail]/norm([1 tail])*sqrt(1^2+0.9^2);
% H_D=[0.9 1];
[main_tap,index_main]=max(H_D);

%%% capacity calculation %%%
% set the gap to 0 dB
gap_dB=0;
[c_bar,cn_bar,en_bar,Nstar,gn]=waterfill_soln(Ex_bar,Noise_var,Ntot,gap_dB,H_D,index_main)

Pe=1e-4;
%%% QAM calculation %%%
% set the gap
gap_dB_qam=10*log10(1/3*qfuncinv(Pe)^2);
[b_bar_qam,bn_bar_qam,en_bar_qam,Nstar_qam,gn_qam]=waterfill_soln(Ex_bar,Noise_var,Ntot,gap_dB_qam,H_D,index_main)

%%% PSK calculation %%%
% set the gap
gap_dB_psk=10*log10(1/pi^2*qfuncinv(Pe)^2);
[b_bar_psk,bn_bar_psk,en_bar_psk,Nstar_psk,gn_psk]=waterfill_soln(Ex_bar,Noise_var,Ntot,gap_dB_psk,H_D,index_main)
b_bar_psk=1/2*b_bar_psk % take into account the additional squaring factor in PSK gap derivation
bn_bar_psk=1/2*bn_bar_psk

Ntot=2:2:200;
for k=1:length(Ntot),
[c_bar(k),c(k).cn_bar,en_bar,Nstar,gn]=waterfill_soln(Ex_bar,Noise_var,Ntot(k),gap_dB,H_D,index_main);
[b_bar_qam(k),b(k).bn_bar_qam,en_bar_qam,Nstar_qam,gn_qam]=waterfill_soln(Ex_bar,Noise_var,Ntot(k),gap_dB_qam,H_D,index_main);
[b_bar_psk(k),bn_bar_psk,en_bar_psk,Nstar_psk,gn_psk]=waterfill_soln(Ex_bar,Noise_var,Ntot(k),gap_dB_psk,H_D,index_main);
b_bar_psk(k)=b_bar_psk(k)/2;
end    
figure(1)
hold off;
plot(Ntot,c_bar,'r.-', Ntot,b_bar_qam,'bo-',Ntot,b_bar_psk,'gv-');
legend('capacity','QAM','PSK');