classdef FFE_Kalman < 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(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,showviz);

            % 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,showviz);

            % 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,showviz)

            arguments
                obj
                x
                d
                mio
                epochs
                N
                training
                showviz
            end

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

            for epoch = 1 : epochs



                symbol = 0;
                % Initialization of Kalman filter variables
                A = 1;  % State transition matrix
                H = 1;  % Observation matrix
                Q = 1e-4;  % Process noise covariance
                R = 1e-1;  % Measurement noise covariance
                P = 1;  % Initial error covariance
                mpi_est = 0;  % Initial estimate for MPI noise
                K = 0;  % Kalman gain
                subtract_mpi_est = 1;

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

                    % Get the current input sample and the equalizer output
                    U = x(obj.order + sample - 1 : -1 : sample);

                    if subtract_mpi_est || ~training
                        y(symbol,1) = (obj.e.' * U) - mpi_est .* 1 ;    % Subtract MPI estimate
                    else
                        y(symbol,1) = obj.e.' * U;
                    end

                    % Decision and error calculation
                    if training
                        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
                    
                    err(symbol) = y(symbol) - d_hat(symbol);   % Instantaneous residual error

                    true_err(symbol) = y(symbol) - d(symbol);

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

                    % Update step
                    mpi_est_new = mpi_est + K * (err(symbol) - H * mpi_est);  % MPI noise estimation
                    alpha=0;
                    mpi_est = alpha * mpi_est + (1 - alpha) * mpi_est_new;

                    P = (1 - K * H) * P;  % Update error covariance

                    % Subtract MPI noise from the signal
                    if subtract_mpi_est || ~training
                        y(symbol) = y(symbol);% - mpi_est;
                    else

                    end

                    % Equalizer weight update (LMS or 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

                    % Store MPI estimate for visualization
                    mpi_estimates(symbol) = mpi_est;

                    Kgain(symbol) = K;
                    P_(symbol) = P;
                    H_(symbol) = H;

                    
                end %symbols
            end %epoch
        
            if ~training

                if 1
                    % figure;
                    % subplot(2,2,1)
                    % hold on
                    % scatter(1:numel(y),y,1,'.');
                    % plot(1:numel(mpi_estimates), mpi_estimates);
                    % subplot(2,2,3)
                    % plot(1:numel(true_err), true_err);
                    % subplot(2,2,2)
                    % scatter(1:numel(y),y'-mpi_estimates,1,'.');
                    % xlabel('Sample Index');
                    % ylabel('MPI Noise Estimate');
                    % title('MPI Noise Estimation Over Time (Kalman Filter)');
                    % grid on;

                    figure(111);
                    
                    subplot(3,1,1)
                    hold on
                    cla
                    scatter(1:numel(y),y,1,'.');
                    plot(1:numel(mpi_estimates), mpi_estimates,'DisplayName','mpi est');
                    ylim([-2,2]);

                    subplot(3,1,2)
                    cla
                    plot(1:numel(true_err), true_err,'DisplayName','true err');
                    ylim([-1,1]);

                    subplot(3,1,3)
                    plot(1:numel(true_err), true_err-mpi_estimates,'DisplayName',['diff rms: ',num2str(rms(true_err-mpi_estimates))]);
                    ylim([-1,1]);
                    legend

                    figure(333)
                    hold on
                    von = 5000;
                    bis = 15000;
                    plot(von:bis,true_err(von:bis),'DisplayName','Error')
                    plot(von:bis,movmean(true_err(von:bis), [30,30]),'DisplayName','Error')
                    plot(von:bis,mpi_estimates(von:bis),'DisplayName','Est','LineWidth',2);
                    legend

                    figure(222)
                    hold on
                    crr = xcorr(mpi_estimates,true_err,'normalized');
                    plot(crr);
                
                end

                %  y = y-mpi_estimates';


            end

            

        end

    end
end