classdef EQ_silas < handle
    %EQ_SILAS FFE and DFE Equalizer Playground

    properties
        % Important Signals
        x_in %Input Sequence to be equalized
        x_length
        x_norm

        d %reference signal
        d_norm
        d_constellation %constellation points of the reference

        y_out %equalizer output signal

        % FFE coefficients always named with "e"
        Ne
        Ce %memory length FFE
        Ie1 %Indice Combination of 1nd order FFE
        Ie2 %Indice Combination of 2nd order FFE
        Ie3 %Indice Combination of 3nd order FFE
        e %coefficients for FFE

        % DFE coefficients always named with "b"
        Nb
        Cb %memory length DFE
        Ib1 %Indice Combination of 1nd order DFE
        Ib2 %Indice Combination of 2nd order DFE
        Ib3 %Indice Combination of 3nd order DFE
        b %coefficients for DFE

        error
        e_ffe
        e_dfe
        e_dc

        % coefficients
        mu_dc_train
        mu_ffe_train
        mu_dfe_train

        mu_dc_dd
        mu_ffe_dd
        mu_dfe_dd
        mu_combined_dd % [1st order FFE, 2nd order FFE, 3rd order FFE, all orders DFE]

        delay
        trainlength
        sps

        trainloops
        ddloops

        eq_parallelization_blocklength % block lengt of EQ (until now, only the dc subtraction is affected by this)
        eq_updatelatency % time in symbols until the calculated updates reach the signal again (until now, only the dc subtraction is affected by this)
        eq_avg_blocklength


    end

    methods
        function obj = EQ_silas(options)
            %EQ_SILAS Construct an instance of this class
            arguments(Input)
                options.Ne = [50 5 0] %Number of FFE coefficients (1st, 2nd and 3rd order)
                options.Nb = [30 5 3] %Number of DFE coefficients (1st, 2nd and 3rd order)
                options.trainloops = 2;
                options.trainlength = 4096;
                options.ddloops = 2;

                options.delay = 0;
                options.sps = 2;

                options.mu_dc_train = 0.01;
                options.mu_ffe_train = 0.005;
                options.mu_dfe_train = 0.005;

                options.mu_dc_dd = 0.01;
                options.mu_ffe_dd =  [0.0004 0.0005 0.0006];
                options.mu_dfe_dd =  0.0005;

                options.eq_parallelization_blocklength = 1;
                options.eq_updatelatency = 1;
                options.eq_avg_blocklength = 0;
            end

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

            % Generate helpful vectors and initialize the filters with
            % correct length:

            obj.Ce = obj.calcVNLEMemoryLength(obj.Ne);

            [obj.Ie2,obj.Ie3] = obj.calcIndiceVectors(obj.Ne);

            obj.e = zeros(sum(obj.Ce),1);


            obj.Cb = obj.calcVNLEMemoryLength(obj.Nb);

            [obj.Ib2,obj.Ib3] = obj.calcIndiceVectors(obj.Nb);

            obj.b = zeros(sum(obj.Cb),1);

        end

        function [signalclass_out] = process(obj,signalclass_in, reference_signalclass_in)

            % actual processing of the signal (steps 1. - 3.)
            % 1 normalize RMS
            signalclass_in = signalclass_in.normalize("mode","rms");

            % Process the EQ optimization
            obj.process_(signalclass_in.signal', reference_signalclass_in.signal');

            signalclass_in.signal = obj.y_out';

            %change sampling frequency of outgoing signal
            signalclass_in.fs = reference_signalclass_in.fs;

            % append to logbook
            lbdesc = ['EQ von Silas ist gelaufen '];
            signalclass_in  = signalclass_in.logbookentry(lbdesc);

            % write to output
            signalclass_out = signalclass_in;

        end

        function process_(obj,x_in,d_in)

            % 1) prepare signals
            obj.e_dc = mean(x_in);

            % 1.1) Input Signal
            obj.x_in = [zeros(1,floor(obj.Ne(1)/2)) x_in zeros(1,obj.Ne(1))];
            obj.x_length = length(x_in);
            obj.x_norm = obj.calcPowerNormalization(x_in);

            % 1.2 Reference Signal // Constellation
            obj.d = [zeros(1,obj.Nb(1)-1) d_in zeros(1,obj.Nb(1))];
            obj.d_constellation = unique(d_in);
            obj.d_norm = obj.calcPowerNormalization(d_in);

            % 1.3 Training
            obj.trainingMode();

            % 1.4 Decision Directed Mode
            obj.decisionDirectedMode();

        end

        %% Adaptive Equalization Modes

        function trainingMode(obj)

            dc_block = ones(obj.eq_parallelization_blocklength,1);

            for tloop = 1:obj.trainloops
                m = 1+obj.delay;
                dc_cnt =  0;
                for n = obj.sps*obj.delay+1:obj.sps:obj.sps*obj.trainlength
                    m = m+1;
                    dc_cnt = dc_cnt+1;

                    %get Sigal input vectors with correct length for VNLE
                    x_in_block = obj.x_in(obj.Ne(1)+n+(obj.sps-1):-1:n+obj.sps).';

                    x_in_vnle_format = obj.calcVNLENonlinVecs(x_in_block,obj.Ie2,obj.Ie3,obj.Ne,obj.x_norm);

                    %get Reference input vectors with correct length for VNLE
                    d_block = obj.d(obj.Nb(1)-obj.delay+m-2:-1:m-obj.delay-1).';
                    d_vnle_format = obj.calcVNLENonlinVecs(d_block,obj.Ib2,obj.Ib3,obj.Nb,obj.d_norm);

                    obj.e_ffe = obj.e.' * x_in_vnle_format;

                    obj.e_dfe = obj.b.' * d_vnle_format;

                    % Calculate the Error
                    obj.error = obj.e_dc + obj.e_ffe - obj.e_dfe - obj.d(obj.Nb(1)-1+m-obj.delay);


                    if obj.mu_ffe_train ~= 0
                        %update FFE coefficients with LMS
                        obj.e = obj.e - obj.error*conj(x_in_vnle_format)*obj.mu_ffe_train;
                    else
                        %update FFE coefficients with NLMS
                        obj.e = obj.e - obj.error*x_in_vnle_format/(x_in_vnle_format.'*x_in_vnle_format);
                    end

                    %update DFE coefficients with LMS
                    obj.b = obj.b + obj.mu_dfe_train*obj.error*d_vnle_format;

                    %update DC error
                    dc_block(dc_cnt) = obj.error .* obj.mu_dc_train;

                    if dc_cnt == obj.eq_parallelization_blocklength
                        obj.e_dc = obj.e_dc - mean(dc_block(dc_cnt));
                        dc_cnt = 0;
                    end

                    
                end

            end

        end


        function decisionDirectedMode(obj)

            %start the dd mode with coefficients from training
            coeff = [obj.e;obj.b];
            obj.e_dc = ones(obj.eq_updatelatency,1);%.*obj.e_dc;
            dc_block = ones(obj.eq_parallelization_blocklength,1);

            for ddloop = 1:obj.ddloops

                m = 0;
                dc_cnt = 0;

                mu_mat = diag([ones(1,obj.Ce(1))*obj.mu_ffe_dd(1)... %1st order ffe
                    ones(1,obj.Ce(2))*obj.mu_ffe_dd(2)... %2nd order ffe
                    ones(1,obj.Ce(3))*obj.mu_ffe_dd(3)... %3rd order ffe
                    ones(1,sum(obj.Cb))*obj.mu_dfe_dd]); %all order dfe

                mu_ffe = [ones(1,obj.Ce(1))*obj.mu_ffe_dd(1)... %1st order ffe
                    ones(1,obj.Ce(2))*obj.mu_ffe_dd(2)... %2nd order ffe
                    ones(1,obj.Ce(3))*obj.mu_ffe_dd(3)];

                mu_dfe = ones(1,sum(obj.Cb))*obj.mu_dfe_dd;

                y               = zeros(1,floor(obj.x_length/obj.sps));
                d_feedback      = zeros(obj.Cb(1),1);
                d_vnle          = obj.calcVNLENonlinVecs(d_feedback,obj.Ib2,obj.Ib3,obj.Nb,obj.d_norm);
                d_hat           = zeros(obj.x_length,1);

                m_reg = 0;

                if obj.eq_avg_blocklength > 0
                    averaging_window = zeros(obj.eq_avg_blocklength,1);
                end

                for k = 1:obj.sps:obj.x_length
                    dc_cnt = dc_cnt+1;
                    m=m+1;

                    %get Sigal input vectors with correct length for VNLE
                    x = obj.x_in(obj.Ne(1)+k-1:-1:k).';
% 
                    if obj.eq_avg_blocklength > 0 %% Das läuft gut mit 400er Fenster!!
                        averaging_window = circshift(averaging_window,obj.sps);
                        averaging_window(1:obj.sps,1) = x(1:obj.sps);
                        avg_(k) = mean(averaging_window);
                        x = x-avg_(k);
                    end

                    %bring this signal to "special" VNLE format
                    x_vnle = obj.calcVNLENonlinVecs(x,obj.Ie2,obj.Ie3,obj.Ne,obj.x_norm);

                    %combine FFE with DFE to one vector (cursor between the two sequences)
                    x_d = [x_vnle;-d_vnle];

                    %Apply filter
                    %y(m) = (m_reg(end)*dc_cnt + obj.e_dc(end)) + x_d.'* coeff;
                    if obj.mu_dc_dd > 0
                        y(m) =  obj.e_dc(end) + x_d.'* coeff;
                    else
                        y(m) =  x_d.'* coeff;
                    end

                    %Decision
                    [~,symbol_idx] = min(abs(y(m) - obj.d_constellation)); % decision for closest constellation point
                    d_hat(k) = obj.d_constellation(symbol_idx);

                    %Error between FFE & DFE filtered signal and Decision
                    obj.error(k) = y(m) - d_hat(k);

                    %Update FFE and DFE coefficients
                    coeff = coeff - (mu_mat * (obj.error(k) * conj(x_d)));

                    %Update DC error
                    dc_block(dc_cnt) = obj.error(k) ;

                    if dc_cnt == obj.eq_parallelization_blocklength
                        if obj.eq_updatelatency > 1

                            obj.e_dc = circshift(obj.e_dc,1);

%                             m_reg(end+1) = ((1:obj.eq_parallelization_blocklength)' \ (cumsum(dc_block)));
% 
%                             obj.e_dc(1) = obj.e_dc(2) - sign(m_reg(end)) .* (sum(dc_block).* m_reg(end) .* obj.mu_dc_dd);
                            obj.e_dc(1) = obj.e_dc(2) - sum(dc_block) .* obj.mu_dc_dd;

                        else
                            %m_reg(end+1) = ((1:obj.eq_parallelization_blocklength)' \ (cumsum(dc_block)));

%                             obj.e_dc = obj.e_dc - sign(m_reg(end)) .* (sum(dc_block).* m_reg(end) .* obj.mu_dc_dd);

                             obj.e_dc = obj.e_dc - sum(dc_block) .* obj.mu_dc_dd;
                            
%                             obj.e_dc = obj.e_dc - obj.mu_dc_dd * obj.error(k); %newapril
                        end

                        dc_cnt = 0;

                    end


                    % Append new decision to decision feedback
                    if obj.Nb(1) > 0

                        %shift up one index
                        d_feedback(2:end) = d_feedback(1:end-1);
                        %replace 1st index with current estimation
                        d_feedback(1) = d_hat(k);
                        %build memorylike VNLE version
                        d_vnle = obj.calcVNLENonlinVecs(d_feedback,obj.Ib2,obj.Ib3,obj.Nb,obj.d_norm);

                    end

                end
            end

            obj.y_out = (circshift( y.'  ,-(obj.delay))).';

        end

        %% Functions needed During Adaption
        function x_in_vnle_format = calcVNLENonlinVecs(~,x_in_block,I_2,I_3,N_,norm_)
            % These are the second and third order input signal products of the VNLE EQ
            % ∑ h1 x_in(k-n1) + ∑∑ h2 x_in(k-n1)*x_in(k-n2) + ∑∑∑ h3 x_in(k-n1)*x_in(k-n2)*x_in(k-n3)
            l1=length(x_in_block);
            l2=length(I_2);
            l3=length(I_3);
            final_length = l1+l2+l3;

            x_in_vnle_format = zeros(final_length,1);

            idx = l1;
            x_in_vnle_format(1:idx) = x_in_block;

            if N_(2) > 0
                delta_2 = round((N_(1)-N_(2)) / 2);
                input_vec_se = x_in_block(delta_2:end) / norm_(2); %TODO normalization step

                % Extract columns from I_2
                col1 = input_vec_se(I_2(:,1));
                col2 = input_vec_se(I_2(:,2));

                x2 = col1 .* col2;
                x_in_vnle_format(idx+1:idx+l2) = x2;
            end

            if N_(3) > 0
                delta_3 = round((N_(1)-N_(3))/2);
                input_vec_th = x_in_block(delta_3:end) / norm_(3);

                % Extract columns from I_3
                col1 = input_vec_th(I_3(:,1));
                col2 = input_vec_th(I_3(:,2));
                col3 = input_vec_th(I_3(:,3));

                % Perform matrix multiplication
                x3 = col1 .* col2 .* col3;

                idx = idx+l2;
                x_in_vnle_format(idx+1:idx+l3) = x3;
            end

        end


        %% Functions needed for Preparation
        function [C] = calcVNLEMemoryLength(~,N)

            %calculates the memory length of VNLE
            C = zeros(size(N));

            for o = 1:numel(N)
                switch o
                    case 1
                        C(o) = N(o);
                    case 2
                        C(o) = N(o)*(N(o)+1) / 2;
                    case 3
                        C(o) = N(o)*(N(o)+1)*(N(o)+2) / 6;
                end
            end

        end

        function [indvec2nd, indvec3rd] = calcIndiceVectors(~,N)

            % Init vectors of 2nd and 3rd order coefficient indices ->
            % yield combination with

            for order = 2:numel(N)
                n = N(order);
                v = 1:n;  % Ursprünglicher Vektor
                row = 1;

                % Schleifen zur Generierung des Indize Vektors
                switch order

                    case 2

                        indvec2nd = zeros(n*(n+1)/2, order);
                        for i = 1:n
                            for j = i:n
                                indvec2nd(row, :) = [v(i) v(j)];
                                row = row + 1;
                            end
                        end

                    case 3

                        indvec3rd = zeros(n*(n+1)*(n+2)/6, 3);
                        for i = 1:n
                            for j = i:n
                                for k = j:n
                                    indvec3rd(row, :) = [v(i) v(j) v(k)];
                                    row = row + 1;
                                end
                            end
                        end
                end
            end

        end

        function powerNorm  = calcPowerNormalization(~,v)

            powerNorm(1) = sqrt(mean(abs(v   ).^2));
            powerNorm(2) = sqrt(mean(abs(v.^2).^2));
            powerNorm(3) = sqrt(mean(abs(v.^3).^2));

        end


    end
end