M_format = [2,4,6,8];

for m = 1:length(M_format)
    % --- Parameters ---
    M        = M_format(m);  % PAM order (e.g., 2,4,8)
    Nsym     = 1e5;          % number of symbols
    h        = [1, 0.5];     % Impulse response to remove

    b = log2(M);
    if M == 6 b = 5; end
    rng(1);
    bits_tx = logical(randi([0 1], Nsym, b, 'uint8'));

    tx_symbols = pammap(bits_tx,M);

    if M == 6
        states = unique(tx_symbols);
        pam6transitions = combvec(states',states')'; % pam6transitions =
        bitmapping = pamdemap(reshape(pam6transitions',1,[])',M);
    else
        bitmapping = pamdemap(unique(tx_symbols),M);
    end

    scaling = sqrt(sum(unique(tx_symbols).^2)/numel(unique(tx_symbols)));
    tx_symbols = tx_symbols ./ scaling;

    % apply impulse response to signal
    y_filt = filter(h, 1, tx_symbols);

    sir = 10:25;
    for s = 1:length(sir)

        % apply noise
        y = awgn(y_filt,sir(s),"measured",1);

        % apply bcjr
        BCJR = bcjr_pam("DIR",h,"duobinary_output",0,"M",M,"trellis_states",unique(tx_symbols));
        [viterbi_estimate,LLR,GMI(m,s)] = BCJR.process(y,tx_symbols,bits_tx,bitmapping);

        % decode LLR's
        bits_LLR = LLR > 0;

        % demap viterbi symbols sequence
        rx_symbols = viterbi_estimate .* scaling;
        bits_rx = pamdemap(rx_symbols,M);

        % BER calc
        BER_vit(m,s) = nnz(bits_tx ~= bits_LLR) / numel(bits_tx);
        fprintf('BER LLR = %.2e \n', BER_vit);

        BER_llr(m,s) = nnz(bits_tx ~= bits_rx) / numel(bits_tx);
        fprintf('BER = %.2e \n', BER_llr);
    end
end

figure();hold on
for m = 1:length(M_format)
    plot(sir,BER_llr(m,:),'DisplayName',sprintf('PAM %d',M_format(m)))
    % plot(sir,BER_vit(m,:),'DisplayName',sprintf('PAM %d',M_format(m)),'LineStyle',':','LineWidth',0.1,'HandleVisibility','off');
end
ylabel('BER');
xlabel('SNR')
title('BER vs. SNR');
set(gca, 'XScale', 'linear', ...
    'YScale', 'log', ...
    'TickLabelInterpreter', 'latex', ...
    'FontSize', 11);


figure();hold on
for m = 1:length(M_format)
    plot(sir,GMI(m,:),'DisplayName',sprintf('GMI PAM %d',M_format(m)))
end
ylabel('GMI');
xlabel('SNR')
title('GMI vs. SNR');
set(gca, 'XScale', 'linear', ...
    'YScale', 'linear', ...
    'TickLabelInterpreter', 'latex', ...
    'FontSize', 11);

function symbols = pammap(bits,M)
bits = logical(bits);
if M == 2
    symbols = bits;
elseif M == 4
    symbols= 2*bits(:,1) + (bits(:,1)==bits(:,2));
    symbols=2*symbols-3;

elseif M == 6

    m = 1;

    if size(bits,2)>size(bits,1)
        bits = bits'; %vector aufrecht stellen
    end
    bits = reshape(bits',1,[])';
    thres = [-3 5;-1 5;-3 -5;-1 -5;-5 3;-5 1;-5 -3;-5 -1;-1 3;-1 1;-1 -3;-1 -1;-3 3;-3 1;-3 -3;-3 -1;3 5;1 5;3 -5;1 -5;5 3;5 1;5 -3;5 -1;1 3;1 1;1 -3;1 -1;3 3;3 1;3 -3;3 -1];
    % LUT based mapping
    for k = 1:5:fix(length(bits)/5)*5
        symbols(m:m+1,1) = thres(bin2dec(int2str(bits(k:k+4)'))+1,:);
        m = m+2;
    end

elseif M == 8
    x1 = bits(:,1);
    x2 = (bits(:,1)==bits(:,3));
    x3 = x2~=bits(:,2);

    symbols = 4*x1 + 2*x2 + x3;
    symbols=2*symbols-7;
end
end

function bits = pamdemap(symbols,M)

if M == 2
    thres=0;
elseif M == 4
    thres=[-2,0,2];
elseif M == 6
    thres = [-3 5;-1 5;-3 -5;-1 -5;-5 3;-5 1;-5 -3;-5 -1;-1 3;-1 1;-1 -3;-1 -1;-3 3;-3 1;-3 -3;-3 -1;3 5;1 5;3 -5;1 -5;5 3;5 1;5 -3;5 -1;1 3;1 1;1 -3;1 -1;3 3;3 1;3 -3;3 -1];
elseif M == 8
    thres=-6:2:6;
end

if M ~= 6
    symbols = symbols';
    a = squeeze(repmat(real(symbols),[1 1 length(thres)])); %Eingangssignal in 3 spalten
    b = squeeze(repmat(reshape(thres(:).',[1 1 length(thres)]),[1 length(symbols) 1])); %Threshold in 3 Spalten
    comp_real = a > b; %check for each symbol/ sampling if it exeeds the obj.thresholdseshold 1, 2 or 3
    comp_real=repmat(real(symbols),[1 1 length(thres)]) > repmat(reshape(thres(:).',[1 1 length(thres)]),[1 length(symbols) 1]);
    s1=size(comp_real,1);
    s2=size(comp_real,2);
end

if M == 2
    data_out=abs(comp_real(:,:,1));
elseif M == 4
    data_out=[comp_real(:,:,2); ones(s1,s2) - comp_real(:,:,1) + comp_real(:,:,3)];
elseif M == 6

    if size(symbols,2) > 1
        symbols = symbols.';
    end

    if length(symbols)/2 ~= round(length(symbols)/2)
        symbols = [symbols;0];
    end

    m = 1;
    for n = 1:2:length(symbols)
        dist = sqrt((symbols(n)-thres(:,1)).^2+(symbols(n+1)-thres(:,2)).^2);
        [~,dd_idx] = min(dist);
        %                 dec_out(n:n+1) = LUT(dd_idx,:);
        data_out(m:m+4) = bitget(dd_idx-1,5:-1:1);
        m = m+5;
    end

    data_out = reshape(data_out',5,[]);

elseif M == 8
    data_out=[comp_real(:,:,4);
        comp_real(:,:,1)-comp_real(:,:,3)+comp_real(:,:,5)-comp_real(:,:,7);
        1-comp_real(:,:,2)+comp_real(:,:,6)];
end

bits = data_out';

end