classdef FFE < 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.001 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);

    % eq_ = FFE("epochs_tr",5,"epochs_dd",5,"len_tr",len_tr, ...
    %     "mu_dd",1e-1,"mu_tr",0.4,"order",50, ...
    %     "sps",2,"decide",0,"optmize_mus",1,"dd_mode",1, ...
    %     "adaption_technique","nlms","mu_dc",1.021e-05);


    properties
        sps % usually 2
        order
        e
        e_tr
        error
        

        len_tr
        mu_tr
        epochs_tr

        adaption_technique % nlms, lms, rls
        dd_mode % 1 or 0 to set DD-mode on or off
        mu_dd %weight update in dd mode
        epochs_dd
        dd_len_fraction
        mu_dc
        e_dc

        P % covariance matrix of rls

        constellation % symbol constellation

        decide %wether to return the (hard) decisions or the result after FFE (soft)

        save_debug = 0;
        debug_struct

        optmize_mus = 0;
        mu_optimization
        mu_optimization_iter = 0;
        mu_optimization_len
    end

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

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

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

                options.adaption_technique adaption_method = adaption_method.lms;
                options.dd_mode = 1; 
                options.mu_dd = 1e-5;
                options.epochs_dd = 5;
                options.dd_len_fraction = 1;
                options.mu_dc = 0;

                options.decide = false;

                options.save_debug = 0;
                options.optmize_mus = 0;
                options.mu_optimization_len = 2^15;

            end

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

            obj.e = zeros(obj.order,1);
            obj.e_dc = 0;
            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);
            obj.e_dc = 0;


            delta   = 0.05;
            obj.P = (1/delta) * eye(obj.order);

            if obj.optmize_mus
                obj.optimizeMus(X.signal,D.signal);
                obj.e = zeros(obj.order,1);
                obj.e_dc = 0;
                obj.P = (1/delta) * eye(obj.order);
            end

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

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

            % 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,mu,epochs,N,training,showviz)

            arguments
                obj
                x
                d
                mu
                epochs
                N
                training
                showviz
            end

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

            if training
                mask = ones(obj.order,1);
            else
                mask = zeros(obj.order,1);
                mask(900:end) = 1;
                mask(ceil(length(obj.e)/2)) = 1;
            end

            mask = ones(obj.order,1);
            maincursor_pos=ceil(length(obj.e)/2);
            always_ideal_decision = 0;
            grad =0;
            weight = 0;
            update = 0;

            if mu == 0 || (~obj.dd_mode && ~training)
                epochs = 1;
            end

            for epoch = 1 : epochs
                symbol = 0;
                % obj.e_dc = 0;
                for sample = 1 : obj.sps : N

                    symbol = symbol+1;

                    U = x(obj.order+sample-1:-1:sample);

                    y(symbol,1) = obj.e_dc + (obj.e.*mask).' * U;    % Calculating output of LMS  __ * |

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

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

                    true_err(symbol) =  y(symbol) - d(symbol);   % Instantaneous error

                    if training || obj.dd_mode 
                        switch obj.adaption_technique
    
                            case adaption_method.lms
    
                                % mu used as update weight (suggestion: 0.001)
                                weight = mu;
                                grad = err(symbol) * U;
                                update = grad * weight;
                                obj.e = obj.e + update;
    
                            case adaption_method.nlms
    
                                % mu used as update weight (suggestion: 0.01-0.05; bit higher during tr)
                                normU = ((U.'*U)) + eps;
                                weight = mu / normU;
                                grad = err(symbol) * U;
                                update = grad * weight;
                                obj.e = obj.e + update;
    
                            case adaption_method.rls

    
                                % RLS‐Gain:
                                denom = lambda + U.' * obj.P * U;
                                k = (obj.P * U) / denom;
    
                                % Gewichtsupdate:
                                update = k * err(symbol);
                                obj.e = obj.e + update;
    
                                % P-Matrix‐Update:
                                obj.P = (1/lambda) * (obj.P - k * (U.' * obj.P));
   
    
                        end

                        if obj.mu_dc ~= 0
                            obj.e_dc = obj.e_dc + obj.mu_dc * err(symbol);
                        end
                    end


                    if obj.save_debug
                        obj.debug_struct.error(epoch,symbol) = err(symbol) * err(symbol)';  % Instantaneous square error
                        
                        if training
                            obj.debug_struct.error_tr(epoch,symbol) = err(symbol) * err(symbol)';  % Instantaneous square error
                            obj.debug_struct.update_tr(epoch,symbol) = update.'*update ./ rms(obj.e);
                        else
                            % if symbol == length(d)
                            %     M = numel(unique(d));
                            %     bits_ref = PAMmapper(M,0).demap(d);
                            %     bits_eq  = PAMmapper(M,0).demap(d_hat);
                            %     [~,~,obj.debug_struct.ber(epoch)] = calc_ber(bits_eq,bits_ref);
                            % end
                        end

                        % obj.debug_struct.main_cursor(epoch,symbol) = abs(obj.e(maincursor_pos));
                        % obj.debug_struct.mu_nlms(epoch,symbol) = weight;
                        % obj.debug_struct.update_gradient(epoch,symbol) = grad.'*grad;
                        % obj.debug_struct.update(epoch,symbol) = update.'*update ./ rms(obj.e);
                    
                    end

                end
            end

        end

        function N_dd = ddLength(obj,N)
            if isempty(obj.dd_len_fraction) || obj.dd_len_fraction <= 0 || obj.dd_len_fraction >= 1
                N_dd = N;
                return
            end

            N_dd = floor(N * obj.dd_len_fraction);
            N_dd = max(obj.sps,N_dd);
            N_dd = min(N,N_dd);
        end

        function optimizeMus(obj,x,d)
            [x_opt,d_opt,N_opt] = obj.optimizationSignals(x,d);

            switch obj.adaption_technique
                case adaption_method.lms
                    mu_range = [1e-5, 1e-2];
                case adaption_method.nlms
                    mu_range = [1e-3, 5e-1];
                case adaption_method.rls
                    mu_range = [0.98, 0.99999];
            end
            mu_dc_range = [1e-5, 1e-1];

            mu_tr_var = optimizableVariable("mu_tr",mu_range,"Transform","log");
            vars = mu_tr_var;
            if obj.dd_mode
                vars = [vars, optimizableVariable("mu_dd",mu_range,"Transform","log")];
            end
            optimize_mu_dc = obj.mu_dc ~= 0;
            if optimize_mu_dc
                vars = [vars, optimizableVariable("mu_dc",mu_dc_range,"Transform","log")];
            end
            obj.mu_optimization_iter = 0;
            fprintf("FFE mu opt uses %d samples / %d symbols\n",N_opt,numel(d_opt));
            obj.mu_optimization = bayesopt(@(p)obj.muObjective(p,x_opt,d_opt),vars, ...
                "MaxObjectiveEvaluations",20, ...
                "AcquisitionFunctionName","expected-improvement-plus", ...
                "IsObjectiveDeterministic",false, ...
                "Verbose",0, ...
                "PlotFcn",[]);
            obj.mu_tr = obj.mu_optimization.XAtMinObjective.mu_tr;
            if obj.dd_mode
                obj.mu_dd = obj.mu_optimization.XAtMinObjective.mu_dd;
            end
            if optimize_mu_dc
                obj.mu_dc = obj.mu_optimization.XAtMinObjective.mu_dc;
            end
            if obj.dd_mode && optimize_mu_dc
                fprintf("\nFFE mu opt done: mu_tr=%9.3e, mu_dd=%9.3e, mu_dc=%9.3e, BER=%9.3e\n", ...
                    obj.mu_tr,obj.mu_dd,obj.mu_dc,obj.mu_optimization.MinObjective);
            elseif obj.dd_mode
                fprintf("\nFFE mu opt done: mu_tr=%9.3e, mu_dd=%9.3e, BER=%9.3e\n", ...
                    obj.mu_tr,obj.mu_dd,obj.mu_optimization.MinObjective);
            elseif optimize_mu_dc
                fprintf("\nFFE mu opt done: mu_tr=%9.3e, mu_dc=%9.3e, BER=%9.3e\n", ...
                    obj.mu_tr,obj.mu_dc,obj.mu_optimization.MinObjective);
            else
                fprintf("\nFFE mu opt done: mu_tr=%9.3e, BER=%9.3e\n", ...
                    obj.mu_tr,obj.mu_optimization.MinObjective);
            end
        end

        function [x_opt,d_opt,N_opt] = optimizationSignals(obj,x,d)
            N_available = min(numel(x),numel(d) * obj.sps);

            if isempty(obj.mu_optimization_len) || obj.mu_optimization_len <= 0 || isinf(obj.mu_optimization_len)
                N_opt = N_available;
            else
                N_opt = min(N_available,max(obj.len_tr,obj.mu_optimization_len));
            end

            N_opt = obj.sps * floor(N_opt / obj.sps);
            N_opt = max(obj.sps,N_opt);

            n_symbols = N_opt / obj.sps;
            x_opt = x(1:N_opt);
            d_opt = d(1:n_symbols);
        end

        function objective = muObjective(obj,params,x,d)
            old_debug = obj.save_debug;
            old_mu_dc = obj.mu_dc;
            obj.save_debug = 0;
            has_mu_dc = any(strcmp(params.Properties.VariableNames,"mu_dc"));
            if has_mu_dc
                obj.mu_dc = params.mu_dc;
            end
            obj.e = zeros(obj.order,1);
            obj.e_dc = 0;
            obj.P = (1/0.05) * eye(obj.order);
            obj.debug_struct = struct();
            N_tr = min(obj.len_tr,numel(x));
            [signal,~] = obj.equalize(x,d,params.mu_tr,obj.epochs_tr,N_tr,1,0);
            if obj.dd_mode
                [signal,~] = obj.equalize(x,d,params.mu_dd,obj.epochs_dd,numel(x),0,0);
            end

            M = numel(unique(d));
            mapper = PAMmapper(M,0);
            eq_signal_sd = Signal(signal);
            eq_signal_hd = mapper.quantize(eq_signal_sd);
            tx_symbols = Signal(d);
            rx_bits = mapper.demap(eq_signal_hd);
            tx_bits = mapper.demap(tx_symbols);
            [~,errors,ber,~] = calc_ber(rx_bits.signal,tx_bits.signal, ...
                "skip_front",10, ...
                "skip_end",10, ...
                "returnErrorLocation",1);
            objective = ber;
            if ~isfinite(objective)
                objective = inf;
            end
            obj.mu_optimization_iter = obj.mu_optimization_iter + 1;
            if obj.dd_mode && has_mu_dc
                fprintf("\rFFE mu opt %02d: mu_tr=%9.3e, mu_dd=%9.3e, mu_dc=%9.3e, BER=%9.3e, errors=%d", ...
                    obj.mu_optimization_iter,params.mu_tr,params.mu_dd,params.mu_dc,ber,errors);
            elseif obj.dd_mode
                fprintf("\rFFE mu opt %02d: mu_tr=%9.3e, mu_dd=%9.3e, BER=%9.3e, errors=%d", ...
                    obj.mu_optimization_iter,params.mu_tr,params.mu_dd,ber,errors);
            elseif has_mu_dc
                fprintf("\rFFE mu opt %02d: mu_tr=%9.3e, mu_dc=%9.3e, BER=%9.3e, errors=%d", ...
                    obj.mu_optimization_iter,params.mu_tr,params.mu_dc,ber,errors);
            else
                fprintf("\rFFE mu opt %02d: mu_tr=%9.3e, BER=%9.3e, errors=%d", ...
                    obj.mu_optimization_iter,params.mu_tr,ber,errors);
            end
            obj.save_debug = old_debug;
            obj.mu_dc = old_mu_dc;
        end
    end
end