classdef ML_MLSE < handle
    % ---------------------------------------------------------------------
    % W. Lanneer and Y. Lefevre,
    % “Machine Learning-Based Pre-Equalizers for Maximum Likelihood
    %  Sequence Estimation in High-Speed PONs,” EUSIPCO 2023
    % ---------------------------------------------------------------------
    % This implementation reproduces the closed-loop ML-based
    % pre-equalizer training for MLSE, supporting both training and
    % detection (decision-directed) modes.
    % ---------------------------------------------------------------------

    properties
        sps
        order
        e
        e_tr
        error

        len_tr
        mu_tr
        epochs_tr

        dd_mode
        mu_dd
        epochs_dd
        adaptive_mu

        constellation
        L
        alpha
        DIR
        DIR_flip
        trellis_states
        traceback_depth
        delta

        % Internal variables
        S
        Nf
        nStates
        nFeasible
        combs
        first_sym
        last_sym
        valid
        valid_to_idx
        valid_from_idx
        w

        % Fast lookup
        nSym
        key_table
        trans_index
        true_to_state_idx

        % Debug metrics
        ber = []
        ce  = ones(1,1)
    end

    methods
        function obj = ML_MLSE(options)
            arguments(Input)
                options.sps = 2;
                options.order = 15;
                options.len_tr = 4096;
                options.mu_tr = 0.001;
                options.epochs_tr = 5;
                options.dd_mode = 1;
                options.mu_dd = 1e-5;
                options.epochs_dd = 5;
                options.adaptive_mu = 1;
                options.delta = 0;
                options.traceback_depth = 1024;
                options.L = 1;
            end

            fn = fieldnames(options);
            for n = 1:numel(fn)
                obj.(fn{n}) = options.(fn{n});
            end

            obj.e = zeros(obj.order,1);
            obj.error = 0;
        end

        % ==============================================================
        % PROCESS
        % ==============================================================
        function [X,X_viterbi] = process(obj, X, D)
            % Normalize input RMS
            X = X.normalize("mode","rms");
            obj.constellation = sort(unique(D.signal),'ascend');
            obj.nSym = numel(obj.constellation);

            if length(X)/length(D) ~= obj.sps
                warning('Signal length does not fit to reference!');
            end

            % --- Parameters
            obj.S = obj.nSym;
            obj.Nf = obj.order * obj.sps;
            obj.nStates = obj.S^obj.L;
            obj.nFeasible = obj.nStates * obj.S;

            % --- Trellis mapping
            obj.trellis_states = reshape(obj.constellation,1,[]);
            pre_comb_mat  = repmat(obj.trellis_states, obj.L, 1);
            pre_comb_cell = mat2cell(pre_comb_mat, ones(1,obj.L), size(pre_comb_mat,2));
            obj.combs     = fliplr(combvec(pre_comb_cell{:}).');
            obj.first_sym = obj.combs(:,1);
            obj.last_sym  = obj.combs(:,end);
            obj.nStates   = size(obj.combs,1);

            % --- Valid transitions
            obj.valid = false(obj.nStates);
            for from = 1:obj.nStates
                for to = 1:obj.nStates
                    if all(obj.combs(to,2:end) == obj.combs(from,1:end-1))
                        obj.valid(to,from) = true;
                    end
                end
            end
            [obj.valid_to_idx,obj.valid_from_idx] = find(obj.valid);

            % --- Initialize weights
            if isempty(obj.w) || any(size(obj.w) ~= [obj.Nf+1,obj.nFeasible])
                obj.w = randn(obj.Nf+1,obj.nFeasible);
            end

            % --- Fast lookup tables
            [~, sym_idx_mat] = ismember(obj.combs, obj.constellation);
            key_vals = 1 + sum((sym_idx_mat - 1) .* (obj.nSym .^ (0:obj.L-1)), 2);
            max_key = obj.nSym^obj.L;
            obj.key_table = zeros(max_key,1,'uint32');
            obj.key_table(key_vals) = 1:obj.nStates;

            obj.trans_index = sparse(obj.nStates,obj.nStates);
            for i = 1:length(obj.valid_from_idx)
                f = obj.valid_from_idx(i);
                t = obj.valid_to_idx(i);
                obj.trans_index(t,f) = i;
            end

            % ==============================================================
            % TRAINING
            % ==============================================================
            fprintf('\n--- Training mode ---\n');
            obj.equalize(X.signal, D.signal, obj.mu_tr, obj.epochs_tr, obj.len_tr, true);
            obj.e_tr = obj.e;

            % ==============================================================
            % DECISION-DIRECTED / TESTING
            % ==============================================================
            fprintf('--- Decision-directed / detection mode ---\n');
            [y, y_vit] = obj.equalize(X.signal, D.signal, obj.mu_dd, obj.epochs_dd, X.length, false);

            X_viterbi = X;
            X.signal = y;
            X_viterbi.signal = y_vit;
        end

        % ==============================================================
        % EQUALIZE
        % ==============================================================
        function [y,y_ref] = equalize(obj,x,d,mu,epochs,N,training)
            debug = 0;
            showPlots = 0;
            y = zeros(N,1);
            nSymbols = ceil(N/obj.sps);

            for epoch = 1:epochs
                pm = zeros(obj.nStates,1);
                pred = zeros(nSymbols,obj.nStates,'uint32');
                pm_sto = nan(obj.nStates,nSymbols,'like',pm);
                CE_accum = 0;

                start_sample = 1;
                end_sample   = N;
                start_symbol = 1 + floor((start_sample - 1)/obj.sps);

                % --- initialize true state
                if numel(d) >= obj.L && start_symbol >= obj.L
                    init_seq = d(start_symbol-obj.L+1:start_symbol);
                    key_init = obj.seq2key(init_seq);
                    true_to_state_idx = obj.key_table(key_init);
                    if true_to_state_idx==0, true_to_state_idx=1; end
                else
                    true_to_state_idx = uint32(1);
                end

                for sample = start_sample:obj.sps:end_sample
                    symbol = (sample - start_sample)/obj.sps + 1;
                    sym_idx = start_symbol + (symbol - 1);

                    % --- Observation window (with delta)
                    i1 = sample - obj.Nf + 1 + obj.delta;
                    i2 = sample + obj.delta;
                    buf = x(max(1,i1):min(length(x),i2));
                    padL = max(0,1 - i1);
                    padR = max(0,i2 - length(x));
                    yk = [zeros(padL,1); buf(:); zeros(padR,1)];
                    yk = [yk;1];

                    % --- Branch metrics
                    c_hat = (yk.' * obj.w).';
                    pm = pm - min(pm);
                    v_tilde = pm(obj.valid_from_idx) + c_hat;

                    % --- allocate once
                    if epoch==1 && symbol==1
                        obj.true_to_state_idx = ones(ceil(N/obj.sps),1,'uint32');
                    end

                    % --- previous "to" becomes "from"
                    if symbol>1
                        true_from_state_idx = obj.true_to_state_idx(symbol-1);
                    else
                        true_from_state_idx = 1;
                    end

                    % --- compute or reuse "to" state
                    if epoch==1
                        if sym_idx>=obj.L
                            key_to = obj.seq2key(d(sym_idx-obj.L+1:sym_idx));
                            state_idx = obj.key_table(key_to);
                            if state_idx==0
                                state_idx = true_from_state_idx;
                            end
                            obj.true_to_state_idx(symbol) = state_idx;
                        else
                            obj.true_to_state_idx(symbol) = true_from_state_idx;
                        end
                    end
                    true_to_state_idx = obj.true_to_state_idx(symbol);

                    % --- fast Dirac creation
                    dirac = zeros(obj.nFeasible,1);
                    trans_idx = obj.trans_index(true_to_state_idx,true_from_state_idx);
                    if trans_idx~=0
                        dirac(trans_idx)=1;
                    end

                    % ===================================================================
                    % TRAINING MODE (weight update)
                    % ===================================================================
                    if training
                        % --- Softmax and CE
                        v_shift = -(v_tilde - min(v_tilde));
                        v_shift = min(v_shift,100);
                        expv = exp(v_shift);
                        p = expv./(sum(expv)+eps);
                        CE_symbol(symbol) = -log(p(dirac==1)+eps);

                        % --- CE smoothing and adaptive μ
                        if sym_idx>obj.L
                            CE_smooth(symbol)=0.01*CE_symbol(symbol)+0.99*CE_symbol(symbol-1);
                        else
                            CE_smooth(symbol)=CE_symbol(symbol);
                        end
                        CE_accum=CE_accum+CE_symbol(symbol);

                        % --- Gradient update
                        dmp=(dirac-p)';
                        dL_Dw=(yk).*dmp;
                        if sym_idx>=obj.L
                            if obj.adaptive_mu
                                mu_eff=CE_smooth(symbol);
                                mu_eff=max(min(mu_eff,0.2),1e-4);
                            else
                                mu_eff=mu;
                            end
                            obj.w=obj.w - mu_eff.*dL_Dw;
                        end
                    end

                    % ===================================================================
                    % DECODING MODE (Viterbi only)
                    % ===================================================================
                    % Compare-Select (always executed)
                    vmat=inf(obj.nStates,obj.nStates);
                    vmat(obj.valid)=v_tilde;
                    [pm_next,pred(symbol,:)]=min(vmat,[],2);
                    pm_next=pm_next-min(pm_next);
                    pm=pm_next;
                    pm_sto(:,symbol)=pm;
                end

                % --- Traceback
                [~,s_end]=min(pm);
                vpath=zeros(symbol,1,'uint32');
                vpath(symbol)=s_end;
                for n=symbol:-1:2
                    vpath(n-1)=pred(n,vpath(n));
                end

                y_ref=d(start_symbol:end);
                y=obj.first_sym(vpath);

                % --- BER/CE reporting and plots
                if training
                    err=sum(y~=y_ref(1:length(y)));
                    ser=err/length(y);
                    try
                        ref_bits=PAMmapper(obj.S,0).demap(y_ref(1:length(y)));
                        eq_bits=PAMmapper(obj.S,0).demap(y);
                        [~,~,ber,~]=calc_ber(ref_bits,eq_bits,"skip_front",10,"skip_end",10,"returnErrorLocation",1);
                        fprintf('Epoch %d - BER: %.2e\n',epoch,ber);
                        obj.ber(epoch)=ber;
                    catch
                        fprintf('Epoch %d - SER: %.2e\n',epoch,ser);
                        obj.ber(epoch)=ser;
                    end
                    obj.ce(epoch)=CE_accum/symbol;

                    if debug && mod(epoch,10)==1 && showPlots
                        figure(10);clf
                        subplot(3,2,1:2);
                        imagesc(obj.w);axis xy;colorbar;title('Filter W');
                        subplot(3,2,3);
                        vtilde_mat=NaN(obj.nStates,obj.nStates);
                        vtilde_mat(obj.valid)=v_tilde;
                        imagesc(vtilde_mat);axis xy;colorbar;title('Path Metrics (v\_tilde)');
                        subplot(3,2,4);
                        plot(1:symbol,pm_sto);title('Path Metric Evolution');
                        subplot(3,2,5);hold on;
                        scatter(1:symbol,CE_symbol,1,'.');
                        scatter(1:symbol,CE_smooth,1,'.');
                        title('Cross Entropy');
                        subplot(3,2,6);hold on;
                        yyaxis left
                        scatter(1:length(obj.ce),obj.ce,10,'s','filled');
                        ylabel('Cross Entropy');
                        yyaxis right
                        scatter(1:length(obj.ber),obj.ber,10,'d','filled');
                        set(gca,'YScale','log');
                        ylabel('BER (log)');
                        xlabel('Epoch');grid on;
                        title('Convergence');
                        drawnow;
                    end
                end
            end
        end

        % ==============================================================
        % Helper: Sequence → key (always scalar)
        % ==============================================================
        function key = seq2key(obj, seq)
            [~, idx] = ismember(flip(seq), obj.constellation);
            pow = (obj.nSym .^ (0:obj.L-1)).';
            key = 1 + sum((idx(:) - 1) .* pow);
        end
    end
end