imdd_silas/Functions/Metrics/calc_gmi_bitwise.m

function [GMI,NGMI] = calc_gmi_bitwise(test_signal,reference_signal,options)
%  https://cioffi-group.stanford.edu/doc/book/AppendixG.pdf
%  Supports PAM2/4/8 (per-symbol) and PAM6 (pairwise: 5 bits per 2 symbols)

arguments(Input)
    test_signal
    reference_signal
    options.skip_front (1,1) double = 0
    options.skip_end   (1,1) double = 0
    options.returnErrorLocation (1,1) double = 0 %#ok<NASGU>
end

options.skip_end   = abs(options.skip_end);
options.skip_front = abs(options.skip_front);
assert((options.skip_end+options.skip_front) < numel(test_signal), ...
    "You can not skip more samples than overall length.");

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);

% Common precompute
N       = length(test_signal);
M       = numel(unique(reference_signal));
noise   = test_signal - reference_signal;
sigma2  = var(noise);                       % real AWGN variance

% --- Constellation (levels observed)
constellation = unique(reference_signal);
M = numel(constellation);

% --- Empirical P_X (uniform is fine too; keep your original behavior)
N = numel(reference_signal);
counts = arrayfun(@(c) sum(reference_signal==c), constellation);
priors = counts / N;

% ----- PAM6: 5 bits mapped to 2 consecutive symbols (36 transitions) -----
if M == 6
    % Pair the stream (drop last if odd)
    numPairs = floor(N/2);
    if numPairs == 0, GMI = 0; NGMI = 0; return; end

    y1 = test_signal(1:2:2*numPairs);
    y2 = test_signal(2:2:2*numPairs);

    % Use the same level order as PAMmapper to be consistent with labeling
    levels = constellation;%PAMmapper(6,0,"eth_style",0).levels;     % 1x6 (actual amplitudes used)
    % Build all 36 transitions (i,j)
    [S1,S2] = ndgrid(levels, levels);                  % 6x6
    trans_pairs = [S1(:), S2(:)];                      % 36x2

    % 5-bit labels for each transition (ETH style).
    % NOTE: keep the normalization you use with your PAMmapper (often /sqrt(10)).
    bits72   = PAMmapper(6,0,"eth_style",0).demap( reshape(trans_pairs.',[],1)  );
    pairBits = reshape(bits72.', 5, []).';             % 36 x 5

    % Transmitted 5-bit labels for each observed pair (from the reference)
    tx_bits_pair = PAMmapper(6,0,"eth_style",0).demap( reshape(reference_signal(1:2* numPairs).',[],1) );
    tx_bits_pair = reshape(tx_bits_pair.', 5, []).';   % numPairs x 5

    % Precompute symbol log-priors (uniform => constant; included for shaped inputs)
    logP = log(priors);

    % Precompute masks for bit=0 / bit=1 in the 6x6 grid
    mask1 = false(6,6,5);
    mask0 = false(6,6,5);
    for b = 1:5
        tmp = false(6,6);
        tmp(:) = pairBits(:,b)==1;
        mask1(:,:,b) = tmp;
        tmp = false(6,6);
        tmp(:) = pairBits(:,b)==0;
        mask0(:,:,b) = tmp;
    end

    % Exact LLRs via log-sum-exp over the 6x6 grid
    LLR = zeros(numPairs,5);
    cst = -0.5*log(2*pi*sigma2);   % cancels in LLR but harmless
    for k = 1:numPairs
        li = cst - ((y1(k) - levels).^2)/(2*sigma2) + logP;  % 1x6
        lj = cst - ((y2(k) - levels).^2)/(2*sigma2) + logP;  % 1x6
        logW = li(:) + lj(:).';                               % 6x6 (adds logP twice)

        for b = 1:5
            L1 = logsumexp(logW(mask1(:,:,b)));
            L0 = logsumexp(logW(mask0(:,:,b)));
            LLR(k,b) = L1 - L0;                               % natural-log LLR
        end
    end

    % Bit-wise MI using consistency relation
    MI_bits = zeros(1,5);
    for b = 1:5
        idx0 = (tx_bits_pair(:,b)==0);
        idx1 = ~idx0;
        r0 = LLR(idx0,b);
        r1 = LLR(idx1,b);
        I0 = mean( log2(1 + exp( r0)) );     % natural LLR inside exp()
        I1 = mean( log2(1 + exp(-r1)) );
        MI_bits(b) = 1 - 0.5*(I0 + I1);
    end

    % Per-symbol outputs (5 bits per 2 symbols)
    GMI  = sum(MI_bits)/2;
    NGMI = GMI / 2.5;

% ----- Generic PAM2/4/8…: per-symbol, Gray bits directly on symbols -----
else
    nBits   = log2(M);
    levels  = PAMmapper(M,0).levels' / PAMmapper(M,0).scaling;  % match PAMmapper ordering
    grayBits= PAMmapper(M,0).showBitMapping;                    % M x nBits

    % Allocate
    LLR     = zeros(N, nBits);

    for n = 1:N
        y = test_signal(n);
        ll_table = -((y - levels).^2)/(2*sigma2) + log(priors);   % 1xM (log-domain)
        for b = 1:nBits
            idx0 = grayBits(:,b)==0;
            idx1 = ~idx0;
            L0   = logsumexp(ll_table(idx0));
            L1   = logsumexp(ll_table(idx1));
            LLR(n,b) = L1 - L0;                                   % natural-log LLR
        end
    end

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

    % Bit-wise MI
    MI_bits = zeros(1,nBits);
    for b = 1:nBits
        r0 = LLR(tx_bits(:,b)==0,b);
        r1 = LLR(tx_bits(:,b)==1,b);
        I0 = mean( log2(1 + exp( r0)) );
        I1 = mean( log2(1 + exp(-r1)) );
        MI_bits(b) = 1 - 0.5*(I0 + I1);
    end

    GMI  = sum(MI_bits);          % bits / symbol
    NGMI = GMI / nBits;
end

% ===== Helpers =====
function s = logsumexp(a)
    if isempty(a), s = -inf; return; end
    m = max(a(:));
    s = m + log(sum(exp(a(:) - m)));
end

function [data_,reference_] = trimseq(data,reference,skipstart,skip_end)
    data_  = data(skipstart+1:end-skip_end,:);
    delta  = numel(reference) - numel(data_);
    reference_ = reference(skipstart+1:end-(skip_end+delta),:);
end
end