classdef MLSE_viterbi < handle
    %MLSE calculates the most probable sequence for an input signal with given/ known channel impulse response of any length

    properties(Access=public)
        M %PAM-M
        DIR
        trellis_states
        duobinary_output
    end

    methods (Access=public)

        function obj = MLSE_viterbi(options)
            %NAME Construct an instance of this class
            %   Detailed explanation goes here

            arguments
                options.M double = 4;
                options.DIR double = [1];
                options.trellis_states double = [-3 -1 1 3];
                options.duobinary_output logical = false;

            end

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

            % do more stuff

        end

        function signalclass = process(obj,signalclass)

            data_in =   signalclass.signal;

            data_out = obj.process_(data_in);

            signalclass.signal = data_out;

        end

        function data_out = process_(obj,data_in)


            % remove unnecessary zeros at start of impulse response to keep
            % number of trellis states minimal
            DIR_nonzero = find(obj.DIR ~= 0);
            if DIR_nonzero(1) > 1
                obj.DIR(1:DIR_nonzero(1)-1) = [];
            end

            if isscalar(obj.DIR)
                obj.DIR = [0 obj.DIR];
            end


            %%%%%%%%%%%%%%%%%%%%%%%%%%%% WORKING
            % impulse respnse i.e. [0.5, 1.0000]
            obj.DIR = flip(obj.DIR);

            % RMS normalization of input data
            data_in = data_in ./ rms(data_in);

            % seems to be the only way to use combvec for a flexible amount
            % of vectors. 'combs' contains all trellis states
            pre_comb_mat = repmat(obj.trellis_states,length(obj.DIR)-1,1);
            pre_comb_cell = mat2cell(pre_comb_mat,ones(1,size(pre_comb_mat,1)),size(pre_comb_mat,2));
            combs = fliplr(combvec(pre_comb_cell{:}).');

            % Save first and last symbol of each state
            first_sym = combs(:,1);
            last_sym = combs(:,end);
            states = sum(combs,2);

            % Calculate all possible input symbols for the desired impulse
            % response. Row number is the index of the previous state,
            % column number is the index of the next state
            % noise free received == branch metrics
            noise_free_received = zeros(length(states),length(states));
            count_row = 1;
            count_col = 1;
            for l1 = 1:length(states)
                for l2 = 1:length(states)
                    if sum(combs(l2,2:end) == combs(l1,1:end-1)) == size(combs,2)-1
                        noise_free_received(count_row,count_col) = sum(combs(l2,:).*obj.DIR(end:-1:2)) + last_sym(l1)*obj.DIR(1);
                    else
                        noise_free_received(count_row,count_col) = inf;
                    end
                    count_row = count_row + 1;
                end
                count_col = count_col + 1;
                count_row = 1;
            end

            % match amplitude levels of input signal to those of the calculated ideal symbols
            % i.e. match the rms values of data_in to noise_free_received
            if isreal(data_in)
                if obj.M == round(obj.M)
                    data_in = data_in * rms(noise_free_received(noise_free_received ~= inf),'all','omitnan');
                end
            end

            % OLD
            % initilaize the output vector
            data_out = NaN(size(data_in));

            sum_path_metrics = zeros(length(states),length(states));

            % first trellis path
            % euclidian distance as path metric
            path_metrics = (abs(repmat(data_in(1),size(noise_free_received)) - noise_free_received)).^2;
            sum_path_metrics = sum_path_metrics + path_metrics;                                     % calculation of all possible sum path metrics
            [sum_path_metrics_res(:,1),path_idx(:,1)] = min(sum_path_metrics,[],2);                 % find the best path to each state, store sum path metric and predecessor for each state

            % remaining trellis paths
            for n = 2:length(data_in)
                sum_path_metrics = repmat(sum_path_metrics_res(:,n-1).',length(states),1);
                %                 path_metrics = (repmat(data_in(n),size(noise_free_received)) - noise_free_received).^2;%.*prob_mat;
                path_metrics = (abs(repmat(data_in(n),size(noise_free_received)) - noise_free_received)).^2;
                sum_path_metrics = sum_path_metrics + path_metrics;
                [sum_path_metrics_res(:,n),path_idx(:,n)] = min(sum_path_metrics,[],2);
            end

            %% trace back
            ideal_path = NaN(1,length(data_in)+1);
            % find ideal trellis path by going through the trellis
            % backwards
            [~,ideal_path(length(data_in)+1)] = min(sum_path_metrics_res(:,length(data_in)));

            % starting with the state that has the lowest sum path
            % metric, follow the stored information about the
            % predecessor
            for h = length(data_in):-1:1
                ideal_path(h) = path_idx(ideal_path(h+1),h);
            end

            idx_out = ideal_path(2:length(data_in)+1);
            %%%%%%%%%%%%%%%%%%%%%%%%%%%% WORKING




            if obj.duobinary_output
                %use duobinary encoder, output is already scaled inside this
                %one
                data_out = Duobinary().encode(first_sym(idx_out));

            else
                %
                data_out(1:length(data_in)) = first_sym(idx_out);
                % scale to rms = 1 using the standard sqrt() expressions
                if obj.M == 4
                    data_out = data_out./sqrt(5);
                elseif obj.M == 6
                    data_out = data_out./sqrt(10);
                elseif obj.M == 8
                    data_out = data_out./sqrt(21);
                end

            end


        end
    end

    methods (Access=private)
        % Cant be seen from outside! So put all your functions here that can/
        % shall not be called from outside


    end
end