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

    % Eq = FFE_DFE("epochs_tr",5,"epochs_dd",5,"len_tr",4096*2,"ffe_mu_dd",1e-4,"dfe_mu_dd",5e-4,"ffe_mu_tr",0,"dfe_mu_tr",0,"ffe_order",21,"dfe_order",0,"sps",2,"decide",1);

    properties
        sps % usually 2
        ffe_order
        dfe_order
        e
        b
        error

        len_tr
        ffe_mu_tr
        dfe_mu_tr
        epochs_tr

        ffe_mu_dd
        dfe_mu_dd
        epochs_dd

        constellation

        decide
    end

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

                options.sps = 2;

                options.ffe_order = 15;
                options.dfe_order = 2;

                options.len_tr = 4096;
                options.ffe_mu_tr = 0;
                options.dfe_mu_tr = 0;
                options.epochs_tr = 5;

                options.ffe_mu_dd = 1e-5;
                options.dfe_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.ffe_order,1);
            obj.b = zeros(obj.dfe_order,1);
            obj.error = 0;

        end

        function [X] = 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.ffe_mu_tr,obj.dfe_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.ffe_mu_dd,obj.dfe_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.ffe_order),' tap FFE'];
            X  = X.logbookentry(lbdesc); % append to logbook
            

        end

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

            arguments
                obj 
                x 
                d 
                ffe_mu 
                dfe_mu
                epochs 
                N 
                training 
                showviz
            end


            mu = diag([ones(1,obj.ffe_order(1))*ffe_mu(1) ...
                       ones(1,obj.dfe_order(1))*dfe_mu(1) ]);


            x = [zeros(floor(obj.ffe_order/2),1); x; zeros(obj.ffe_order,1)];
            %d = [zeros(obj.dfe_order-1,1);        d; zeros(obj.dfe_order,1)];
            d_ = zeros(obj.dfe_order(1),1);
            coeff = [obj.e;obj.b];

            if showviz
                f = figure(111);
                subplot(2,2,1:2);
                hold on
                a = scatter(1:numel(x),x,1,'.');
                a2 = scatter(1,1,1,'.');
                a3 = scatter(1,1,2,'.');
                a4 = xline(1);
                ylim([-3 3])
                xlim([0 length(x)]);
                subplot(2,2,3:4)
                c = stem(obj.e);
                ylim([-1 1])
                drawnow
            end
            
            for epoch = 1 : epochs
                symbol = 0;
                for sample = 1 : obj.sps : N

                    symbol = symbol+1;

                    x_ = x(obj.ffe_order+sample-1:-1:sample);
                    v = [x_;d_];

                    y(symbol,1) = coeff.' * v;    % Calculating output of LMS  __ * |

                    if training
                        d_hat(symbol,1) = d(symbol);  
                    else
                        [~,symbol_idx] = min(abs(y(symbol) - obj.constellation)); % decision for closest constellation point
                        d_hat(symbol,1) = obj.constellation(symbol_idx);
                    end

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

                    if ~all(mu == 0,'all') %not all mu values are zero
                        coeff = coeff - (mu * err(symbol) * v) ;    % Weight update rule of LMS
                    else
                        normalizationfactor =  (v.' * v);
                        coeff = coeff - err(symbol) * v / normalizationfactor;    % Weight update rule of NLMS
                    end

                    % Append new decision to decision feedback
                    if obj.dfe_order(1) > 0
                        %shift up one index
                        d_(2:end) = d_(1:end-1);
                        %replace 1st index with current estimation
                        d_(1) = d_hat(symbol);
                    end
                    
                    if mod(sample,100) == 1 && showviz
                        a2.XData = 1:2*numel(y);
                        a2.YData = repelem(y, 2);
                        a3.XData = 1:2*numel(d_hat);
                        a3.YData = repelem(d_hat, 2);
                        a4.Value = sample;
                        c.YData = obj.e;
                        drawnow;
                    end

                    obj.error(epoch,symbol) = err(symbol) * err(symbol)';  % Instantaneous square error

                end
            end

            obj.e = coeff(1:obj.ffe_order);
            obj.b = coeff(obj.ffe_order+1:end);



        end

    end
end