function [GMI,NGMI] = calc_ngmi(test_signal,reference_signal,options)
% Silas implementation of (N)GMI calculation according to: J. Cho, L. Schmalen, und P. J. Winzer,
% „Normalized Generalized Mutual Information as a Forward Error Correction Threshold for Probabilistically Shaped QAM“,
% in 2017 European Conference on Optical Communication (ECOC), Sep. 2017, doi: 10.1109/ECOC.2017.8345872.

% This implementation assumes the same normal distributed noise for each
% channel (sigma2 is calculated once for the whole signal)

arguments(Input)
    test_signal;
    reference_signal;
    options.skip_front = 0;
    options.skip_end = 0;
    options.returnErrorLocation = 0;
end

options.skip_end = abs(options.skip_end);
options.skip_front = abs(options.skip_front);

assert((options.skip_end+options.skip_front)<length(test_signal),"You can not skip more bits than overall length of data! Set skip_front or skip_end to lower value or check data_in");

if isa(reference_signal,'Signal')
    reference_signal = reference_signal.signal;
end

if isa(test_signal,'Signal')
    test_signal = test_signal.signal;
end

% TRIM
[test_signal,reference_signal]=trimseq(test_signal,reference_signal,options.skip_front,options.skip_end);

% CALC EVM
[GMI,NGMI] = calc_ngmi_(test_signal,reference_signal);


    function [GMI,NGMI] = calc_ngmi_(test_signal,reference_signal)

        assert(length(test_signal) == length(reference_signal),"Sequence length does not match");

        %%% implemented according to [1] J. Cho, L. Schmalen, und P. J. Winzer,
        % „Normalized Generalized Mutual Information as a Forward Error Correction Threshold for Probabilistically Shaped QAM“,
        % in 2017 European Conference on Optical Communication (ECOC), Sep. 2017, S. 1–3. doi: 10.1109/ECOC.2017.8345872.

        error_vector = (test_signal-reference_signal);
        sigma2 = var(error_vector); %noise variance

        %%% Separate Classes
        constellation = unique(reference_signal);
        received_sd = NaN(numel(constellation),length(reference_signal));
        lvlcol = cbrewer2('Set1',numel(constellation));
        for lvl = 1:numel(constellation)
            %Separate the equalized signal into the
            %respective levels based on the actually
            %transmitted level!
            received_sd(lvl,reference_signal==constellation(lvl)) = test_signal(reference_signal==constellation(lvl));
        end

        N = length(test_signal);           % Number of received samples

        M = length(constellation); %P
        m = log2(M); %bits per symbol

        entries = sum(~isnan(received_sd),2)';
        P_X = ones(1,M)/M;%entries./N;

        % Parameters
        symbols = constellation';     % PAM-4 symbols

        gray_bits = PAMmapper(M,0).demap_(constellation);       % Gray coding (bits per symbol)

        % Conditional probability function for AWGN
        q_Y_given_X = @(y, x) (1 / sqrt(2 * pi * sigma2)) * exp(-(y - x).^2 / (2 * sigma2));

        % Entropy term
        H_X = -sum(P_X .* log2(P_X)); % Entropy of input distribution

        tx_bits = PAMmapper(M,0,"eth_style",0).demap(reference_signal);

        % GMI computation
        bit_llr_sum  = 0;
        for k = 1:N %loop over signal
            y_k = test_signal(k); % Current received sample
            % [~, closest_symbol_idx] = min(abs(symbols - y_k)); % Closest symbol index

            for i = 1:m %loop over bit position
                % Extract i-th bit for each symbol
                bit_mask = gray_bits(:, i); % Binary column for i-th bit of all symbols

                % matching_symbols = symbols(bit_mask == gray_bits(closest_symbol_idx, i)); %old, based on decision that mihht be wrong
                matching_symbols = symbols(bit_mask == tx_bits(k,i)); %new, based on tx bits

                % Numerator: Sum over x in x_{b_{k, i}}
                numerator = sum(q_Y_given_X(y_k, matching_symbols) .* P_X(ismember(symbols, matching_symbols)));

                % Denominator: Sum over all x
                denominator = sum(q_Y_given_X(y_k, symbols) .* P_X);

                % Logarithmic contribution
                bit_llr_sum = bit_llr_sum + log2(numerator / denominator);
            end
        end

        % Normalize the noise impact term by N
        bit_llr_sum = bit_llr_sum / N;

        % GMI
        GMI = H_X + bit_llr_sum;
        NGMI = GMI / m;

    end

    function [data_,reference_]=trimseq(data,reference,skipstart,skip_end)

        data_  =  data(skipstart+1:end-skip_end,:);

        delta_bits = length(reference) - length(data);

        skip_end = delta_bits + skip_end;

        reference_ = reference(skipstart+1:end-skip_end,:);

    end

end