clear all; close all; clc;

%%% System parameters %%%
P=1e-6;
fp=1e6;
cp_coeff=0.99; %e.g. 99% energy in pulse response taps determines the CP length
Pe=1e-6;
NF_dB=20;

Tsamp_vec=10.^-(3:0.25:10);
N_vec=2.^(4:12);
for n=1:length(Tsamp_vec),
    for k=1:length(N_vec),
        Tsamp=Tsamp_vec(n);
        N=N_vec(k);
        
        %%% Get channel data %%%
        Pnoise=1e-3*10^((-174+10*log10(1/Tsamp)+NF_dB)/10);
        [Hn,f,Ncp]=channel_model(Tsamp,N,fp,cp_coeff);

        %%% Energy calculations %%%
        Ex_tot=N*P*Tsamp; % total signal energy
        Ex_bar=Ex_tot/N; % average energy per dimension
        E_noise=Pnoise*Tsamp; % noise energy per dimension

        %%% Waterfilling %%%
        Tsym_cap=N*Tsamp;
        gap_dB=0;
        [c_bar(n,k),cn_bar,en_bar,Nstar(n,k),gn]=waterfill(Ex_bar,E_noise,gap_dB,Hn);
        C(n,k)=sum(cn_bar)/Tsym_cap; % i.e. C=c_bar/Tsamp;

        gap_dB_qam=10*log10(1/3*qfuncinv(Pe)^2);
        [b_bar(n,k),bn_bar,en_bar,Nstar_qam(n,k),gn]=waterfill(Ex_bar,E_noise,gap_dB_qam,Hn);
        R_qam(n,k)=b_bar(n,k)/Tsamp; % i.e. C=c_bar/Tsamp;
        
        %%% Levin-Campello %%%
        Hn_ss=Hn(length(Hn)/2:end); % single-sided channel (positive frequencies only)
        [b_bar_lc(n,k),bn,En,gn]=LC_RA(Ex_bar,E_noise,gap_dB_qam,Hn_ss);
        R_qam_lc(n,k)=N*b_bar_lc(n,k)/((N+Ncp)*Tsamp); % i.e. C=c_bar/Tsamp;
        if N<Ncp, R_qam_lc(n,k)=0; end
    end
end

figure(1)
hold off;
plot(1e-9./Tsamp_vec, 1e-6*C(:,end)','b.-');
hold on;
plot(1e-9./Tsamp_vec, 1e-6*R_qam(:,end)','g.-');
plot(1e-9./Tsamp_vec, 1e-6*max(R_qam_lc,[],2)','r.-');
xlabel('sampling rate [GHz]');
ylabel('[Mb/s]');
legend('Capacity', 'Waterfilling with QAM gap','QAM rate with gap, LC and CP');