clear; clc;

M = 4;
randkey = 1;
duob_mode = db_mode.no_db;

mlse = 1;
dbtgt = 1;

link_lengths = 0:1:8;          % [m] --- outer loop
rops          = linspace(-10, 0, 12);  % [dBm] --- inner sweep
FEC_thr       = 3.8e-3;               % BER target
baudrate = 200e9;

% --- allocate results
reqROP_FFE      = nan(size(link_lengths));
reqROP_MLSE     = nan(size(link_lengths));
reqROP_DBTGT    = nan(size(link_lengths));
reqROP_ML_MLSE2 = nan(size(link_lengths));
reqROP_ML_MLSE3 = nan(size(link_lengths));

%% ====================== OUTER LOOP ======================
for L = 1:numel(link_lengths)
    link_length_m = link_lengths(L);
    fprintf('\n=== %.0f m fiber length ===\n', link_length_m);

    ber_ffe  = nan(size(rops));
    ber_mlse = nan(size(rops));
    ber_dbtgt = nan(size(rops));
    ber_ml2  = nan(size(rops));
    ber_ml3  = nan(size(rops));

    %% -------- inner ROP loop --------
    parfor i = 1:length(rops)
        rop = rops(i);

        [Rx_sig_2sps_v1, Symbols_v1, Tx_bits_v1] = standard_link_model( ...
            "M",M,"fsym",baudrate,"rop",rop,"laser_wavelength",1290, ...
            "link_length_km",link_length_m,"random_key",1);

        Rx_sig_2sps_v1.spectrum("displayname",'Rx Sig','normalizeTo0dB',1);

        %% FFE + MLSE
        if mlse
            pf_ncoeffs  = 1;
            ffe_order   = [50, 0, 0];
            mu_ffe      = [0.0001, 0.0008, 0.001];
            mu_dfe      = 0.0004;
            eq_  = EQ("Ne",ffe_order,"Nb",[0,0,0],"training_length",2^13, ...
                "training_loops",5,"dd_loops",5,"K",2,"DCmu",0.005, ...
                "DDmu",[mu_ffe mu_dfe],"DFEmu",0.005,"FFEmu",0, ...
                "plotfinal",0,"ideal_dfe",1);
            pf_  = Postfilter("ncoeff",pf_ncoeffs,"useBurg",1);
            mlse_ = MLSE("duobinary_output",0,'M',M, ...
                'trellis_states',PAMmapper(M,0).levels,'scale_mode',2, ...
                'trellis_exclusion',0,'trellis_state_mode',2,'debug',0, ...
                'DIR',pf_.coefficients);

            [ffe_results, mlse_results] = vnle_postfilter_mlse(eq_, pf_, mlse_, M, ...
                Rx_sig_2sps_v1, Symbols_v1, Tx_bits_v1, ...
                "precode_mode", duob_mode,'showAnalysis', 0, "postFFE", [], ...
                "eth_style_symbol_mapping", 0);

            ber_ffe(i)  = ffe_results.metrics.BER;
            ber_mlse(i) = mlse_results.metrics.BER;
        end

        %% FFE + duobinary target MLSE
        if dbtgt
            mlse_db_ = MLSE("DIR",[1,1],"duobinary_output",0,"M",M, ...
                "trellis_states",PAMmapper(M,0).levels,'scale_mode',2, ...
                'trellis_exclusion',0,'trellis_state_mode',3);
            ffe_order = [50, 0, 0];
            mu_ffe = [0.0001, 0.0008, 0.001];
            mu_dfe = 0.0004;
            eq_ = EQ("Ne",ffe_order,"Nb",[0,0,0],"training_length",2^13, ...
                "training_loops",5,"dd_loops",5,"K",2,"DCmu",0.005, ...
                "DDmu",[mu_ffe mu_dfe],"DFEmu",0.005,"FFEmu",0, ...
                "plotfinal",0,"ideal_dfe",1);

            dbt_results = duobinary_target(eq_, mlse_db_, M, Rx_sig_2sps_v1, ...
                Symbols_v1, Tx_bits_v1, "precode_mode", duob_mode, ...
                'showAnalysis', 0, "postFFE", []);
            ber_dbtgt(i) = dbt_results.metrics.BER;
        end

        %% ML-based MLSE (L=2)
        mu_ml = 0.1; training_epochs = 100;
        ml_mlse_equalizer = ML_MLSE("epochs_tr",training_epochs,"epochs_dd",1, ...
            "len_tr",2^16,"mu_dd",mu_ml,"mu_tr",mu_ml,"order",11,"sps",2, ...
            "traceback_depth",128,"L",2,"delta",4,"adaptive_mu",0);
        [y_ml_mlse,y_ref] = ml_mlse_equalizer.process(Rx_sig_2sps_v1,Symbols_v1);
        ref_bits = PAMmapper(M,0).demap(y_ref);
        ml_bits  = PAMmapper(M,0).demap(y_ml_mlse);
        [~,~,ber_ml2(i)] = calc_ber(ml_bits.signal, ref_bits.signal, ...
            "skip_front",10,"skip_end",10);

        %% ML-based MLSE (L=3)
        ml_mlse_equalizer = ML_MLSE("epochs_tr",training_epochs,"epochs_dd",1, ...
            "len_tr",2^16,"mu_dd",mu_ml,"mu_tr",mu_ml,"order",11,"sps",2, ...
            "traceback_depth",128,"L",3,"delta",4,"adaptive_mu",0);
        [y_ml_mlse,y_ref] = ml_mlse_equalizer.process(Rx_sig_2sps_v1,Symbols_v1);
        ref_bits = PAMmapper(M,0).demap(y_ref);
        ml_bits  = PAMmapper(M,0).demap(y_ml_mlse);
        [~,~,ber_ml3(i)] = calc_ber(ml_bits.signal, ref_bits.signal, ...
            "skip_front",10,"skip_end",10);
    end % ROP loop

    %% --- find required ROP (FEC crossing)
    reqROP_FFE(L)      = interp_fec_cross(rops, ber_ffe,   FEC_thr);
    reqROP_MLSE(L)     = interp_fec_cross(rops, ber_mlse,  FEC_thr);
    reqROP_DBTGT(L)    = interp_fec_cross(rops, ber_dbtgt, FEC_thr);
    reqROP_ML_MLSE2(L) = interp_fec_cross(rops, ber_ml2,   FEC_thr);
    reqROP_ML_MLSE3(L) = interp_fec_cross(rops, ber_ml3,   FEC_thr);

    fprintf('Length %.0f m: FFE %.1f, MLSE %.1f, DB %.1f, ML2 %.1f, ML3 %.1f\n', ...
        link_length_m, reqROP_FFE(L), reqROP_MLSE(L), reqROP_DBTGT(L), ...
        reqROP_ML_MLSE2(L), reqROP_ML_MLSE3(L));
end

%% ====================== PLOT REQUIRED ROP ======================
cols = cbrewer2('Set1',8);
colFFE      = cols(1,:); 
colMLSE     = cols(2,:);
colDBTGT    = cols(4,:);
colML_MLSE  = cols(3,:);

figure(); hold on
plot(link_lengths, reqROP_FFE,      '-o','Color',colFFE,     'DisplayName','FFE');
plot(link_lengths, reqROP_MLSE,     '-s','Color',colMLSE,    'DisplayName','FFE+PF+MLSE');
plot(link_lengths, reqROP_DBTGT,    '--^','Color',colDBTGT,  'DisplayName','DB tgt. MLSE');
plot(link_lengths, reqROP_ML_MLSE2, '-v','Color',colML_MLSE, 'DisplayName','ML-based MLSE (L=2)');
plot(link_lengths, reqROP_ML_MLSE3, '-d','Color',colML_MLSE*0.8,'DisplayName','ML-based MLSE (L=3)');

xlabel('Link length [km]');
ylabel('Required ROP [dBm]');
title(sprintf('Required ROP the reach FEC threshold (3.8e-3); %.0f GBd PAM-%d', baudrate.*1e-9, M));
legend('Location','northwest');
grid on;
beautifyBERplot("logscale",0,"polyfit",1,"polyorder",3,"fitmethod",'smoothingspline');