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
        e_dc
        error

        len_tr
        mu_tr
        epochs_tr

        mu_dd
        epochs_dd
        mu_dc

        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 = 3020;

        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.mu_dc = 0;

                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 = 3020;

            end

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

            
            obj.error = 0;
            obj.e_dc = 0;

        end

        function [X,N] = 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);
            obj.e_dc = 0;

            if obj.optmize_mus
                obj.optimizeMus(X.signal,D.signal);
                obj.e = zeros(sum(obj.ce),1);
                obj.e_dc = 0;
            end

            % 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

            N = 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

            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 = x(:);
            d = d(:);
            x = [zeros(floor(obj.order(1)/2),1); x; zeros(obj.order(1),1)];
            n_symbols = floor(N / obj.sps);
            y = zeros(n_symbols,1);
            d_hat = zeros(n_symbols,1);

            if showviz
                figure(111);
                subplot(2,2,1:2);
                hold on
                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_dc + 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 * 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 obj.mu_dc ~= 0
                        obj.e_dc = obj.e_dc - obj.mu_dc * err;
                    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
                    if obj.save_debug
                        obj.debug_struct.error(epoch,symbol) = err * err';
                        if training
                            obj.debug_struct.error_tr(epoch,symbol) = err * err';
                        end
                    end

                end
            end

        end

        function optimizeMus(obj,x,d)
            mu_range = [1e-5, 1e-2];
            mu_dc_range = [1e-5, 1e-1];

            [x_opt,d_opt,N_opt] = obj.optimizationSignals(x,d);

            vars = obj.muOptimizableVariables("mu_tr",mu_range);
            vars = [vars, obj.muOptimizableVariables("mu_dd",mu_range)];
            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;
            fprintf("VNLE 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.muVectorFromParams(obj.mu_optimization.XAtMinObjective,"mu_tr");
            obj.mu_dd = obj.muVectorFromParams(obj.mu_optimization.XAtMinObjective,"mu_dd");
            if optimize_mu_dc
                obj.mu_dc = obj.mu_optimization.XAtMinObjective.mu_dc;
            end

            if optimize_mu_dc
                fprintf("\nVNLE mu opt done: mu_tr=[%s], mu_dd=[%s], mu_dc=%9.3e, BER=%9.3e\n", ...
                    obj.formatMuVector(obj.mu_tr),obj.formatMuVector(obj.mu_dd), ...
                    obj.mu_dc,obj.mu_optimization.MinObjective);
            else
                fprintf("\nVNLE mu opt done: mu_tr=[%s], mu_dd=[%s], BER=%9.3e\n", ...
                    obj.formatMuVector(obj.mu_tr),obj.formatMuVector(obj.mu_dd), ...
                    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;
            optimize_mu_dc = ismember("mu_dc",string(params.Properties.VariableNames));
            if optimize_mu_dc
                obj.mu_dc = params.mu_dc;
            end

            obj.e = zeros(sum(obj.ce),1);
            obj.e_dc = 0;
            obj.debug_struct = struct();
            muTrCandidate = obj.muVectorFromParams(params,"mu_tr");
            muDdCandidate = obj.muVectorFromParams(params,"mu_dd");
            N_tr = min(obj.len_tr,numel(x));
            obj.equalize(x,d,muTrCandidate,obj.epochs_tr,N_tr,1,0);
            [signal,~] = obj.equalize(x,d,muDdCandidate,obj.epochs_dd,numel(x),0,0);

            [ber,errors] = obj.berObjective(signal,d);
            objective = ber;
            if ~isfinite(objective)
                objective = inf;
            end

            obj.mu_optimization_iter = obj.mu_optimization_iter + 1;
            if optimize_mu_dc
                fprintf("\rVNLE mu opt %02d: mu_tr=[%s], mu_dd=[%s], mu_dc=%9.3e, BER=%9.3e, errors=%d", ...
                    obj.mu_optimization_iter,obj.formatMuVector(muTrCandidate), ...
                    obj.formatMuVector(muDdCandidate),params.mu_dc,ber,errors);
            else
                fprintf("\rVNLE mu opt %02d: mu_tr=[%s], mu_dd=[%s], BER=%9.3e, errors=%d", ...
                    obj.mu_optimization_iter,obj.formatMuVector(muTrCandidate), ...
                    obj.formatMuVector(muDdCandidate),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);
            skip_front = min(1000,max(0,floor(numel(rx_bits.signal) / 4)));
            [~,errors,ber,~] = calc_ber(rx_bits.signal,tx_bits.signal, ...
                "skip_front",skip_front, ...
                "skip_end",0, ...
                "returnErrorLocation",1);
        end

        function vars = muOptimizableVariables(obj,prefix,mu_range)
            vars = optimizableVariable.empty;
            activeOrders = find(obj.order > 0);

            for idx = 1:numel(activeOrders)
                orderIdx = activeOrders(idx);
                varName = sprintf("%s_%d",prefix,orderIdx);
                vars = [vars, optimizableVariable(varName,mu_range,"Transform","log")]; %#ok<AGROW>
            end
        end

        function mu = muVectorFromParams(obj,params,prefix)
            mu = zeros(1,3);
            for orderIdx = 1:numel(mu)
                varName = sprintf("%s_%d",prefix,orderIdx);
                if ismember(varName,string(params.Properties.VariableNames))
                    mu(orderIdx) = params.(varName);
                elseif numel(obj.(char(prefix))) >= orderIdx
                    mu(orderIdx) = obj.(char(prefix))(orderIdx);
                else
                    mu(orderIdx) = obj.(char(prefix))(1);
                end
            end
        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,"VNLE 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_dc = obj.e_dc;
            state.error = obj.error;
            state.mu_dc = obj.mu_dc;
            state.save_debug = obj.save_debug;
            state.debug_struct = obj.debug_struct;
        end

        function restoreObjectiveState(obj,state)
            obj.e = state.e;
            obj.e_dc = state.e_dc;
            obj.error = state.error;
            obj.mu_dc = state.mu_dc;
            obj.save_debug = state.save_debug;
            obj.debug_struct = state.debug_struct;
        end

        function s = formatMuVector(~,mu)
            s = strtrim(sprintf("%9.3e ",mu));
        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)
            x_in_block = x_in_block(:);
            l1=length(x_in_block);
            l2=size(I_2,1);
            l3=size(I_3,1);
            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)
            v = 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