function [GMI,NGMI] = calc_ngmi(test_signal,reference_signal,options)
% (N)GMI for AWGN, supports PAM2/4/8 (per-symbol) and PAM6 (pair-based, 5 bits / 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 the 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 head/tail consistently
[test_signal,reference_signal] = trimseq(test_signal,reference_signal,options.skip_front,options.skip_end);

% --- Noise variance (real AWGN): use mean-square error for robustness
err    = test_signal - reference_signal;
sigma2 = mean((err(:)).^2);

% --- 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);
P_X = counts / N;

% --- Dispatch: PAM6 needs pairwise treatment (5 bits across 2 symbols)
if M==6
    % ===== PAM6: 5 bits mapped to two consecutive symbols =====
    % Pair the sequence (drop last symbol 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);
    x1 = reference_signal(1:2:2*numPairs);
    x2 = reference_signal(2:2:2*numPairs);

    % Build the 36 transition table and its 5-bit labels (ETH-style)
    symbols = sort(constellation); % keep fixed order
    % all pairs using ndgrid (no dependency on combvec)
    [S1,S2] = ndgrid(symbols, symbols);             % MxM
    trans_pairs = [S1(:), S2(:)];                   % (M^2) x 2

    % Map those pairs to 5-bit labels using your PAMmapper (same as in your BCJR code)
    % NOTE: keep your normalization used by PAMmapper (sqrt(10)) if that's what it expects.
    bits72 = PAMmapper(6,0,"eth_style",0).demap( reshape(trans_pairs',[],1) );
    pam6bits = reshape(bits72', 5, []).';           % (M^2) x 5, row-aligned with trans_pairs

    % Precompute mapping (pair -> bit label) and (symbol -> index)
    [ok1, idx_s1] = ismember(trans_pairs(:,1), symbols);
    [ok2, idx_s2] = ismember(trans_pairs(:,2), symbols);
    assert(all(ok1)&all(ok2), 'Transition level not found in symbols.');
    % For fast masking: linear indices of each (i,j) pair in an MxM grid
    lin_idx_pairs = sub2ind([M,M], idx_s1, idx_s2); % (M^2) x 1

    % Build pairwise prior P_pair = P_X(i)*P_X(j) on the MxM grid
    P_pair = P_X(:) * P_X(:).';                     % MxM
    logP_pair = log(P_pair);                        % for numerical stability

    % Precompute constant term for log-likelihood
    logc = -0.5*log(2*pi*sigma2);

    % Prepare transmitted 5-bit labels per observed pair (to split llr0/llr1)
    % Find each (x1(k),x2(k)) in trans_pairs:
    [tf, row_in_table] = ismember([x1(:), x2(:)], trans_pairs, 'rows');
    assert(all(tf), 'A reference pair was not found in the transition table.');
    tx_bits_pair = pam6bits(row_in_table, :);       % numPairs x 5

    % Precompute masks for bit=0 / bit=1 on the MxM grid for each bit position
    mask1 = false(M,M,5);
    mask0 = false(M,M,5);
    for b = 1:5
        tmp = false(M,M);
        tmp(lin_idx_pairs) = pam6bits(:,b) == 1;
        mask1(:,:,b) = tmp;
        tmp = false(M,M);
        tmp(lin_idx_pairs) = pam6bits(:,b) == 0;
        mask0(:,:,b) = tmp;
    end

    % Loop pairs: compute exact LLRs via log-sum-exp over the MxM grid
    LLR_exact = zeros(numPairs,5);
    for k = 1:numPairs
        % log-likelihood vectors along each axis
        li = logc - ((y1(k) - symbols).^2) / (2*sigma2);  % 1xM
        lj = logc - ((y2(k) - symbols).^2) / (2*sigma2);  % 1xM
        % joint log-weights over all (i,j)
        logW = (li(:) + lj(:).') + logP_pair;             % MxM

        for b = 1:5
            % log P(bit=1 | y) ∝ logsumexp(logW over mask1)
            lnum = logsumexp(logW(mask1(:,:,b)));
            % log P(bit=0 | y) ∝ logsumexp(logW over mask0)
            lden = logsumexp(logW(mask0(:,:,b)));
            LLR_exact(k,b) = lnum - lden;                 % natural-log LLR
        end
    end

    % --- Bitwise MI via consistency relation
    MI = zeros(1,5);
    for b = 1:5
        llr_b = LLR_exact(:,b);
        idx0 = (tx_bits_pair(:,b)==0);
        idx1 = ~idx0;
        I0 = mean( log2(1 + exp( llr_b(idx0))) );   % natural LLR inside exp()
        I1 = mean( log2(1 + exp(-llr_b(idx1))) );
        MI(b) = 1 - 0.5*(I0 + I1);
    end

    % GMI per symbol (5 bits per 2 symbols)
    GMI  = sum(MI)/2;
    NGMI = GMI / 2.5;

else
    % ===== Generic PAM2/4/8 etc.: per-symbol bit labeling =====
    symbols = constellation;                         % 1xM
    % Gray labels per symbol (your helper)
    gray_bits = PAMmapper(M,0).demap_(symbols);     % M x log2(M)
    m = log2(M);

    % Likelihood (real AWGN)
    q = @(y,x) (1/sqrt(2*pi*sigma2)) * exp(-(y - x).^2/(2*sigma2));

    % Transmitted bits per sample
    tx_bits = PAMmapper(M,0,"eth_style",0).demap(reference_signal); % N x m

    % Exact bit-LLRs and MI
    LLR_exact = zeros(N,m);
    for b = 1:m
        mask1 = (gray_bits(:,b)==1);
        mask0 = ~mask1;
        for k = 1:N
            yk = test_signal(k);
            den = sum(q(yk, symbols) .* P_X);
            num1 = sum(q(yk, symbols(mask1)) .* P_X(mask1));
            num0 = sum(q(yk, symbols(mask0)) .* P_X(mask0));
            % Use the ratio of posteriors P(bit=1|y)/P(bit=0|y)
            LLR_exact(k,b) = log(num1) - log(num0); % natural log
        end
    end

    % Bitwise MI
    MI = zeros(1,m);
    for b = 1:m
        llr_b = LLR_exact(:,b);
        idx0 = (tx_bits(:,b)==0);
        idx1 = ~idx0;
        I0 = mean( log2(1 + exp( llr_b(idx0))) );
        I1 = mean( log2(1 + exp(-llr_b(idx1))) );
        MI(b) = 1 - 0.5*(I0 + I1);
    end

    % Per-symbol GMI/NGMI
    GMI  = sum(MI);
    NGMI = GMI / m;
end

% ---------- helpers ----------
function s = logsumexp(a)
    if isempty(a), s = -inf; return; end
    amax = max(a(:));
    s = amax + log(sum(exp(a(:) - amax)));
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