clear; clc;

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

mlse = 1;
dbtgt = 1;

baudrates = 180e9:2e9:220e9;        % outer loop
rops      = linspace(-10,0,12);     % inner sweep
FEC_thr   = 3.8e-3;                  % BER target

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

%% ====================== OUTER LOOP ======================
for b = 1:numel(baudrates)
    baudrate = baudrates(b);
    fprintf('\n=== %.0f GBd ===\n', baudrate/1e9);

    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 --------
    for 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_linewidth",1310, ...
            "link_length_km",0,"random_key",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(b)      = interp_fec_cross(rops, ber_ffe,   FEC_thr);
    reqROP_MLSE(b)     = interp_fec_cross(rops, ber_mlse,  FEC_thr);
    reqROP_DBTGT(b)    = interp_fec_cross(rops, ber_dbtgt, FEC_thr);
    reqROP_ML_MLSE2(b) = interp_fec_cross(rops, ber_ml2,   FEC_thr);
    reqROP_ML_MLSE3(b) = interp_fec_cross(rops, ber_ml3,   FEC_thr);

    % --- diagnostic
    fprintf('Baud %.0f GBd: FFE %.1f, MLSE %.1f, DB %.1f, ML2 %.1f, ML3 %.1f\n', ...
        baudrate/1e9, reqROP_FFE(b), reqROP_MLSE(b), reqROP_DBTGT(b), ...
        reqROP_ML_MLSE2(b), reqROP_ML_MLSE3(b));
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(baudrates/1e9, reqROP_FFE,      '-o','Color',colFFE,     'DisplayName','FFE');
plot(baudrates/1e9, reqROP_MLSE,     '-s','Color',colMLSE,    'DisplayName','FFE+PF+MLSE');
plot(baudrates/1e9, reqROP_DBTGT,    '--^','Color',colDBTGT,  'DisplayName','DB tgt. MLSE');
plot(baudrates/1e9, reqROP_ML_MLSE2, '-v','Color',colML_MLSE, 'DisplayName','ML-based MLSE (L=2)');
plot(baudrates/1e9, reqROP_ML_MLSE3, '-d','Color',colML_MLSE*0.8,'DisplayName','ML-based MLSE (L=3)');

xlabel('Baud rate [GBd]');
ylabel('Required ROP [dBm]');
title('ROP required for FEC threshold');
grid on; legend('Location','northwest');
beautifyBERplot("logscale",0,"polyfit",1,"polyorder",4,"fitmethod",'polyfit');