classdef FFE_plain < handle
    % Plain feed-forward equalizer with training, DD mode, and BER mu tuning.
    %
    % This class intentionally contains only the core FFE tap adaptation logic.
    % MPI-specific DC tracking, A1/A2 suppression, and delayed update buffers
    % stay in FFE and the MPI_reduction classes.

    properties
        sps
        order
        e
        e_tr
        error

        len_tr
        mu_tr
        epochs_tr

        adaption_technique
        dd_mode
        mu_dd
        epochs_dd
        dd_len_fraction

        P
        constellation

        decide

        save_debug = 0
        debug_struct

        optmize_mus = 0
        mu_optimization
        mu_optimization_iter = 0
        mu_optimization_len
        plot_mu_optimization = 0
        mu_optimization_fignum = 3010
    end

    methods
        function obj = FFE_plain(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.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

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

            obj.sps = floor(obj.sps);
            obj.order = floor(obj.order);
            obj.dd_len_fraction = min(max(obj.dd_len_fraction,0),1);
            obj.resetState();
        end

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

            if obj.optmize_mus
                obj.optimizeMus(X.signal,D.signal);
                obj.resetState();
            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;

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

            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

            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();

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

            debug_enabled = obj.save_debug;
            if debug_enabled
                obj.initializeDebug(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;

                    U = x(obj.order+sample-1:-1:sample);
                    y(symbol,1) = (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);

                    if mu ~= 0 && (training || obj.dd_mode)
                        switch adaption_code
                            case 1
                                weight = mu / ((U.'*U) + eps);
                                grad = err * U;
                                update = grad * weight;
                                obj.e = obj.e + update;

                            case 2
                                weight = mu;
                                grad = err * U;
                                update = grad * weight;
                                obj.e = obj.e + update;

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

                    if debug_enabled && epoch == 1
                        obj.debug_struct.error_first_epoch(1,symbol) = err * err';
                    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;

                        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

        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_plain: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

            obj.mu_optimization_iter = 0;
            fprintf("FFE_plain 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;
                fprintf("\nFFE_plain 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);
            else
                fprintf("\nFFE_plain 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.captureObjectiveState();
            cleanup = onCleanup(@()obj.restoreObjectiveState(old_state));

            obj.save_debug = 0;
            obj.resetState();
            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;
            if obj.dd_mode
                fprintf("\rFFE_plain 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);
            else
                fprintf("\rFFE_plain 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 [x_opt,d_opt,N_opt] = optimizationSignals(obj,x,d,opt_len)
            if nargin < 4
                opt_len = obj.mu_optimization_len;
            end

            N_available = min(numel(x),numel(d) * obj.sps);

            if isempty(opt_len) || opt_len <= 0 || isinf(opt_len)
                N_opt = N_available;
            else
                N_opt = min(N_available,max(obj.len_tr,opt_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 [ber,errors] = berObjective(~,signal,d)
            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",1000, ...
                "skip_end",0, ...
                "returnErrorLocation",1);
        end

        function plotMuOptimization(obj)
            if isempty(obj.mu_optimization)
                return
            end

            X = obj.mu_optimization.XTrace;
            objective = obj.mu_optimization.ObjectiveTrace;
            objective = objective(:);
            valid = isfinite(objective);

            if isempty(X) || ~any(valid)
                return
            end

            var_names = X.Properties.VariableNames;
            n_vars = numel(var_names);
            eval_idx = (1:numel(objective)).';
            objective_plot = obj.positiveObjectiveForLogPlot(objective);
            best_plot = obj.positiveObjectiveForLogPlot(cummin(objective));

            figure(obj.mu_optimization_fignum);
            clf;
            t = tiledlayout(2,2,"TileSpacing","compact","Padding","compact");
            title(t,"FFE_plain Bayesian mu optimization");

            nexttile;
            h_candidate = semilogy(eval_idx,objective_plot,"o-","DisplayName","candidate");
            obj.addOptimizationDataTips(h_candidate,X,objective,objective_plot,eval_idx,var_names);
            hold on;
            h_best_trace = semilogy(eval_idx,best_plot,"k-","LineWidth",1.2,"DisplayName","best so far");
            h_best_trace.DataTipTemplate.DataTipRows(end+1) = dataTipTextRow("Evaluation",eval_idx);
            h_best_trace.DataTipTemplate.DataTipRows(end+1) = dataTipTextRow("Best BER",best_plot);
            grid on;
            xlabel("Evaluation");
            ylabel("BER");
            legend("Location","best");

            if n_vars < 2
                return
            end

            pairs = nchoosek(1:n_vars,2);
            n_pair_plots = min(size(pairs,1),3);
            [~,best_idx] = min(objective);

            for pair_idx = 1:n_pair_plots
                nexttile;
                x_name = var_names{pairs(pair_idx,1)};
                y_name = var_names{pairs(pair_idx,2)};
                x_data = X.(x_name);
                y_data = X.(y_name);
                c_data = log10(objective_plot);

                h_scatter = scatter(log10(x_data),log10(y_data),35,c_data,"filled");
                obj.addOptimizationDataTips(h_scatter,X,objective,objective_plot,eval_idx,var_names);
                hold on;
                h_best = plot(log10(x_data(best_idx)),log10(y_data(best_idx)),"kp", ...
                    "MarkerSize",12, ...
                    "MarkerFaceColor","y", ...
                    "DisplayName","best");
                obj.addOptimizationDataTips(h_best,X(best_idx,:),objective(best_idx),objective_plot(best_idx),eval_idx(best_idx),var_names);
                grid on;
                xlabel("log10(" + string(x_name) + ")");
                ylabel("log10(" + string(y_name) + ")");
                cb = colorbar;
                cb.Label.String = "log10(BER)";
                title(string(x_name) + " vs " + string(y_name));
            end
        end

        function addOptimizationDataTips(~,plot_handle,X,objective,objective_plot,eval_idx,var_names)
            plot_handle.DataTipTemplate.DataTipRows(end+1) = dataTipTextRow("Evaluation",eval_idx);
            plot_handle.DataTipTemplate.DataTipRows(end+1) = dataTipTextRow("BER",objective);
            plot_handle.DataTipTemplate.DataTipRows(end+1) = dataTipTextRow("BER shown",objective_plot);

            for var_idx = 1:numel(var_names)
                var_name = var_names{var_idx};
                plot_handle.DataTipTemplate.DataTipRows(end+1) = dataTipTextRow(var_name,X.(var_name));
            end
        end

        function objective_plot = positiveObjectiveForLogPlot(~,objective)
            objective_plot = objective;
            positive_values = objective(isfinite(objective) & objective > 0);

            if isempty(positive_values)
                floor_value = 1e-12;
            else
                floor_value = min(positive_values) / 10;
            end

            objective_plot(~isfinite(objective_plot) | objective_plot <= 0) = floor_value;
        end

        function state = captureObjectiveState(obj)
            state.e = obj.e;
            state.e_tr = obj.e_tr;
            state.error = obj.error;
            state.P = obj.P;
            state.save_debug = obj.save_debug;
            state.debug_struct = obj.debug_struct;
        end

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

    methods (Access = private)
        function resetState(obj)
            obj.e = zeros(obj.order,1);
            obj.e_tr = zeros(obj.order,1);
            obj.error = 0;
            obj.P = (1/0.05) * eye(obj.order);
            obj.debug_struct = struct();
        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_plain:InvalidAdaptionTechnique", ...
                    "Unsupported FFE adaption technique.");
            end
        end

        function initializeDebug(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);

            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
    end
end