clear all; close all; clc;
%%% 6.973 HW1 %%%
% A simple equalization script

% Define channel pulse response
p=[0.9 1];
sigma2_noise=0.181; % assume unit energy

% Define equalizer parameters
Lmax=20; % equalizer length
% delta=4; % equalizer delay

for L=1:Lmax,
    for index=1:L,
        [wzfe,sigma2_zfe(L,index), sigma2_zfe_res_isi(L,index), SINRzfe(L,index), wmmse, sigma2_mmse(L,index), sigma2_mmse_res_isi(L,index), SINRmmse(L,index)] = linear_equalizer(p,L,index-1, sigma2_noise);
    end
end

[SINRzfe_vec, index_opt_zfe]=max(SINRzfe,[],2);
[SINRmmse_vec, index_opt_mmse]=max(SINRmmse,[],2);

for k=1:Lmax,
    sigma2_zfe_res_isi_vec(k)=sigma2_zfe_res_isi(k,index_opt_zfe(k));
    sigma2_mmse_res_isi_vec(k)=sigma2_mmse_res_isi(k,index_opt_zfe(k));
end
for k=1:Lmax,
    sigma2_zfe_vec(k)=sigma2_zfe(k,index_opt_zfe(k));
    sigma2_mmse_vec(k)=sigma2_mmse(k,index_opt_zfe(k));
end
figure(1)
plot(1:Lmax,10*log10(SINRzfe_vec),'b.-'); hold on;
plot(1:Lmax,10*log10(SINRmmse_vec),'rsq-');
xlabel('Equalizer length');
ylabel('SINR [dB]');
title('Linear equalizers, mmse vs. zfe');
legend('zfe-le','mmse-le');
figure(2)
plot(1:Lmax,sigma2_zfe_vec,'b.-'); hold on;
plot(1:Lmax,sigma2_mmse_vec,'rsq-');
xlabel('Equalizer length');
ylabel('Residual ISI power');
title('Linear equalizers, mmse vs. zfe');
legend('zfe-le','mmse-le');
figure(3)
plot(1:Lmax,index_opt_zfe-1,'b.-'); hold on;
plot(1:Lmax,index_opt_mmse-1,'rsq-');
xlabel('Equalizer length');
ylabel('Optimal delay');
title('Linear equalizers, mmse vs. zfe');
legend('zfe-le','mmse-le');

%%% Punch the appropriate columns of P matrix to implement the DFE
%%% equalizer (try both ZFE-DFE and MMSE-DFE)
Lmax_dfe=20;
for L=1:Lmax_dfe,
    for index=1:L,
    [wzfe_ff, wzfe_fb,sigma2_zfedfe(L,index), SINR_zfedfe(L,index), wmmse_ff,wmmse_fb,sigma2_mmsedfe(L,index), SINR_mmsedfe(L,index)]=dfe_equalizer(p,L,min(Lmax_dfe-L,length(p)-1+L-index-1),index, sigma2_noise);
    end
end

[SINR_zfedfe_vec, index_opt_zfe]=max(SINR_zfedfe,[],2);
[SINR_mmse_dfe_vec, index_opt_mmse]=max(SINR_mmsedfe,[],2);
figure(4)
plot(1:Lmax_dfe,10*log10(SINR_zfedfe_vec),'b.-'); hold on;
plot(1:Lmax_dfe,10*log10(SINR_mmse_dfe_vec),'rsq-');
xlabel('Total Equalizer length');
ylabel('SINR [dB]');
title('DFE equalizers, mmse vs. zfe');
legend('zfe-dfe','mmse-dfe');
figure(5)
plot(1:Lmax_dfe,index_opt_zfe-1,'b.-'); hold on;
plot(1:Lmax_dfe,index_opt_mmse-1,'rsq-');
xlabel('Total Equalizer length');
ylabel('Optimal delay');
title('DFE equalizers, mmse vs. zfe');
legend('zfe-dfe','mmse-dfe');

L=10; delta=6;Lff=9; Lfb=1;deltaff=8;
[wzfe,sigma2_zfe, sigma2_zfe_res_isi, SINRzfe, wmmse, sigma2_mmse, sigma2_mmse_res_isi, SINRmmse] = linear_equalizer(p,L,delta, sigma2_noise);
[wzfe_ff, wzfe_fb,sigma2_zfedfe, SINR_zfedfe, wmmse_ff,wmmse_fb,sigma2_mmsedfe, SINR_mmsedfe]=dfe_equalizer(p,Lff,Lfb,deltaff, sigma2_noise);
peq_zfe=conv(wzfe,p); [maxp,indexMain_zfe]=max(peq_zfe);
peq_mmse=conv(wmmse,p); [maxp,indexMain_mmse]=max(peq_mmse);
peq_zfedfe=conv(wzfe_ff,p); [maxp,indexMain_zfedfe]=max(peq_zfedfe); peq_zfedfe(indexMain_zfedfe:indexMain_zfedfe)=0; %%% impact of feedback section
peq_mmsedfe=conv(wmmse_ff,p); [maxp,indexMain_mmsedfe]=max(peq_mmsedfe); peq_mmsedfe(indexMain_mmsedfe+1:indexMain_mmsedfe+1)=0; %%% impact of feedback section

w=-pi:pi/100:pi;
for k=1:length(w),
    D=exp(-j*w(k));
    dvec_zfe(k,:)=D.^((1:length(peq_zfe))-indexMain_zfe);
    dvec_mmse(k,:)=D.^((1:length(peq_mmse))-indexMain_mmse);
    dvec_zfedfe(k,:)=D.^((1:length(peq_zfedfe))-indexMain_zfedfe);
    dvec_mmsedfe(k,:)=D.^((1:length(peq_mmsedfe))-indexMain_mmsedfe);

    dvec_wzfe(k,:)=D.^(1:length(wzfe));
    dvec_wmmse(k,:)=D.^(1:length(wmmse));
    dvec_wzfedfe(k,:)=D.^(1:length(wzfe_ff));
    dvec_wmmsedfe(k,:)=D.^(1:length(wmmse_ff));

    Peq_zfe(k)=20*log10(abs(peq_zfe*dvec_zfe(k,:)'));
    Peq_mmse(k)=20*log10(abs(peq_mmse*dvec_mmse(k,:)'));
    Peq_zfedfe(k)=20*log10(abs(peq_zfedfe*dvec_zfedfe(k,:)'));
    Peq_mmsedfe(k)=20*log10(abs(peq_mmsedfe*dvec_mmsedfe(k,:)'));
    
    Weq_zfe(k)=20*log10(abs(wzfe*dvec_wzfe(k,:)'));
    Weq_mmse(k)=20*log10(abs(wmmse*dvec_wmmse(k,:)'));
    Weq_zfedfe(k)=20*log10(abs(wzfe_ff*dvec_wzfedfe(k,:)'));
    Weq_mmsedfe(k)=20*log10(abs(wmmse_ff*dvec_wmmsedfe(k,:)'));
end
figure(6)
subplot(221);stem(peq_zfe,'b'); xlabel('sample'); ylabel('amplitude'); title('ZFE');
subplot(222);stem(peq_mmse,'r'); xlabel('sample'); ylabel('amplitude'); title('MMSE');
subplot(223);stem(peq_zfedfe,'g'); xlabel('sample'); ylabel('amplitude'); title('ZFE-DFE');
subplot(224);stem(peq_mmsedfe,'m'); xlabel('sample'); ylabel('amplitude'); title('MMSE-DFE');

figure(7)
subplot(221);plot(w,Peq_zfe,'b'); xlabel('frequency'); ylabel('amplitude'); title('ZFE');
subplot(222);plot(w,Peq_mmse,'r'); xlabel('frequency'); ylabel('amplitude'); title('MMSE');
subplot(223);plot(w,Peq_zfedfe,'g'); xlabel('frequency'); ylabel('amplitude'); title('ZFE-DFE');
subplot(224);plot(w,Peq_mmsedfe,'m'); xlabel('frequency'); ylabel('amplitude'); title('MMSE-DFE');

figure(8)
subplot(221);stem(wzfe,'b'); xlabel('sample'); ylabel('amplitude'); title('ZFE');
subplot(222);stem(wmmse,'r'); xlabel('sample'); ylabel('amplitude'); title('MMSE');
subplot(223);stem(wzfe_ff,'g'); xlabel('sample'); ylabel('amplitude'); title('ZFE-DFE');
subplot(224);stem(wmmse_ff,'m'); xlabel('sample'); ylabel('amplitude'); title('MMSE-DFE');

figure(9)
subplot(221);plot(w,Weq_zfe,'b'); xlabel('frequency'); ylabel('amplitude'); title('ZFE');
subplot(222);plot(w,Weq_mmse,'r'); xlabel('frequency'); ylabel('amplitude'); title('MMSE');
subplot(223);plot(w,Weq_zfedfe,'g'); xlabel('frequency'); ylabel('amplitude'); title('ZFE-DFE');
subplot(224);plot(w,Weq_mmsedfe,'m'); xlabel('frequency'); ylabel('amplitude'); title('MMSE-DFE');