clear;
col = linspecer(6);

% GENERATE SIR CURVE
if ismac
    foldername = '/Users/silasoettinghaus/Documents/MATLAB/Labor_Datensatz_PAM4_MPI/pam4_10km_1km';
else
    foldername = 'C:\Users\Silas\Documents\MATLAB\Datensätze\Labor_Datensatz_PAM4_MPI_OFC2023\pam4_10km_1km';
end


% SETTINGS

optimize_mudc = 0;
run_sir_sweep = 0;
run_feed_forward = 1;
run_baseline  = 0;
run_ideal_dc_tap = 0;
plot_timesignal = 1;

block_loop = [1];

for bl = 1:numel(block_loop)

    eq_parallelization_blocklength = block_loop(bl);
    eq_updatelatency = 3;
    eq_avg_blocklength =0;
    
    % FIND BEST MU DC
    
    if optimize_mudc
        mudc_loop = [0, 0.0001,0.0005, 0.001,0.005, 0.01:0.01:0.1, 0.2:0.1:1];
        
        current_filename = 'pam4__loop_14_92Gbd_19092023_1513.mat';
        current_filename = 'pam4__loop_29_56Gbd_19092023_1546.mat';
        recorded_data = load([foldername,filesep, current_filename]);
    
        [best_mudc,ber]= optimizeMuDc(recorded_data,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength,mudc_loop);
    
%         figure(16)
%         hold on
%         plot(mudc_loop,ber);
%         set(gca, 'YScale', 'log');
%         set(gca, 'XScale', 'log');
%         scatter(best_mudc,ber(mudc_loop==best_mudc),100,'Marker','x','LineWidth',2)
%         yline(ber(1));
%         xlim([0,1])
    end
    
    if run_sir_sweep
    
        mudc = best_mudc;
        [ber,sir,fsym] = runSIRsweep(eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength,mudc);

        plotSirSweep(ber,sir,fsym,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength,mudc)

        crossing = calculateCrossing([92e9, 56e9],ber,sir,fsym);
    
        req_sir_56(bl) = crossing(2,2) ;
        req_sir_92(bl) = crossing(2,1) ;

    end


end

if run_feed_forward

        current_filename = 'pam4__loop_14_92Gbd_v19092023_1513.mat';
        recorded_data = load([foldername,filesep, current_filename]);
        eq_avg_blocklength = [50,100,1000,3000];
        [best_block,ber]= optimizeAvgBlocklength(recorded_data,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength,0);


        best_block = 400;

        [ber_base,sir_base,fsym_base] = runSIRsweep(eq_parallelization_blocklength,eq_updatelatency,best_block,0);

        plotSirSweep(ber_base,sir_base,fsym_base,eq_parallelization_blocklength,eq_updatelatency,best_block,0)

        crossing = calculateCrossing([92e9, 56e9],ber_base,sir_base,fsym_base);

        [ber_base,sir_base,fsym_base] = runSIRsweep(eq_parallelization_blocklength,eq_updatelatency,0,0);

        plotSirSweep(ber_base,sir_base,fsym_base,eq_parallelization_blocklength,eq_updatelatency,0,0)

        crossing = calculateCrossing([92e9, 56e9],ber_base,sir_base,fsym_base);

end


if run_baseline

    mudc = 0;
    [ber_base,sir_base,fsym_base] = runSIRsweep(eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength,mudc);

    plotSirSweep(ber_base,sir_base,fsym_base,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength,mudc)

    crossing = calculateCrossing([92e9, 56e9],ber_base,sir_base,fsym_base);

    req_sir_56_baseline = crossing(2,2) ;
    req_sir_92_baseline = crossing(2,1) ;

end

if run_ideal_dc_tap

    current_filename = 'pam4__loop_29_56Gbd_19092023_1546.mat';
    recorded_data = load([foldername,filesep, current_filename]);
    [mudc,~]= optimizeMuDc(recorded_data,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength,mudc_loop);

    eq_parallelization_blocklength = 1;
    eq_updatelatency = 1;
    [ber_ideal,sir_ideal,fsym_ideal] = runSIRsweep(eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength,mudc);

    plotSirSweep(ber_ideal,sir_ideal,fsym_ideal,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength,mudc)

    crossing = calculateCrossing([92e9, 56e9],ber_ideal,sir_ideal,fsym_ideal);

    req_sir_56_ideal = crossing(2,2) ;
    req_sir_92_ideal = crossing(2,1) ;

end



if plot_timesignal
    
    current_filename = 'pam4__loop_26_92Gbd_19092023_1542.mat';
    recorded_data = load([foldername,filesep, current_filename]);

    sir = recorded_data.saveStructTemp.awg2scope_keysight_state.eigenlight_mpi - 3.3 + 7.2;

    fsym = 92e9;
    eq_parallelization_blocklength = 1;
    eq_updatelatency = 1;
    eq_avg_blocklength = 0;

    % TRUE SYMBOLS
    correct_symbols     = recorded_data.saveStructTemp.digi_mod_out';

        %3) PLAIN RECEIVED SIGNAL
    disp('RX Signal')
    y_rx = Electricalsignal(recorded_data.saveStructTemp.dp_tsynch_out');
    y_rx.fs = 2.*fsym;
    rx_symbols          = y_rx.resample("fs_in",2.*fsym,"fs_out",fsym);
    rx_symbols          = rx_symbols.normalize("mode","rms").signal;
    disp(std(rx_symbols));
    scatterleveldependent(rx_symbols,correct_symbols,fsym);

    %1) PLOT WITH WITH ALGORITHM
    if 0
    mudc_loop = [0, 0.0001,0.0005, 0.001,0.005, 0.01:0.01:0.1, 0.2:0.1:1];
    [mudc,~]= optimizeMuDc(recorded_data,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength,mudc_loop);
    end 

    disp('ALGORITHM ')
    mudc = 0.05;
    results_alg = runPostProcessing(recorded_data,mudc,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength);
    disp(results_alg.ber);
    rx_symbols = results_alg.EQ_out.signal; 
    disp(std(rx_symbols));
    scatterleveldependent(rx_symbols,correct_symbols,fsym);

    %2) JUST EQ
    disp('Just EQ')
    mudc = 0;
    results = runPostProcessing(recorded_data,mudc,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength);
    disp(results.ber)
    rx_symbols = results.EQ_out.signal; 
    disp(std(rx_symbols));
    scatterleveldependent(rx_symbols,correct_symbols,fsym);
    



end


function plotSirSweep(ber,sir,fsym,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength,mudc)


        rate = [92e9, 56e9];
        figure(95)

        for r = 1:length(rate)

            sorted = sortrows([sir(fsym == rate(r)); ber(fsym == rate(r))]',1)';
            hold on
            plot(abs(sorted(1,:)),sorted(2,:),'DisplayName',[num2str(rate(r).*1e-9),' GBd; PAM4; mudc: ',num2str(mudc)],'LineWidth',1,'Marker','o','MarkerSize',5,'LineStyle','-','HandleVisibility','on');

        end

        set(gca, 'YScale', 'log');
        yline(3.8e-3,'LineWidth',1, 'LineStyle','--','HandleVisibility','off');
        xlim([15,35]);
        xlabel('SIR in dB');
        ylabel('BER');

end

function [crossing] = calculateCrossing(rate,ber,sir,fsym)

    for r = 1:length(rate)
        sorted = sortrows([sir(fsym == rate(r)); ber(fsym == rate(r))]',1)';
        berVals = sorted(2,:);
        xAxis = sorted(1,:);
        
        hdfec = 3.8e-3 .* ones(1,length(sorted));
        crossing(:,r) = InterX([hdfec(:)';xAxis],[berVals;xAxis]) ; 
    end

end

function [ber,sir,fsym] = runSIRsweep(eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength,mudc)

    % GENERATE SIR CURVE
    if ismac
        foldername = '/Users/silasoettinghaus/Documents/MATLAB/Labor_Datensatz_PAM4_MPI/pam4_10km_1km';
    else
        foldername = 'C:\Users\Silas\Documents\MATLAB\Labor_Datensatz_PAM4_MPI\pam4_10km_1km';
    end
    
    allfiles = dir(foldername);
    
    parfor i = 1:length(allfiles)
    
        if allfiles(i).bytes ~= 0
            current_filename = allfiles(i).name;
            recorded_data = load([foldername,filesep, current_filename]);
        else
            continue
        end
   
        results = runPostProcessing(recorded_data,mudc,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength);
        ber(i) = results.ber;
        fsym(i) = results.fsym;
        sir(i) = results.sir
    
        disp(['fsym: ',num2str(results.fsym*1e-9),'GBd, SIR: ',num2str(results.sir),' ->> BER: ',sprintf('%2E',results.ber)]);
    
    end
end

function [best_blocklength,ber] = optimizeAvgBlocklength(recorded_data,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength_loop,mudc)

    for j = 1:length(eq_avg_blocklength_loop)    
        eq_avg_blocklength = eq_avg_blocklength_loop(j);

    
        results = runPostProcessing(recorded_data,mudc,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength);
        ber(j) = results.ber;
        
        disp(['fsym: ',num2str(results.fsym*1e-9),'GBd, SIR: ',num2str(results.sir),' ->> BER: ',sprintf('%2E',results.ber)]);
    
    end
    
    [~,pos]=min(ber);
    best_blocklength = eq_avg_blocklength_loop(pos);

end

function [best_mudc,ber] = optimizeMuDc(recorded_data,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength,mudc_loop)

    parfor j = 1:length(mudc_loop)    
        mudc = mudc_loop(j);

    
        results = runPostProcessing(recorded_data,mudc,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength);
        ber(j) = results.ber;
        
        disp(['fsym: ',num2str(results.fsym*1e-9),'GBd, SIR: ',num2str(results.sir),' ->> BER: ',sprintf('%2E',results.ber)]);
    
    end
    
    [~,pos]=min(ber);
    best_mudc = mudc_loop(pos);

end

function results = runPostProcessing(recorded_data,mudc,eq_parallelization_blocklength,eq_updatelatency,eq_avg_blocklength)

        results.sir  = recorded_data.saveStructTemp.awg2scope_keysight_state.eigenlight_mpi - 3.3 + 7.2;
        results.fsym = recorded_data.saveStructTemp.common.f_sym;

        results.fdac = recorded_data.saveStructTemp.common.f_DAC;
        results.fadc = recorded_data.saveStructTemp.common.f_ADC;


        % 0) build RX Signal
        y_rx        = Electricalsignal(recorded_data.saveStructTemp.dp_tsynch_out');
        y_rx.fs     = 2.*results.fsym;

        % 0) build tx reference Signal for eq training
        y_digimod   = Electricalsignal(recorded_data.saveStructTemp.digi_mod_out');
        y_digimod.fs = results.fsym;


        % 1) normlaize
        y_rx = y_rx.normalize("mode","rms");

        eq = EQ_silas("Ne",[25,0,0],"Nb",[2,0,0],"trainlength",4096,...
            "sps",2,...
            "mu_dc_dd",mudc,...
            "mu_dc_train",mudc,...
            "mu_ffe_train",0,...
            "mu_dfe_train",0.005,...
            "mu_ffe_dd",[0.0004 0.0004 0.0004],...
            "mu_dfe_dd",0.005,...
            "ddloops",3,...
            "trainloops",4,...
            "eq_parallelization_blocklength",eq_parallelization_blocklength, ...
            "eq_updatelatency",eq_updatelatency,...
            "eq_avg_blocklength",eq_avg_blocklength);

        [Eq_out] = eq.process(y_rx,y_digimod);

        results.EQ_out = Eq_out;

        % 3) digital demodulation object
        digimod = PAMmapper(2^recorded_data.saveStructTemp.common.M,0);
        
        %estimated/ equalized symbol sequence
        d_estimated = digimod.demap(Eq_out);

        %correct data symbols
        d_correct = recorded_data.saveStructTemp.prms_out;

        % 4) BER calculation
        [totalbits,errors,results.ber,loc]     = calc_ber(d_estimated.signal(1:end,:)    ,d_correct(1:end,:)',"skip",0,"returnErrorLocation",1);
end

function std_per_lvl = calcleveldependentstd(rx_symbols,correct_symbols)
        error_of_rx_signal =  rx_symbols - correct_symbols;
        levels = unique(correct_symbols);
        for l = 1:4
            level_amplitude = levels(l);
            std_per_lvl(l) = var(( 1/64 .* movsum(error_of_rx_signal(correct_symbols==level_amplitude),[64/2,64/2])  ));
            %std_per_lvl(l) = std(error_of_rx_signal(correct_symbols==level_amplitude));
        end
end

function scatterleveldependent(rx_symbols,correct_symbols,f_sym)
        col = cbrewer2('Paired',8);
        ccnt = -1;
        figure1 = figure();
        levels = unique(correct_symbols);

        start = 1;
        ende = length(correct_symbols);

        start = 30000;
        ende = 40000;
        
        for l = 1:4
            ccnt = ccnt+2;

            level_amplitude = levels(l);

            symbols_for_lvl = NaN(1,length(correct_symbols));
            symbols_for_lvl(correct_symbols==level_amplitude) = rx_symbols(correct_symbols==level_amplitude);
            std_lvl(l) = std(symbols_for_lvl,'omitnan');
            xax_in_sec = ((1:length(correct_symbols)) / f_sym) * 1e6;
            xax_in_sec = 1:length(correct_symbols);

            scatter(xax_in_sec(start:ende),symbols_for_lvl(start:ende),10,'.','MarkerFaceAlpha',0.5,'MarkerEdgeAlpha',0.5,'MarkerEdgeColor',col(ccnt,:));
            hold on;
        end

        std_lvl = round(std_lvl,2);
        disp(std_lvl);

        ccnt = 0;

        % Add the windowed/ smoothed curves
        for l = 1:4
            ccnt = ccnt+2;
            level_amplitude = levels(l);

            symbols_for_lvl = NaN(1,length(correct_symbols));

            movmean = 1/250 .* movsum(rx_symbols(correct_symbols==level_amplitude),[250/2,250/2]);

            symbols_for_lvl(correct_symbols==level_amplitude) = movmean;

            nanx = isnan(symbols_for_lvl);
            t    = 1:numel(symbols_for_lvl);
            symbols_for_lvl(nanx) = interp1(t(~nanx), symbols_for_lvl(~nanx), t(nanx));

            xax_in_sec = ((1:length(correct_symbols)) / f_sym) * 1e6;
            xax_in_sec = 1:length(correct_symbols);

            plot(xax_in_sec(start:ende),symbols_for_lvl(start:ende),'Color',col(ccnt,:));
            hold on
        end

        %yline(max(rx_symbols(correct_symbols==levels(2))))

        if 0
        annotation(figure1,'textbox',...
            [0.660523809523809 0.844444444444448 0.133523809523809 0.0603174603174607],...
            'String',['\sigma = ',num2str(std_lvl(4))],...
            'LineWidth',1.8,...
            'LineStyle','none',...
            'FontSize',12,...
            'FitBoxToText','off');
        
        % Create textbox
        annotation(figure1,'textbox',...
            [0.667666666666665 0.642857142857147 0.133523809523809 0.0603174603174607],...
            'String',['\sigma = ',num2str(std_lvl(3))],...
            'LineWidth',1.8,...
            'LineStyle','none',...
            'FontSize',12,...
            'FitBoxToText','off');
        
        % Create textbox
        annotation(figure1,'textbox',...
            [0.671238095238093 0.442857142857148 0.133523809523809 0.0603174603174608],...
            'String',['\sigma = ',num2str(std_lvl(2))],...
            'LineWidth',1.8,...
            'LineStyle','none',...
            'FontSize',12,...
            'FitBoxToText','off');
        
        % Create textbox
        annotation(figure1,'textbox',...
            [0.670047619047616 0.265079365079371 0.133523809523809 0.0603174603174608],...
            'String',['\sigma = ',num2str(std_lvl(1))],...
            'LineWidth',1.8,...
            'LineStyle','none',...
            'FontSize',12,...
            'FitBoxToText','off');
        end

        xlim([0, 2.6])
        ylim([-2 2])
        xlabel('Time in $\mu$s');
        ylabel('Normalized Amplitude');

end