classdef FFE_DCremoval_adaptive_mu < handle
    % FFE variant for MPI/DC-removal experiments.
    % With dc_buffer_len <= 1, ffe_buffer_len <= 1 and no smoothing, this
    % follows FFE.m semantics so MPI-reduction changes can be isolated.

    properties
        sps
        order
        e
        e_tr
        error

        len_tr
        mu_tr
        epochs_tr

        adaption_technique
        dd_mode
        mu_dd
        epochs_dd
        dd_len_fraction
        mu_dc
        e_dc

        P

        dc_buffer_len
        adaptive_mu_mode
        ffe_buffer_len
        smoothing_buffer_length
        smoothing_buffer_update

        constellation

        decide

        save_debug = 0;
        debug_struct

        optmize_mus = 0;
        mu_optimization
        mu_optimization_iter = 0;
    end

    methods
        function obj = FFE_DCremoval_adaptive_mu(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 = 0.25;

                options.mu_dc = 0.05;
                options.dc_buffer_len = 1;
                options.adaptive_mu_mode = 1;
                options.ffe_buffer_len = 1;
                options.smoothing_buffer_length = 0;
                options.smoothing_buffer_update = 0;

                options.decide = false;

                options.save_debug = 0;
                options.optmize_mus = 0;
            end

            assert(options.dc_buffer_len >= 0);
            assert(options.ffe_buffer_len >= 0);
            assert(options.smoothing_buffer_length >= 0);
            if options.smoothing_buffer_length > 0
                assert(options.smoothing_buffer_update > 0);
            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;

            obj.dc_buffer_len = floor(obj.dc_buffer_len);
            obj.ffe_buffer_len = floor(obj.ffe_buffer_len);
            obj.smoothing_buffer_length = floor(obj.smoothing_buffer_length);
            obj.smoothing_buffer_update = floor(obj.smoothing_buffer_update);
        end

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

            n = X.length;
            training = 0;
            if obj.dd_mode
                n_dd = obj.ddLength(n);
                obj.equalize(X.signal,D.signal,obj.mu_dd,obj.epochs_dd,n_dd,training,showviz);
            end
            [signal,decision] = obj.applyCurrentTaps(X.signal,n);

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

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

            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 = 0 %#ok<INUSD>
            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);
            always_ideal_decision = 0;
            grad = 0;
            weight = 0;
            update = 0;

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

            dc_buffer_enabled = obj.mu_dc ~= 0 && obj.dc_buffer_len > 1;
            adaptive_dc_enabled = dc_buffer_enabled && obj.adaptive_mu_mode;
            if dc_buffer_enabled
                e_dc_buffer = NaN(obj.dc_buffer_len,1);
            end

            ffe_buffer_enabled = ~training && obj.ffe_buffer_len > 1 && ...
                obj.adaption_technique ~= adaption_method.rls;
            if ffe_buffer_enabled
                grad_buffer = NaN(obj.order,obj.ffe_buffer_len);
            end

            if obj.smoothing_buffer_length > 0
                smoothing_buffer = zeros(1,obj.smoothing_buffer_length);
                smoothing_mean = 0;
            end

            P_err = 0;
            alpha = 0.98;
            err_prev = 0;
            gamma_dc = 1e-6;
            mu_min = 1e-6;
            mu_max = 3e-1;

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

                    if obj.smoothing_buffer_length > 0
                        smoothing_buffer = circshift(smoothing_buffer,1,2);
                        smoothing_buffer(1) = x(sample);
                        if mod(symbol,obj.smoothing_buffer_update) == 0
                            smoothing_mean = mean(smoothing_buffer);
                        end
                        x(sample:sample+obj.sps-1) = x(sample:sample+obj.sps-1) - smoothing_mean;
                    end

                    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
                        if ~always_ideal_decision
                            [~,symbol_idx] = min(abs(y(symbol) - obj.constellation));
                            d_hat(symbol,1) = obj.constellation(symbol_idx);
                        else
                            d_hat(symbol,1) = d(symbol);
                        end
                    end

                    err(symbol) = d_hat(symbol) - y(symbol); %#ok<AGROW>
                    true_err(symbol) = y(symbol) - d(symbol); %#ok<AGROW,NASGU>

                    if training || obj.dd_mode
                        switch obj.adaption_technique
                            case adaption_method.lms
                                weight = mu;
                                grad = err(symbol) * U;
                                update = grad * weight;

                            case adaption_method.nlms
                                normU = (U.'*U) + eps;
                                weight = mu / normU;
                                grad = err(symbol) * U;
                                update = grad * weight;

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

                        if ffe_buffer_enabled
                            grad_buffer = circshift(grad_buffer,1,2);
                            grad_buffer(:,1) = update;
                            if mod(symbol,obj.ffe_buffer_len) == 0
                                obj.e = obj.e + mean(grad_buffer,2,"omitnan");
                            end
                        else
                            obj.e = obj.e + update;
                        end

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

                        if obj.mu_dc ~= 0
                            if adaptive_dc_enabled
                                delta_mu = gamma_dc * err(symbol) * err_prev * (U.'*U);
                                obj.mu_dc = min(max(obj.mu_dc + delta_mu,mu_min),mu_max);
                                err_prev = err(symbol);
                                P_err = alpha*P_err + (1-alpha)*err(symbol)^2;
                                mu_dc_eff = obj.mu_dc / (P_err + eps);
                            else
                                mu_dc_eff = obj.mu_dc;
                            end

                            if dc_buffer_enabled
                                e_dc_buffer = circshift(e_dc_buffer,1);
                                e_dc_buffer(1) = obj.e_dc + mu_dc_eff * err(symbol);
                                if mod(symbol,obj.dc_buffer_len) == 0
                                    obj.e_dc = median(e_dc_buffer,"omitnan");
                                end
                            else
                                obj.e_dc = obj.e_dc + mu_dc_eff * err(symbol);
                            end
                        end
                    end

                    if obj.save_debug
                        obj.debug_struct.error(epoch,symbol) = err(symbol) * err(symbol)';

                        if training
                            obj.debug_struct.error_tr(epoch,symbol) = err(symbol) * err(symbol)';
                            obj.debug_struct.update_tr(epoch,symbol) = update.'*update ./ rms(obj.e);
                        end
                    end

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

        function [y,d_hat] = applyCurrentTaps(obj,x,N)
            x = [zeros(floor(obj.order/2),1); x; zeros(obj.order,1)];
            for sample = 1 : obj.sps : N
                symbol = (sample - 1) / obj.sps + 1;
                U = x(obj.order+sample-1:-1:sample);
                y(symbol,1) = obj.e_dc + obj.e.' * U;
                [~,symbol_idx] = min(abs(y(symbol) - obj.constellation));
                d_hat(symbol,1) = obj.constellation(symbol_idx);
            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)
            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;
            obj.mu_optimization = bayesopt(@(p)obj.muObjective(p,x,d),vars, ...
                "MaxObjectiveEvaluations",10, ...
                "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
            objective_db = 10*log10(obj.mu_optimization.MinObjective);
            if obj.dd_mode && optimize_mu_dc
                fprintf("\nFFE_DCremoval_adaptive_mu opt done: mu_tr=%9.3e, mu_dd=%9.3e, mu_dc=%9.3e, MSE=%9.3e, MSE_dB=%7.2f dB\n", ...
                    obj.mu_tr,obj.mu_dd,obj.mu_dc,obj.mu_optimization.MinObjective,objective_db);
            elseif obj.dd_mode
                fprintf("\nFFE_DCremoval_adaptive_mu opt done: mu_tr=%9.3e, mu_dd=%9.3e, MSE=%9.3e, MSE_dB=%7.2f dB\n", ...
                    obj.mu_tr,obj.mu_dd,obj.mu_optimization.MinObjective,objective_db);
            elseif optimize_mu_dc
                fprintf("\nFFE_DCremoval_adaptive_mu opt done: mu_tr=%9.3e, mu_dc=%9.3e, MSE=%9.3e, MSE_dB=%7.2f dB\n", ...
                    obj.mu_tr,obj.mu_dc,obj.mu_optimization.MinObjective,objective_db);
            else
                fprintf("\nFFE_DCremoval_adaptive_mu opt done: mu_tr=%9.3e, MSE=%9.3e, MSE_dB=%7.2f dB\n", ...
                    obj.mu_tr,obj.mu_optimization.MinObjective,objective_db);
            end
        end

        function objective = muObjective(obj,params,x,d)
            old_debug = obj.save_debug;
            old_mu_dc = obj.mu_dc;
            obj.save_debug = 1;
            if isprop(params,"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();
            obj.equalize(x,d,params.mu_tr,obj.epochs_tr,obj.len_tr,1,0);
            if obj.dd_mode
                obj.equalize(x,d,params.mu_dd,obj.epochs_dd,obj.ddLength(numel(x)),0,0);
                objective = mean(obj.debug_struct.error(end,:),"omitnan");
            else
                objective = mean(obj.debug_struct.error_tr(end,:),"omitnan");
            end
            if ~isfinite(objective)
                objective = inf;
            end
            objective_db = 10*log10(objective);
            obj.mu_optimization_iter = obj.mu_optimization_iter + 1;
            optimize_mu_dc = isprop(params,"mu_dc");
            if obj.dd_mode && optimize_mu_dc
                fprintf("\rFFE_DCremoval_adaptive_mu opt %02d: mu_tr=%9.3e, mu_dd=%9.3e, mu_dc=%9.3e, MSE=%9.3e, MSE_dB=%7.2f dB", ...
                    obj.mu_optimization_iter,params.mu_tr,params.mu_dd,params.mu_dc,objective,objective_db);
            elseif obj.dd_mode
                fprintf("\rFFE_DCremoval_adaptive_mu opt %02d: mu_tr=%9.3e, mu_dd=%9.3e, MSE=%9.3e, MSE_dB=%7.2f dB", ...
                    obj.mu_optimization_iter,params.mu_tr,params.mu_dd,objective,objective_db);
            elseif optimize_mu_dc
                fprintf("\rFFE_DCremoval_adaptive_mu opt %02d: mu_tr=%9.3e, mu_dc=%9.3e, MSE=%9.3e, MSE_dB=%7.2f dB", ...
                    obj.mu_optimization_iter,params.mu_tr,params.mu_dc,objective,objective_db);
            else
                fprintf("\rFFE_DCremoval_adaptive_mu opt %02d: mu_tr=%9.3e, MSE=%9.3e, MSE_dB=%7.2f dB", ...
                    obj.mu_optimization_iter,params.mu_tr,objective,objective_db);
            end
            obj.save_debug = old_debug;
            obj.mu_dc = old_mu_dc;
        end
    end
end