% Parameters
symbols = [0, 1, 2, 3];                  % PAM-4 symbols
P_X = [0.25, 0.25, 0.25, 0.25];         % Uniform probabilities
sigma2 = 0.1;                           % Noise variance
received_samples = [0.2, 1.1, 1.9, 2.8];% Received symbols (example)
gray_bits = [0 0; 0 1; 1 1; 1 0];       % Gray coding (bits per symbol)
m = size(gray_bits, 2);                 % Bits per symbol
N = length(received_samples);           % Number of received samples

% 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

% GMI computation
noise_impact_term = 0;
for k = 1:N
    y_k = received_samples(k); % Current received sample
    [~, closest_symbol_idx] = min(abs(symbols - y_k)); % Closest symbol index
    closest_symbol = symbols(closest_symbol_idx);     % Closest symbol
    
    for i = 1:m
        % 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));

        % 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
        noise_impact_term = noise_impact_term + log2(numerator / denominator);
    end
end

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

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


% Display the result
fprintf('GMI: %.4f bits\n', GMI);
fprintf('NGMI: %.4f bits\n', NGMI);