classdef FFE_Kalman_Feedback < handle
    % Implementation of plain and simple FFE.
    % 1) Training mode (stable performance when you use NLMS)
    % 2) Decision directed mode

    % Eq     = FFE("epochs_tr",5,"epochs_dd",5,"len_tr",4096*2,"mu_dd",1e-4,"mu_tr",0,"order",25,"sps",2,"decide",0);

    properties
        sps % usually 2
        order
        e
        error

        len_tr
        mu_tr
        epochs_tr

        mu_dd
        epochs_dd

        constellation

        decide
    end

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

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

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

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

                options.decide = false;

            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,Noi] = process(obj, X, D)

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

            obj.constellation = unique(D.signal);

            % Training Mode
            training = 1;
            showviz = 0;
            obj.equalize(X.signal, D.signal,obj.mu_tr,obj.epochs_tr,obj.len_tr,training);

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

            % Output Signal
            if obj.decide
                X.signal = decision;
            else
                X.signal = signal;
            end

            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

            Noi = X;
            Noi = X - D;


        end




        function [y, d_hat] = equalize(obj, x, d, mio, epochs, N, training)

            arguments
                obj
                x
                d
                mio
                epochs
                N
                training
            end

            % Initialize Kalman filter variables
            A = 0.7;  % State transition matrix
            H = 1;  % Observation matrix
            Q = 1e-2;  % Process noise covariance
            R = 1e-3;  % Measurement noise covariance
            P = 1;  % Initial error covariance
            mpi_est = 0;  % Initial estimate for MPI noise
            K = 0;  % Kalman gain

            % Error buffers for parallelized structure
            parallel_depth = 100;  % Depending on your parallelization depth
            err_buffer = zeros(parallel_depth, 1);  % Store collected errors
            mpi_est_buffer = zeros(parallel_depth, 1);  % Store MPI estimates

            x = [zeros(floor(obj.order/2), 1); x; zeros(obj.order, 1)];

            for epoch = 1 : epochs
                symbol = 0;

                for sample = 1 : obj.sps : N
                    symbol = symbol + 1;

                    U = x(obj.order + sample - 1 : -1 : sample);  % Input window for FFE

                    % Subtract the estimated MPI noise from the input signal before FFE
                    x_mpi_reduced = U;
                    y(symbol, 1) = obj.e.' * x_mpi_reduced;  % Output of FFE

                    y(symbol, 1) =  y(symbol, 1) - mpi_est;

                    % Decision and error calculation
                    if training
                        [~, symbol_idx]  = min(abs(d(symbol) - obj.constellation));  % Closest constellation point
                        d_hat(symbol, 1) = d(symbol);
                    else
                        [~, symbol_idx]  = min(abs(y(symbol) - obj.constellation));  % Closest constellation point
                        d_hat(symbol, 1) = obj.constellation(symbol_idx);
                    end

                    % Calculate the residual error
                    err(symbol) = y(symbol) - d_hat(symbol);   % Instantaneous residual error

                    % Collect error for feedback after 'parallel_depth' samples
                    err_buffer(mod(symbol, parallel_depth) + 1) = err(symbol);

                    % Update Kalman filter based on the accumulated errors every parallel_depth samples
                    if mod(symbol, parallel_depth) == 0
                        % Compute the mean error over the last 'parallel_depth' samples
                        avg_error = mean(err_buffer);

                        % Prediction step (Kalman filter)
                        P = A * P * A' + Q;  % Update the error covariance
                        K = P * H' / (H * P * H' + R);  % Compute Kalman gain

                        % Update MPI noise estimate
                        mpi_est = mpi_est + K * (avg_error - H * mpi_est);  % Update based on averaged error
                        P = (1 - K * H) * P;  % Update error covariance

                    end

                    % Store MPI estimate for visualization or debugging
                   k_buffer(symbol) = K;
                   mpi_est_buffer(symbol) = mpi_est;

                    % FFE weight update using NLMS
                    if mio ~= 0
                        obj.e = obj.e - (mio * err(symbol) * U);    % LMS weight update
                    else
                        normalizationfactor = (U.' * U);
                        obj.e = obj.e - err(symbol) * U / normalizationfactor;  % NLMS weight update
                    end

                end
            end
            
            if ~training
                    
                figure(111);
                    
                    subplot(3,1,1)
                    hold on
                    cla
                    scatter(1:numel(y),y,1,'.');
                    plot(1:numel(mpi_est_buffer), mpi_est_buffer,'DisplayName','mpi est');

            end
        end

    end

end