classdef ML_MLSE < handle
    % Implementation of plain and simple FFE.
    % 1) Training mode (stable performance when you use NLMS)
    % 2) Decision directed mode
    %
    %LMS: mu in order of 0.0001 for acceptable convergence speed
    %NLMS: mu in order of 0.01 for acceptable convergence speed
    %RLS: mu is lambda -> 0.99 -> 1 (has a strong dependency on this! use a loop to find out best values)
    %
    %  FFE("epochs_tr",5,"epochs_dd",2,"len_tr",2^13,"mu_dd",mu_dd,"mu_tr",mu_tr,"order",25,"sps",2,"decide",0, "adaption",adaption_method(adaption),"dd_mode",use_dd_mode);

    properties
        sps % usually 2
        order
        e
        e_tr
        error

        len_tr
        mu_tr
        epochs_tr

        dd_mode % 1 or 0 to set DD-mode on or off
        mu_dd %weight update in dd mode
        epochs_dd

        constellation

        L %viterbi memory length

        alpha
        DIR
        DIR_flip
        trellis_states

        traceback_depth

        % --- Added internal class variables used later ---
        S
        Nf
        delta
        nStates
        nFeasible
        combs
        first_sym
        last_sym
        valid
        valid_to_idx
        valid_from_idx
        w

        % --- New: fast state lookup ---
        state_dict              % containers.Map: key(sequence)->state index
        key_fmt = '%.8g_';      % key format for sequence strings
        nSym                      % |constellation|

        ber = []
    end

    methods
        function obj = ML_MLSE(options)
            arguments(Input)

                options.sps = 2;
                options.order = 15;

                options.len_tr = 4096;
                options.mu_tr = 0;
                options.epochs_tr = 5;

                options.dd_mode = 1;
                options.mu_dd = 1e-5;
                options.epochs_dd = 5;

                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

        function [X,X_viterbi] = process(obj, X, D)

            % actual processing of the signal (steps 1. - 3.)
            % 1 normalize RMS
            X = X.normalize("mode","rms");

            % Use sorted constellation for deterministic mapping
            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

            % ==============================================================
            %  INITIALIZATION (only before final epoch and detection mode)
            % ==============================================================

            % --- Parameters
            obj.S = numel(obj.constellation);      % alphabet size
            obj.Nf = obj.order*obj.sps;            % filter length
            % obj.delta = 3;%ceil(obj.Nf/2);       % delay parameter
            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{:}).');  % rows: states, columns: [x_k, x_{k-1}, ...]
            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);

            % --- Allocate vectors and weights
            % !! IF SHAPE FIT, then we already have smth there an we want
            % to start with the existing fitler-set
            if isempty(obj.w) || any(size(obj.w) ~= [obj.Nf+1,obj.nFeasible])
                obj.w = zeros(obj.Nf+1,obj.nFeasible);     % filter weights per transition + bias tap
                obj.w = randn(obj.Nf+1,obj.nFeasible);
            end

            % --- Precompute dictionary for fast state lookup (sequence -> state)
            keys = cell(obj.nStates,1);
            for i = 1:obj.nStates
                keys{i} = obj.seq_key(obj.combs(i,:));  % combs row is already [x_k, x_{k-1}, ...]
            end
            obj.state_dict = containers.Map(keys, 1:obj.nStates);

            % ==============================================================
            %  TRAINING
            % ==============================================================

            % Training Mode
            n = obj.len_tr;
            training = 1;
            obj.equalize(X.signal, D.signal,obj.mu_tr,obj.epochs_tr,n,training);
            obj.e_tr = obj.e;

            % ==============================================================
            %  DD-Mode / Fixed Mode
            % ==============================================================

            % Decision Directed Mode
            n = X.length;
            training = 0;
            [y,y_vit]=obj.equalize(X.signal, D.signal,obj.mu_dd,obj.epochs_dd,n,training);

            X_viterbi = X;

            X.signal = y;
            X.fs = D.fs; %change sampling frequency of outgoing signal from fdac e.g. 2 sps to symbol spaced = fsym
            lbdesc = [num2str(obj.order),' tap FFE'];
            X  = X.logbookentry(lbdesc); % append to logbook

            X_viterbi.signal = y_vit;
            X_viterbi.fs = D.fs; %change sampling frequency of outgoing signal from fdac e.g. 2 sps to symbol spaced = fsym
            lbdesc = [num2str(obj.order),'order FFE + PF + Viterbi'];
            X_viterbi  = X_viterbi.logbookentry(lbdesc); % append to logbook
        end

        function [y,y_vit] = equalize(obj,x,d,mu,epochs,N,training)
            % ==============================================================
            %  FFE + Whitening + ML-Based Branch Metric Estimation + Viterbi
            % ==============================================================
            debug = 1;

            % --- Input padding and preallocation
            y = zeros(N,1);

            % number of symbol steps in this block
            nSymbols = ceil(N/obj.sps);

            for epoch = 1:epochs

                % state metrics (log-domain costs): keep as column [nStates×1]
                pm        = zeros(obj.nStates,1);  % v_{k-1}(s′)
                c_hat     = zeros(1,obj.nFeasible);
                v_tilde   = zeros(1,obj.nFeasible);
                pred      = zeros(nSymbols, obj.nStates, 'uint32');
                pm_sto    = nan(obj.nStates, nSymbols,'like',pm);


                %%% START IDX 
                if training
                    max_start = length(x) - ( (ceil(N/obj.sps)-1)*obj.sps + 1 );
                    max_start = max(1, max_start);  % safety
                    start_sample = randi([1, max_start], 1); %rnd training; not really good
                    start_sample = 1;
                    end_sample   = start_sample + (ceil(N/obj.sps)-1)*obj.sps;
                else
                    start_sample = 1;
                    end_sample   = N;
                end

                start_symbol = 1 + floor((start_sample - 1)/obj.sps);   % ABSOLUTE symbol index

                if numel(d) >= obj.L && start_symbol >= obj.L
                    init_seq = d(start_symbol-obj.L+1 : start_symbol);      % [d_k-L+1 ... d_k]
                    true_to_state_idx = obj.state_dict(obj.seq_key(flip(init_seq))); % [d_k ... d_k-L+1]
                else
                    % Not enough history – fall back to state 1
                    true_to_state_idx = uint32(1);
                end

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

                    % --- Build Δ-delayed observation window y_k
                    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)];   % Nf×1
                    yk  = [yk;1];

                    % --- Predict branch metrics for all feasible transitions: c_hat
                    c_hat = (yk.' * obj.w);     % [1×nFeasible]
                    c_hat = c_hat.';            % [nFeasible×1]

                    % --- Extended path metrics: v_tilde = pm(from) + c_hat
                    % normalize pm to avoid growth (invariant to additive const)
                    pm = pm - min(pm);
                    v_tilde = pm(obj.valid_from_idx) + c_hat;   % [nFeasible×1]

                    % ===== Gradient update (Algorithm 1) =====

                    if 1 %training
                        % previous "to" becomes current "from" (shift-register)
                        true_from_state_idx = true_to_state_idx;

                        % --- Build current "to" state from ABSOLUTE symbol index
                        if sym_idx >= obj.L
                            curr_seq = d(sym_idx-obj.L+1 : sym_idx);     % [d_k-L+1 ... d_k]
                            key_to   = obj.seq_key(flip(curr_seq));      % -> [d_k ... d_k-L+1]
                            if isKey(obj.state_dict, key_to)
                                true_to_state_idx = obj.state_dict(key_to);
                            else
                                % Fall back safely (should not happen with proper constellation)
                                true_to_state_idx = true_from_state_idx;
                            end
                        else
                            % Not enough history yet for a full L-symbol state
                            % keep previous 'to' and 'from'
                            true_to_state_idx   = true_to_state_idx;
                            true_from_state_idx = true_from_state_idx;
                        end

                        % Dirac delta over correct extended transition (from,to)
                        dirac = zeros(obj.nFeasible,1);
                        dirac(obj.valid_from_idx==true_from_state_idx & ...
                            obj.valid_to_idx  ==true_to_state_idx) = 1;

                        % softmax over -v_tilde (numerically safe shift)
                        p = exp(-(v_tilde - max(v_tilde)));
                        p = p./(sum(p)+eps);

                        % gradient term (t - p)
                        dmp = (dirac - p)';  % 1×nFeasible

                        % Per-feature gradient; implicit expansion gives (Nf+1)×nFeasible
                        dL_Dw = (yk) .* dmp;

                        % Start updates only when the ABSOLUTE symbol index has ≥ L history
                        if sym_idx >= obj.L


                            obj.w = obj.w - ones(size(dL_Dw,1),1).*mu .* dL_Dw;   % (Nf+1)×nFeasible
                            % obj.w = obj.w - mu * dL_Dw;   % (Nf+1)×nFeasible
                        end

                        % if debug && epoch > 2
                        %     figure(100);
                        %     subplot(4,1,1);
                        %     heatmap(p');
                        %     title('Probs')
                        %     subplot(4,1,2);
                        %     heatmap(dmp);
                        %     title('Update')
                        %     subplot(4,1,3);
                        %     heatmap(dL_Dw);
                        %     title('Update')
                        %     subplot(4,1,4);
                        %     heatmap(bj.w);
                        %     title('Update')
                        % 
                        % end

                    end



                    % --- Compare-Select (matrix form, min of costs)
                    v_tilde_mat = inf(obj.nStates, obj.nStates);
                    v_tilde_mat(obj.valid) = v_tilde;
                    [pm_next, pred(k,:)] = min(v_tilde_mat, [], 2);

                    % re-center to keep metrics bounded (decision-invariant)
                    pm_next = pm_next - min(pm_next);

                    pm = pm_next;
                    pm_sto(:,symbol) = pm;
                end

                % --- Traceback (full; you can window with traceback_depth if desired)
                [~, s_end] = min(pm);
                viterbi_path = zeros(symbol,1,'uint32');
                viterbi_path(symbol) = s_end;
                for n = symbol:-1:2
                    viterbi_path(n-1) = pred(n, viterbi_path(n));
                end

                y_vit = obj.first_sym(viterbi_path);
                y     = obj.first_sym(viterbi_path);

                if debug %&& training
                    sym_start = start_symbol;
                    sym_end   = start_symbol + symbol - 1;
                    ref_slice = d(sym_start : sym_end);
                    err = sum(y ~= ref_slice(1:numel(y)));

                    ref_bits = PAMmapper(obj.S,0).demap(ref_slice);
                    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: %.1e  \n',epoch, ber);

                    obj.ber(epoch) = ber;
                    
                    % ser = err./length(y);
                    % fprintf('Epoch: %d - SER: %.1e  \n',epoch, ser);

                    figure(10);
                    subplot(2,2,1:2);
                    heatmap(obj.w);
                    title('Filter')

                    subplot(2,2,3);
                    v_tildemat = NaN(obj.nStates, obj.nStates);
                    v_tildemat(obj.valid) = v_tilde;      % log-domain scores
                    heatmap(v_tildemat);
                    title('Path Metrics (v_tilde)')

                    subplot(2,2,4);
                    scatter(1:symbol,pm_sto,1,'.')
                    % plot(1:symbol,pm_sto,'LineStyle','none')
                    title('Path Metric Winners')

                    drawnow
                end
            end
        end
    end

    methods (Access=private)
        function k = seq_key(obj, seq)
            % Build a stable key string for a sequence row vector in the *same order as combs rows* ([x_k, x_{k-1}, ...])
            % Use rounding via sprintf to avoid floating-point issues.
            % seq must be a row vector.
            k = sprintf(obj.key_fmt, seq);
        end
    end
end