classdef VNLE < handle
    % Implementation of plain and simple FFE.
    % 1) Training mode (stable performance when you use NLMS)
    % 2) Decision directed mode

    %  Eq = VNLE("epochs_tr",5,"epochs_dd",5,"len_tr",4096*2,"mu_dd",[0.0004 0.0005 0.0006],"mu_tr",0,"order",[25,2,2],"sps",2,"decide",1);
    % Somehow it is not possible to use only 1 nonlinear order

    properties
        sps % usually 2
        order
        e
        error

        len_tr
        mu_tr
        epochs_tr

        mu_dd
        epochs_dd

        constellation

        decide

        x_norm
        ce
        ie2
        ie3
    end

    methods
        function obj = VNLE(options)
            arguments(Input)

                options.sps = 2;
                options.order = [15,2,2];

                options.len_tr = 4096;
                options.mu_tr = 0;
                options.epochs_tr = 5;

                options.mu_dd = 1e-5;
                options.epochs_dd = 5;

                options.decide = false;

            end

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

            
            obj.error = 0;

        end

        function [X] = process(obj, X, D)

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

            obj.constellation   = unique(D.signal);
            obj.x_norm          = obj.calcPowerNormalization(X.signal);
            obj.ce              = obj.calcVNLEMemoryLength(obj.order);
            [obj.ie2,obj.ie3]   = obj.calcIndiceVectors(obj.order);

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

            % Training Mode
            training = 1;
            showviz = 0;
            obj.equalize(X.signal, D.signal,obj.mu_tr,obj.epochs_tr,obj.len_tr,training,showviz);

            % Decision Directed Mode
            N = X.length;
            training = 0;
            showviz = 0;
            [signal,decision]=obj.equalize(X.signal, D.signal,obj.mu_dd,obj.epochs_dd,N,training,showviz);

            % Output Signal
            if obj.decide
                X.signal = decision;
            else
                X.signal = signal;
            end
            X.fs = D.fs; %change sampling frequency of outgoing signal from fdac e.g. 2 sps to symbol spaced = fsym
            lbdesc = [num2str(obj.order),' tap FFE'];
            X  = X.logbookentry(lbdesc); % append to logbook


        end

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

            arguments
                obj
                x
                d
                mu
                epochs
                N
                training
                showviz
            end

            if all(mu == mu(1))
                % mu = mu(1);
                mu = diag(ones(1,sum(obj.ce))*mu(1));
            else
                mu = diag([ones(1,obj.ce(1))*mu(1) ...
                           ones(1,obj.ce(2))*mu(2) ...
                           ones(1,obj.ce(3))*mu(3) ]);
            end

            x = [zeros(floor(obj.order(1)/2),1); x; zeros(obj.order(1),1)];

            if showviz
                f = figure(111);
                subplot(2,2,1:2);
                hold on
                a = scatter(1:numel(x),x,1,'.');
                a2 = scatter(1,1,1,'.');
                a3 = scatter(1,1,2,'.');
                a4 = xline(1);
                ylim([-3 3])
                xlim([0 length(x)]);
                subplot(2,2,3:4)
                c = stem(obj.e);
                ylim([-1 1])
                drawnow
            end

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

                    symbol = symbol+1;
                    
                    
                    % x_in = x(obj.order(1)+sample+(obj.sps-1):-1:sample+obj.sps);
                    x_in = x(obj.order(1)+sample-1:-1:sample);
                    x_in = obj.calcVNLENonlinVecs(x_in,obj.ie2,obj.ie3,obj.order,obj.x_norm); 

                    y(symbol,1) = obj.e.' * x_in;    % Calculating output of LMS  __ * |

                    if training
                        err =  y(symbol) - d(symbol);   % Instantaneous error
                    else
                        [~,symbol_idx] = min(abs(y(symbol) - obj.constellation)); % decision for closest constellation point
                        d_hat(symbol,1) = obj.constellation(symbol_idx);
                        err =  y(symbol) - d_hat(symbol);   % Instantaneous error
                    end

                    if ~all(mu==0,'all') %mu has not only zeros
                        % obj.e = obj.e - (mu * err * x_in) ;    % Weight update rule of LMS
                        obj.e = obj.e - ( (mu * x_in) * err ) ;    % Weight update rule of LMS
                    else
                        normalizationfactor =  (x_in.' * x_in);
                        obj.e = obj.e - err * x_in / normalizationfactor;    % Weight update rule of NLMS
                    end

                    if mod(sample,100) == 1 && showviz
                        a2.XData = 1:2*numel(y);
                        a2.YData = repelem(y, 2);
                        a3.XData = 1:2*numel(d_hat);
                        a3.YData = repelem(d_hat, 2);
                        a4.Value = sample;
                        % b.YData = x(symbol:symbol+500);
                        c.YData = obj.e;
                        drawnow;
                    end

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

                end
            end

        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 
             indvec2nd=[];
             indvec3rd=[];
            for o = 2:numel(N) 
                n = N(o); 
                v = 1:n;  % Ursprünglicher Vektor 
                row = 1; 
 
                % Schleifen zur Generierung des Indize Vektors 
                switch o 
 
                    case 2 
 
                        indvec2nd = zeros(n*(n+1)/2, o); 
                        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