%%
clear all;
close all;

%% Choose a fitting base leading to the FSO Data
base = "C:\Users\magf\Desktop\Desktop\MATLAB-Zeugs\FSO Equalizer\FSO_FP_QCL_60umUTC\";
mode = 0; %0 oder 1
M = 4;

all_files = dir(fullfile(base, "**/*.mat"));

% data_tr_mf is the already recovered and filtered data
if M == 2
    tx_data_path = fullfile(base, "14G_PAM2\tx_info\tx_info_PAM2_14Gbd0.75RRC.mat");
    filename = fullfile(base, "14G_PAM2\M=2_Rs=1.4e10_Fs=8e10_I=265mA_RoP=46.3mW_L=31m_PS=RRC_rolloff=0.75_Mode=Rise.mat");
    data_tr_mf = load("C:\Users\magf\Desktop\Desktop\MATLAB-Zeugs\FSO Equalizer\FSO_FP_QCL_60umUTC\Already Recovered and Filtered\AfterSync_M=2_Rs=1.4e10_Fs=8e10_I=265mA_RoP=46.3mW_L=31m_PS=RRC_rolloff=0.75_Mode=Rise.mat");
elseif M == 4
    tx_data_path = fullfile(base, "6G_PAM4\tx_info\tx_info_PAM4_6Gbd0.6RRC.mat");
    filename = fullfile(base, "6G_PAM4\M=4_Rs=6e9_Fs=8e10_I=255mA_RoP=42.3mW_L=31m_PS=RRC_rolloff=0.6_Mode=Rise.mat");
    data_tr_mf = load("C:\Users\magf\Desktop\Desktop\MATLAB-Zeugs\FSO Equalizer\FSO_FP_QCL_60umUTC\Already Recovered and Filtered\AfterSync_M=4_Rs=6e9_Fs=8e10_I=255mA_RoP=42.3mW_L=31m_PS=RRC_rolloff=0.6_Mode=Rise.mat");
end

if mode == 1
    [f, p] = uigetfile(fullfile(base, "**/*.mat")); 
    if f~=0
        filename = fullfile(p,f);
    end
end

tx_data = load(tx_data_path);
datas = load(filename);

%%
str = filename;
M_       = str2double(regexp(str, 'M=([^_]+)',       'tokens', 'once'));
assert(M==M_);
fsym      = str2double(regexp(str, 'Rs=([^_]+)',      'tokens', 'once'));
fs      = str2double(regexp(str, 'Fs=([^_]+)',      'tokens', 'once'));
I       = sscanf(char(regexp(str, 'I=([^_]+)',       'tokens', 'once')), '%f');
rop     = sscanf(char(regexp(str, 'RoP=([^_]+)',     'tokens', 'once')), '%f');
L       = sscanf(char(regexp(str, 'L=([^_]+)',       'tokens', 'once')), '%f');
pulseshape = string(    regexp(str, 'PS=([^_]+)',      'tokens', 'once'));
rolloff = str2double(regexp(str, 'rolloff=([^_]+)', 'tokens', 'once'));
mode    = string(    regexp(str, 'Mode=([^\.]+)',   'tokens', 'once'));

%%
% Tx data

Bits = Informationsignal(tx_data.tx_data,"fs",fsym);
Symbols = Informationsignal(real(tx_data.tx_PAM_sym),"fs",fsym);

mapping_style = M==4; % Pam2 is like move-it; PAM-4 is different, same mapping like ETH peopled used in Zurich... hence the "eth_style" argument here and there
PM = PAMmapper(M,0,"eth_style",mapping_style); % one should rename "eth style" as this is simply a different mapping scheme

Symbols_ = PM.map(Bits) .* PM.scaling;
assert(isequal(Symbols.signal,Symbols_.signal));

Bits_ = PM.demap(Symbols);
[bits,errors,ber,errorIndice] = calc_ber(Bits_.signal,Bits.signal);
assert(ber == 0);

%% For comparison, apply pulsef on Tx Symbols
Pform = Pulseformer("fsym",fsym,"fdac",fs,"pulse","rrc","pulselength",16,"alpha",rolloff);
Digi_sig_compare = Pform.process(Symbols);

MF = Pulseformer("fsym",fsym,"fdac",2*fsym,"pulse","rrc","pulselength",16,"alpha",rolloff);
Rx_sig_compare = MF.process(Digi_sig_compare);

%%

% Rx Data
traceData = datas.tr.lastData(2).trace.ch3;

%FYI: Voltage=(RawData−YReference)×YIncrement+YOrigin
scoperead_volts = (traceData.RawData - traceData.YReference) * traceData.YIncrement + traceData.YOrigin;
demystified = isequal(traceData.YData,scoperead_volts);
assert(demystified);

Scope_sig = Electricalsignal(traceData.YData,"fs",fs);

Scope_sig.plot("displayname",'raw','fignum',100);
Scope_sig.spectrum("displayname",'raw','fignum',101)

%Calculate Transfer Function
Tx_spectrum = Digi_sig_compare;
Rx_spectrum = Scope_sig;

[~,Rx_synced_spectrum,inverted_spectrum,sequenceFound_spectrum,sequenceStarts_spectrum] = Rx_spectrum.tsynch("reference", Symbols, "fs_ref", fsym, "debug_plots", 1);
Rx_spectrum = Rx_synced_spectrum{1};

Tx_spectrum = Tx_spectrum.resample("fs_out",fsym);
Tx_spectrum.signal = fft(Tx_spectrum.signal);

Rx_spectrum = Rx_spectrum.resample("fs_out",fsym);
Rx_spectrum.signal = fft(Rx_spectrum.signal);

H_transfer = Rx_spectrum.signal./Tx_spectrum.signal;
H_inv = 1./H_transfer;

%Number of Samples/Symbol after Matched Filter
Kov = 14;

Scope_sig = Scope_sig.resample('fs_in',fs,'fs_out',Kov*fsym);

% 1) matched filter
% pulse is symmetric, hence we can use pulsef firectly as matched filter.
% It feels off (bit I think correct) that the fsym is now the output freq.!!
% -> output 2 sps to omit timing recovery!?
Matched_Filter = Pulseformer("fsym",fsym,"fdac",Kov*fsym,"pulse","rrc","pulselength",16,"alpha",rolloff,"matched",1);
Rx_matched = Matched_Filter.process(Scope_sig);
Rx_matched.spectrum("displayname",'Signal after matched filter','fignum',1);

% Loading and synchronizing their matched filtered and timing recovered data
data_tr_mf = Electricalsignal(data_tr_mf.Results, "fs", fsym);
[~,Rx_synced_cell_tr_mf,inverted_tr_mf,sequenceFound_tr_mf,sequenceStarts_tr_mf] = data_tr_mf.tsynch("reference", Symbols, "fs_ref", fsym, "debug_plots", 1);
Rx_tr_mf = Rx_synced_cell_tr_mf{1};

% Timing sync -> at this point we still have no symbol timing recovery, we
% try to do this with 2sps EQ!
[~,Rx_synced_cell,inverted,sequenceFound,sequenceStarts] = Rx_matched.tsynch("reference", Symbols, "fs_ref", fsym, "debug_plots", 1);
Rx_matched_1 = Rx_synced_cell{1};

% Timing recovery
Time_Rec = 1;
Timing_Mode = 3;
if Time_Rec
    if Timing_Mode == 1 % Zero-Crossing/Gardner/Early-Late/Müller-Mueller Timing Recovery
        [Rx_Time_Rec, Timing_Error] = Timing_Recovery("modulation", 'PAM/PSK/QAM', "timing_error_detector",'Gardner (non-data-aided)','sps',Kov,'damping_factor',1,'normalized_loop_bandwidth',1e-4,'detector_gain',2.7).process(Rx_matched_1);
    elseif Timing_Mode == 2 % Godard Timing Recovery (resampling is required as the Godard timing recovery does not change the number of samples per symbol)
        [Rx_Time_Rec, Timing_Error_MG] = Godard_Timing_Recovery('mode',3,'num_blocks',1,'fft_length',length(Rx_matched_1),'sps',Kov,'rolloff',0.6,'mu',-0.2,'Ki',1e-4).process(Rx_matched_1);
        Rx_Time_Rec = Rx_Time_Rec.resample('fs_in',Kov*fsym,'fs_out',fsym);
    elseif Timing_Mode == 3 % Maximum Variance Timing Recovery
        Rx_Time_Rec = MaxVar_Timing_Recovery('mode',0,'fsym',fsym,'fadc',Kov*fsym,'num_tau',Kov*128,'sps',Kov,'comp_signal',Rx_tr_mf,'comp_mode',0).process(Rx_matched_1);
    end
else
    Rx_Time_Rec = Rx_matched_1.resample('fs_in',Kov*fsym,'fs_out',fsym);
end
sps = 1;
Rx_Time_Rec.fs = fsym;

% % Compensate using the inverse transfer function
% Amax = 10;
% H_inv = min(abs(H_inv), Amax) .* exp(1j*angle(H_inv));
% Rx_Time_Rec.signal = fft(Rx_Time_Rec.signal);
% Rx_Time_Rec.signal = ifft(Rx_Time_Rec.signal .* H_inv);

% Normalization
Rx_Time_Rec = Rx_Time_Rec.normalize('mode','rms');
Rx_tr_mf = Rx_tr_mf.normalize('mode','rms');

% Compare spectra of our and their time signal
Rx_Time_Rec.spectrum("displayname",'Our signal','normalizeTo0dB',1,'fignum',101111);
% Rx_matched_original.spectrum('normalizeTo0dB',1,"displayname",'Our signal','fignum',101111);
Rx_tr_mf.spectrum("displayname",'Their signal','normalizeTo0dB',1,'fignum',101111);

% Compare our and their time signal
Rx_Time_Rec.normalize("mode","rms").plot("displayname",'Our signal','fignum',101311);
% Rx_matched_original.normalize("mode","rms").plot("displayname",'Original signal','fignum',101311);
Rx_tr_mf.normalize("mode","rms").plot("displayname",'Their signal','fignum',101311);

%%
% Use our or their signal
for our_signal = 1
    if our_signal
        Rx_synced = Rx_Time_Rec;
    else
        Rx_synced = Rx_tr_mf;
    end
    % Rx_synced = Rx_Time_Rec;
    % Rx_synced = Rx_synced_cell{1};
    len_tr  = 4096*2;
    mu_ffe1 = 0.0001;
    mu_ffe2 = 0.0008;
    mu_ffe3 = 0.001;
    mu_dc = 0.004;
    mu_ffe = [mu_ffe1 mu_ffe3 mu_ffe3];
    mu_dfe = 0.0004;
    duob_mode = db_mode.no_db;
    
    Rx_synced.plot("displayname",'RX: Matched+Sync+2sps','fignum',103);
    Rx_synced.spectrum("displayname",'RX: Matched+Sync+2sps','fignum',104);
    
    Digi_sig_compare.normalize("mode","rms").spectrum("displayname",'Tx: RC-shaped','fignum',1,'normalizeTo0dB',1);
    Rx_sig_compare.normalize("mode","rms").spectrum("displayname",'Tx: RC-shaped + matched filtered ','fignum',1,'normalizeTo0dB',1);
    Rx_synced.normalize("mode","rms").spectrum("displayname",'RX: matched filtered + synced','fignum',1,'normalizeTo0dB',1);
    
    if M == 2
        ber_in_paper = 10^(-2.6); %fig 3a) 4 Gb/s MWIR FSO Transmission using Directly Modulated QCL and an Uncooled UTC-PD at Room-Temperature
    elseif M == 4
        ber_in_paper = 10^(-2.5);
    end
    
    %% -------------------- FFE --------------------
    % requires some more digging what is going on :-) 
    % eq_ffe      = EQ("Ne",[150, 0, 0],"Nb",[0,0,0], ...
    %     "training_length",len_tr,"training_loops",5,"dd_loops",5, ...
    %     "K",sps,"DCmu",mu_dc,"DDmu",[mu_ffe mu_dfe],"DFEmu",0.005, ...
    %     "FFEmu",0,"plotfinal",0,"ideal_dfe",1);
    % 
    % % eq_ffe = FFE_DFE('ffe_order',99,'dfe_order',99,'len_tr',len_tr,'epochs_tr',5,'epochs_dd',5,'sps',sps,'decide',0);
    % 
    % ffe_results = ffe(eq_ffe,M,Rx_synced,Symbols,Bits, ...
    %     "precode_mode",duob_mode,'showAnalysis',0,"postFFE",[], ...
    %     "eth_style_symbol_mapping",mapping_style);
    % 
    % if our_signal
    %     fprintf('Our signal: %.1e \n',ffe_results.metrics.BER);
    %     fprintf('Paper: %.1e \n \n',ber_in_paper);
    % else
    %     fprintf('Their signal: %.1e \n',ffe_results.metrics.BER);
    %     fprintf('Paper: %.1e \n \n',ber_in_paper);
    % end
    
    %% -------------------- VNLE + MLSE --------------------
    % pf_ncoeffs  = 4;
    % eq_v        = EQ("Ne",[300, 0, 0],"Nb",[0, 0, 0], ...
    %     "training_length",len_tr,"training_loops",5,"dd_loops",5, ...
    %     "K",sps,"DCmu",mu_dc,"DDmu",[mu_ffe mu_dfe],"DFEmu",0.005, ...
    %     "FFEmu",0,"plotfinal",0,"ideal_dfe",1, ...
    %     'weighted_DFE',0,'weighted_DFE_d_min',0.5,'weighted_DFE_mode','R2','weighted_DFE_I_mode',[5,0.5,0.6]);
    % % eq_v = FFE_DFE('ffe_order',300,'dfe_order',5,'len_tr',len_tr,'epochs_tr',5,'epochs_dd',5, ...
    % %     'ffe_mu_dd',mu_ffe,'ffe_mu_tr',0,'dfe_mu_dd',mu_dfe,'dfe_mu_tr',0.005,'sps',sps,'decide',0);
    % pf_         = Postfilter("ncoeff",pf_ncoeffs,"useBurg",1);
    % mlse_ = MLSE("duobinary_output",0,'M',M,'trellis_states',PAMmapper(M,0,"eth_style",mapping_style).levels);
    % 
    % [vnle_results, mlse_results] = vnle_postfilter_mlse(eq_v, pf_, mlse_, M, Rx_synced, Symbols, Bits, ...
    %     "precode_mode", duob_mode, 'showAnalysis', 1, "postFFE", [], "eth_style_symbol_mapping", mapping_style);
    % 
    % mlse_results.metrics.print
    % if our_signal
    %     fprintf('Our Signal: %.1e \n',mlse_results.metrics.BER);
    %     fprintf('Paper: %.1e \n \n',ber_in_paper);
    % else
    %     fprintf('Their Signal: %.1e \n',mlse_results.metrics.BER);
    %     fprintf('Paper: %.1e \n \n',ber_in_paper);
    % end
    
    %% -------------------- ML-based MLSE (L=2) --------------------
    ml_mlse_equalizer = ML_MLSE("epochs_tr",150,"epochs_dd",1, ...
        "len_tr",length(Rx_synced),"mu_dd",0.03,"mu_tr",0.03,"order",210,"sps",sps, ...
        "traceback_depth",256,"L",1,"delta",4,"adaptive_mu",0);

    [ml_mlse_results] = ml_mlse(ml_mlse_equalizer, M, Rx_synced, Symbols, Bits,"precode_mode",duob_mode,"eth_style_symbol_mapping",mapping_style);
    if our_signal
        fprintf('Our Signal: %.1e \n',ml_mlse_results.metrics.BER);
        fprintf('Paper: %.1e \n \n',ber_in_paper);
    else
        fprintf('Their EQ: %.1e \n',ml_mlse_results.metrics.BER);
        fprintf('Paper: %.1e \n \n',ber_in_paper);
    end
end