classdef FFE_DCTracking < FFE_plain
    %FFE_DCTracking Plain FFE with the DC-tracking offset loop.
    %
    % This class reuses the plain FFE process idea and BER mu optimization,
    % but keeps only the DC-tracking path from the larger FFE class. A1, A2,
    % and delayed FFE tap-update buffers are intentionally excluded.

    properties
        dc_tracking_mu
        dc_tracking_adaptive_enabled
        dc_tracking_persistence_gain
        dc_tracking_power_exponent
        dc_tracking_buffer_len
        e_dc
    end

    properties (Access = private)
        dc_tracking_mu_min = 1e-6
        dc_tracking_mu_max = 3e-1
        dc_tracking_mu_eff_min = 0
        dc_tracking_mu_eff_max = inf
    end

    methods
        function obj = FFE_DCTracking(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.dc_tracking_mu = 0
                options.dc_tracking_adaptive_enabled = false
                options.dc_tracking_persistence_gain = 0
                options.dc_tracking_power_exponent = 2
                options.dc_tracking_buffer_len = 1

                options.decide = false

                options.save_debug = 0
                options.optmize_mus = 0
                options.mu_optimization_len = 2^15
                options.plot_mu_optimization = 0
                options.mu_optimization_fignum = 3010
            end

            obj@FFE_plain( ...
                "sps",options.sps, ...
                "order",options.order, ...
                "len_tr",options.len_tr, ...
                "mu_tr",options.mu_tr, ...
                "epochs_tr",options.epochs_tr, ...
                "adaption_technique",options.adaption_technique, ...
                "dd_mode",options.dd_mode, ...
                "mu_dd",options.mu_dd, ...
                "epochs_dd",options.epochs_dd, ...
                "dd_len_fraction",options.dd_len_fraction, ...
                "decide",options.decide, ...
                "save_debug",options.save_debug, ...
                "optmize_mus",options.optmize_mus, ...
                "mu_optimization_len",options.mu_optimization_len, ...
                "plot_mu_optimization",options.plot_mu_optimization, ...
                "mu_optimization_fignum",options.mu_optimization_fignum);

            obj.dc_tracking_mu = options.dc_tracking_mu;
            obj.dc_tracking_adaptive_enabled = options.dc_tracking_adaptive_enabled;
            obj.dc_tracking_persistence_gain = options.dc_tracking_persistence_gain;
            obj.dc_tracking_power_exponent = options.dc_tracking_power_exponent;
            obj.dc_tracking_buffer_len = floor(options.dc_tracking_buffer_len);
            obj.e_dc = 0;

            assert(obj.dc_tracking_buffer_len >= 0);
        end

        function [X,Noi] = process(obj,X,D)
            X = X.normalize("mode","rms");
            obj.constellation = unique(D.signal);
            obj.resetTrackingState();

            if obj.optmize_mus
                obj.optimizeMus(X.signal,D.signal);
                obj.resetTrackingState();
            end

            training = true;
            showviz = false;
            obj.equalize(X.signal,D.signal,obj.mu_tr,obj.epochs_tr,obj.len_tr,training,showviz);
            obj.e_tr = obj.e;

            n = X.length;
            training = false;
            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

            if obj.decide
                X.signal = decision;
            else
                X.signal = signal;
            end

            X.fs = D.fs;
            X = X.logbookentry([num2str(obj.order),' tap FFE DC tracking']);

            Noi = X - D;
        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];
                otherwise
                    builtin("error","FFE_DCTracking:InvalidAdaptionTechnique", ...
                        "Unsupported FFE adaption technique.");
            end

            vars = optimizableVariable("mu_tr",mu_range,"Transform","log");
            if obj.dd_mode
                vars = [vars, optimizableVariable("mu_dd",mu_range,"Transform","log")];
            end
            optimize_dc_tracking_mu = obj.dc_tracking_mu ~= 0;
            if optimize_dc_tracking_mu
                vars = [vars, optimizableVariable("dc_tracking_mu",[1e-5,1e-1],"Transform","log")];
            end

            obj.mu_optimization_iter = 0;
            fprintf("FFE_DCTracking 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_dc_tracking_mu
                obj.dc_tracking_mu = obj.mu_optimization.XAtMinObjective.dc_tracking_mu;
            end

            if obj.dd_mode && optimize_dc_tracking_mu
                fprintf("\nFFE_DCTracking mu opt done: mu_tr=%9.3e, mu_dd=%9.3e, dc_tracking_mu=%9.3e, BER=%9.3e\n", ...
                    obj.mu_tr,obj.mu_dd,obj.dc_tracking_mu,obj.mu_optimization.MinObjective);
            elseif obj.dd_mode
                fprintf("\nFFE_DCTracking 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_dc_tracking_mu
                fprintf("\nFFE_DCTracking mu opt done: mu_tr=%9.3e, dc_tracking_mu=%9.3e, BER=%9.3e\n", ...
                    obj.mu_tr,obj.dc_tracking_mu,obj.mu_optimization.MinObjective);
            else
                fprintf("\nFFE_DCTracking mu opt done: mu_tr=%9.3e, BER=%9.3e\n", ...
                    obj.mu_tr,obj.mu_optimization.MinObjective);
            end

            if obj.plot_mu_optimization
                obj.plotMuOptimization();
            end
        end

        function objective = muObjective(obj,params,x,d)
            old_state = obj.captureTrackingObjectiveState();
            cleanup = onCleanup(@()obj.restoreTrackingObjectiveState(old_state));

            obj.save_debug = 0;
            if any(strcmp(params.Properties.VariableNames,"dc_tracking_mu"))
                obj.dc_tracking_mu = params.dc_tracking_mu;
            end
            obj.resetTrackingState();

            N_tr = min(obj.len_tr,numel(x));
            obj.equalize(x,d,params.mu_tr,obj.epochs_tr,N_tr,true,false);
            if obj.dd_mode
                [signal,~] = obj.equalize(x,d,params.mu_dd,obj.epochs_dd,numel(x),false,false);
            else
                [signal,~] = obj.equalize(x,d,0,1,numel(x),false,false);
            end

            [ber,errors] = obj.berObjective(signal,d);
            objective = ber;
            if ~isfinite(objective)
                objective = inf;
            end

            obj.mu_optimization_iter = obj.mu_optimization_iter + 1;
            has_dc_tracking_mu = any(strcmp(params.Properties.VariableNames,"dc_tracking_mu"));
            if obj.dd_mode && has_dc_tracking_mu
                fprintf("\rFFE_DCTracking mu opt %02d: mu_tr=%9.3e, mu_dd=%9.3e, dc_tracking_mu=%9.3e, BER=%9.3e, errors=%d", ...
                    obj.mu_optimization_iter,params.mu_tr,params.mu_dd,params.dc_tracking_mu,ber,errors);
            elseif obj.dd_mode
                fprintf("\rFFE_DCTracking 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_dc_tracking_mu
                fprintf("\rFFE_DCTracking mu opt %02d: mu_tr=%9.3e, dc_tracking_mu=%9.3e, BER=%9.3e, errors=%d", ...
                    obj.mu_optimization_iter,params.mu_tr,params.dc_tracking_mu,ber,errors);
            else
                fprintf("\rFFE_DCTracking mu opt %02d: mu_tr=%9.3e, BER=%9.3e, errors=%d", ...
                    obj.mu_optimization_iter,params.mu_tr,ber,errors);
            end

            clear cleanup
        end

        function [y,d_hat] = equalize(obj,x,d,mu,epochs,N,training,showviz)
            arguments
                obj
                x
                d
                mu
                epochs
                N
                training
                showviz
            end

            unused_showviz = showviz; %#ok<NASGU>
            x = x(:);
            d = d(:);
            N = obj.validSampleLength(N,x,d);
            n_symbols = N / obj.sps;
            y = zeros(n_symbols,1);
            d_hat = zeros(n_symbols,1);

            if n_symbols == 0
                return
            end

            if isempty(obj.constellation)
                obj.constellation = unique(d);
            end
            decision_constellation = obj.constellation(:);

            x = [zeros(floor(obj.order/2),1); x; zeros(obj.order,1)];
            lambda = mu;
            mask = ones(obj.order,1);
            maincursor_pos = ceil(length(obj.e)/2);
            adaption_code = obj.adaptionCode();
            adaption_is_rls = adaption_code == 3;

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

            dc_tracking_mu_eff_min = obj.dc_tracking_mu_eff_min;
            dc_tracking_mu_eff_max = obj.dc_tracking_mu_eff_max;
            dc_tracking_persistence_gain = 0;
            if obj.dc_tracking_adaptive_enabled
                dc_tracking_persistence_gain = obj.dc_tracking_persistence_gain;
            end
            dc_tracking_enabled = obj.dc_tracking_mu ~= 0;
            if dc_tracking_enabled
                obj.dc_tracking_mu = min(max(obj.dc_tracking_mu, ...
                    obj.dc_tracking_mu_min),obj.dc_tracking_mu_max);
            end
            dc_tracking_mu = obj.dc_tracking_mu;
            dc_tracking_use_persistence = dc_tracking_persistence_gain > 0;
            dc_tracking_base_mu_eff = min(max(dc_tracking_mu, ...
                dc_tracking_mu_eff_min),dc_tracking_mu_eff_max);
            dc_buffer_enabled = dc_tracking_enabled && obj.dc_tracking_buffer_len > 1;
            if dc_buffer_enabled
                dc_tracking_err_buffer = NaN(obj.dc_tracking_buffer_len,1);
                dc_tracking_err_buffer_pos = 0;
                dc_tracking_err_sum = 0;
                dc_tracking_abs_err_sum = 0;
                dc_tracking_valid_count = 0;
            else
                dc_tracking_err_buffer = [];
                dc_tracking_err_buffer_pos = 0;
                dc_tracking_err_sum = 0;
                dc_tracking_abs_err_sum = 0;
                dc_tracking_valid_count = 0;
            end

            debug_enabled = obj.save_debug;
            if debug_enabled
                obj.initializeTrackingDebug(n_symbols,training);
            end

            for epoch = 1:epochs
                symbol = 0;
                for sample = 1:obj.sps:N
                    symbol = symbol + 1;
                    grad = zeros(obj.order,1);
                    update = zeros(obj.order,1);
                    weight = 0;
                    dc_tracking_mu_eff = 0;

                    U = x(obj.order+sample-1:-1:sample);
                    y(symbol,1) = obj.e_dc + (obj.e.*mask).' * U;

                    if training
                        d_hat(symbol,1) = d(symbol);
                    else
                        [~,symbol_idx] = min(abs(y(symbol) - decision_constellation));
                        d_hat(symbol,1) = decision_constellation(symbol_idx);
                    end

                    err = d_hat(symbol) - y(symbol);

                    switch adaption_code
                        case 1
                            weight = mu / ((U.'*U) + eps);
                            grad = err * U;
                            update = grad * weight;

                        case 2
                            weight = mu;
                            grad = err * U;
                            update = grad * weight;

                        case 3
                            denom = lambda + U.' * obj.P * U;
                            k = (obj.P * U) / denom;
                            update = k * err;
                    end

                    obj.e = obj.e + update;

                    if adaption_is_rls
                        obj.P = (1/lambda) * (obj.P - k * (U.' * obj.P));
                    end

                    if dc_tracking_enabled
                        [obj.e_dc,dc_tracking_mu_eff,dc_tracking_err_buffer, ...
                            dc_tracking_err_buffer_pos,dc_tracking_err_sum, ...
                            dc_tracking_abs_err_sum,dc_tracking_valid_count] = ...
                            obj.updateDcTracking( ...
                            err,dc_tracking_mu,dc_tracking_base_mu_eff, ...
                            dc_tracking_persistence_gain,dc_tracking_use_persistence, ...
                            dc_buffer_enabled,dc_tracking_mu_eff_min,dc_tracking_mu_eff_max, ...
                            symbol,dc_tracking_err_buffer,dc_tracking_err_buffer_pos, ...
                            dc_tracking_err_sum,dc_tracking_abs_err_sum, ...
                            dc_tracking_valid_count);
                    end

                    if debug_enabled && epoch == epochs
                        error_power = err * err';
                        update_power = update.'*update ./ (sqrt((obj.e.'*obj.e) / obj.order) + eps);
                        obj.debug_struct.error(1,symbol) = error_power;
                        obj.debug_struct.main_cursor(1,symbol) = abs(obj.e(maincursor_pos));
                        obj.debug_struct.mu_nlms(1,symbol) = weight;
                        obj.debug_struct.update_gradient(1,symbol) = grad.'*grad;
                        obj.debug_struct.dc_tracking_mu_eff(1,symbol) = dc_tracking_mu_eff;
                        obj.debug_struct.dc_tracking_est(1,symbol) = obj.e_dc;

                        if training
                            obj.debug_struct.error_tr(1,symbol) = error_power;
                            obj.debug_struct.update_tr(1,symbol) = update_power;
                        else
                            obj.debug_struct.error_dd(1,symbol) = error_power;
                            obj.debug_struct.update(1,symbol) = update_power;
                        end
                    end
                end
            end
        end
    end

    methods (Access = private)
        function resetTrackingState(obj)
            obj.e = zeros(obj.order,1);
            obj.e_tr = zeros(obj.order,1);
            obj.error = 0;
            obj.e_dc = 0;
            obj.P = (1/0.05) * eye(obj.order);
            obj.debug_struct = struct();
        end

        function state = captureTrackingObjectiveState(obj)
            state.e = obj.e;
            state.e_tr = obj.e_tr;
            state.error = obj.error;
            state.e_dc = obj.e_dc;
            state.P = obj.P;
            state.save_debug = obj.save_debug;
            state.debug_struct = obj.debug_struct;
            state.dc_tracking_mu = obj.dc_tracking_mu;
        end

        function restoreTrackingObjectiveState(obj,state)
            obj.e = state.e;
            obj.e_tr = state.e_tr;
            obj.error = state.error;
            obj.e_dc = state.e_dc;
            obj.P = state.P;
            obj.save_debug = state.save_debug;
            obj.debug_struct = state.debug_struct;
            obj.dc_tracking_mu = state.dc_tracking_mu;
        end

        function N = validSampleLength(obj,N,x,d)
            N_available = min(numel(x),numel(d) * obj.sps);
            N = min(N,N_available);
            N = obj.sps * floor(N / obj.sps);
            N = max(0,N);
        end

        function adaption_code = adaptionCode(obj)
            if obj.adaption_technique == adaption_method.nlms
                adaption_code = 1;
            elseif obj.adaption_technique == adaption_method.lms
                adaption_code = 2;
            elseif obj.adaption_technique == adaption_method.rls
                adaption_code = 3;
            else
                builtin("error","FFE_DCTracking:InvalidAdaptionTechnique", ...
                    "Unsupported FFE adaption technique.");
            end
        end

        function initializeTrackingDebug(obj,n_symbols,training)
            obj.debug_struct.error = NaN(1,n_symbols);
            obj.debug_struct.error_first_epoch = NaN(1,n_symbols);
            obj.debug_struct.main_cursor = NaN(1,n_symbols);
            obj.debug_struct.mu_nlms = NaN(1,n_symbols);
            obj.debug_struct.update_gradient = NaN(1,n_symbols);
            obj.debug_struct.dc_tracking_mu_eff = NaN(1,n_symbols);
            obj.debug_struct.dc_tracking_est = NaN(1,n_symbols);

            if training
                obj.debug_struct.error_tr = NaN(1,n_symbols);
                obj.debug_struct.update_tr = NaN(1,n_symbols);
            else
                obj.debug_struct.error_dd = NaN(1,n_symbols);
                obj.debug_struct.update = NaN(1,n_symbols);
            end
        end

        function [e_dc,mu_eff,err_buffer,err_buffer_pos,err_sum,abs_err_sum,valid_count] = ...
                updateDcTracking(obj,err_current,mu_dc,base_mu_eff,persistence_gain, ...
                use_persistence,buffer_enabled,mu_eff_min,mu_eff_max,symbol, ...
                err_buffer,err_buffer_pos,err_sum,abs_err_sum,valid_count)

            mu_eff = 0;
            e_dc = obj.e_dc;

            if buffer_enabled
                err_buffer_pos = err_buffer_pos + 1;
                if err_buffer_pos > obj.dc_tracking_buffer_len
                    err_buffer_pos = 1;
                end

                old_err = err_buffer(err_buffer_pos);
                if isfinite(old_err)
                    err_sum = err_sum - old_err;
                    valid_count = valid_count - 1;
                    if use_persistence
                        abs_err_sum = abs_err_sum - abs(old_err);
                    end
                end

                if isfinite(err_current)
                    err_buffer(err_buffer_pos) = err_current;
                    err_sum = err_sum + err_current;
                    valid_count = valid_count + 1;
                    if use_persistence
                        abs_err_sum = abs_err_sum + abs(err_current);
                    end
                else
                    err_buffer(err_buffer_pos) = NaN;
                end

                if mod(symbol,obj.dc_tracking_buffer_len) == 0
                    if valid_count == 0
                        err_mean = 0;
                        err_abs_mean = 0;
                    else
                        err_mean = err_sum / valid_count;
                        err_abs_mean = abs_err_sum / valid_count;
                    end

                    mu_eff = obj.effectiveDcMu(mu_dc,base_mu_eff,persistence_gain, ...
                        use_persistence,err_mean,err_abs_mean,mu_eff_min,mu_eff_max);
                    e_dc = e_dc + mu_eff * err_mean;
                end
            else
                if isfinite(err_current)
                    err_mean = err_current;
                    err_abs_mean = abs(err_current);
                else
                    err_mean = 0;
                    err_abs_mean = 0;
                end

                mu_eff = obj.effectiveDcMu(mu_dc,base_mu_eff,persistence_gain, ...
                    use_persistence,err_mean,err_abs_mean,mu_eff_min,mu_eff_max);
                e_dc = e_dc + mu_eff * err_mean;
            end
        end

        function mu_eff = effectiveDcMu(~,mu_dc,base_mu_eff,persistence_gain, ...
                use_persistence,err_mean,err_abs_mean,mu_eff_min,mu_eff_max)

            if use_persistence
                persistence_scale = abs(err_mean) / (err_abs_mean + eps);
                persistence_scale = min(max(persistence_scale,0),1);
                mu_eff = mu_dc * (1 + persistence_gain * persistence_scale);
                mu_eff = min(max(mu_eff,mu_eff_min),mu_eff_max);
            else
                mu_eff = base_mu_eff;
            end
        end
    end
end