classdef FFE_A1 < handle
    %FFE_A1 Feed-forward equalizer with A1 moving-average input suppression.
    %
    % This class is intentionally narrow: it contains the normal adaptive FFE
    % flow plus the A1 block-rate moving-average input correction. A2,
    % DC-tracking, and optimizer logic are kept out of this file.

    properties
        sps
        order
        e
        e_tr
        error

        len_tr
        mu_tr
        epochs_tr

        adaption_technique
        dd_mode
        mu_dd
        epochs_dd
        dd_len_fraction

        % A1 moving-average input suppression
        dc_avg_bufferlength_a1
        dc_avg_update_blocklength_a1
        dc_smoothing_a1

        P
        constellation

        decide

        save_debug = 0
        debug_struct
    end

    methods
        function obj = FFE_A1(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_avg_bufferlength_a1 = 0
                options.dc_avg_update_blocklength_a1 = 0
                options.dc_smoothing_a1 = 0

                options.decide = false
                options.save_debug = 0
            end

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

            assert(obj.dc_avg_bufferlength_a1 >= 0);
            assert(obj.dc_avg_update_blocklength_a1 >= 0);

            obj.e = zeros(obj.order,1);
            obj.error = 0;
            obj.dc_avg_bufferlength_a1 = floor(obj.dc_avg_bufferlength_a1);
            obj.dc_avg_update_blocklength_a1 = floor(obj.dc_avg_update_blocklength_a1);
            obj.dc_smoothing_a1 = min(max(obj.dc_smoothing_a1,0),1);
        end

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

            obj.constellation = unique(D.signal);

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

            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 A1']);

            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 %#ok<INUSA>
            end

            x = [zeros(floor(obj.order/2),1); x; zeros(obj.order,1)];
            lambda = mu;
            n_symbols = ceil(N / obj.sps);
            y = zeros(n_symbols,1);
            d_hat = zeros(n_symbols,1);
            err = zeros(n_symbols,1);
            constellation = obj.constellation;

            adaption_code = obj.adaptionCode();
            adaption_is_rls = adaption_code == 3;

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

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

            dc_avg_enabled = obj.dc_avg_bufferlength_a1 > 1;
            if dc_avg_enabled
                dc_avg_buffer = NaN(obj.dc_avg_bufferlength_a1,1);
                dc_avg_buffer_pos = 0;
                dc_avg_sum = 0;
                dc_avg_valid_count = 0;
                dc_avg_est = 0;
                dc_avg_update_blocklength = obj.dc_avg_update_blocklength_a1;
                if dc_avg_update_blocklength <= 0
                    dc_avg_update_blocklength = obj.dc_avg_bufferlength_a1;
                end
                dc_avg_update_blocklength = max(1,floor(dc_avg_update_blocklength));
                dc_avg_input_offset = floor(obj.order/2);
                % Hardware-like A1 is causal; dc_smoothing_a1 is reserved for offline variants.
            end

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

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

                    U = x(obj.order+sample-1:-1:sample);
                    if dc_avg_enabled
                        dc_avg_offset = dc_avg_est;
                        U = U - dc_avg_offset;
                    end

                    y(symbol,1) = (obj.e.*mask).' * U;

                    if dc_avg_enabled
                        raw_sample_idx = dc_avg_input_offset + sample;
                        [dc_avg_buffer,dc_avg_buffer_pos,dc_avg_sum,dc_avg_valid_count,dc_avg_est] = ...
                            obj.updateA1Buffer(dc_avg_buffer,dc_avg_buffer_pos,dc_avg_sum, ...
                            dc_avg_valid_count,dc_avg_est,x(raw_sample_idx),symbol, ...
                            dc_avg_update_blocklength);
                    end

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

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

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

                            case 2
                                weight = mu;
                                grad = err(symbol) * 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(symbol);
                                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(symbol) * err(symbol)';
                    end

                    if debug_enabled && epoch == epochs
                        error_power = err(symbol) * err(symbol)';
                        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_avg_offset(1,symbol) = dc_avg_offset;

                        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

            if ~adaption_is_rls
                obj.P = [];
            end
        end
    end

    methods (Access = private)
        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_A1:InvalidAdaptionTechnique", ...
                    "Unsupported FFE adaption technique.");
            end
        end
    end

    methods (Static, Access = private)
        function [buffer,buffer_pos,buffer_sum,valid_count,estimate] = ...
                updateA1Buffer(buffer,buffer_pos,buffer_sum,valid_count,estimate, ...
                new_value,symbol,update_blocklength)

            buffer_pos = buffer_pos + 1;
            if buffer_pos > numel(buffer)
                buffer_pos = 1;
            end

            old_value = buffer(buffer_pos);
            if isfinite(old_value)
                buffer_sum = buffer_sum - old_value;
                valid_count = valid_count - 1;
            end

            if isfinite(new_value)
                buffer(buffer_pos) = new_value;
                buffer_sum = buffer_sum + new_value;
                valid_count = valid_count + 1;
            else
                buffer(buffer_pos) = NaN;
            end

            if mod(symbol,update_blocklength) == 0
                if valid_count == 0
                    estimate = 0;
                else
                    estimate = buffer_sum / valid_count;
                end
            end
        end
    end
end