M = 4;
order = 19;
randkey = 1;

bitpattern = [];
s = RandStream('twister','Seed',randkey);
for i = 1:log2(M)
    N = 2^(order-1); %length of prbs
    bitpattern(:,i) = randi(s,[0 1], N, 1);
end

if M == 6
    bitpattern = reshape(bitpattern',[],1);
    bitpattern = bitpattern(1:end-mod(length(bitpattern),5));
end

Bits = Informationsignal(bitpattern);

Symbols = PAMmapper(M,0).map(Bits);
Symbols.fs = 200e9;

Bits_ = PAMmapper(M, 0, "eth_style", 0).demap(Symbols);

% --- Channel: minimal ISI response + AWGN ---
h = [0.3 0.9 0.3];               % impulse response (normalized later if desired)
h = h / norm(h);                          % optional normalization for unit energy

symbols_filt = Symbols.filter(h,1);


%% SHOW Loss during training

mu = logspace(-3,-0.8,12);
ber_ml_mlse = zeros(size(mu));
ber_training = [];
ce_training = [];

parfor i = 1:numel(mu)
    
    symbols_noi = symbols_filt;
    SNR_dB = 20;
    symbols_noi.signal = awgn(symbols_filt.signal, SNR_dB, 'measured');  % AWGN with given SNR
    
    ml_mlse_equalizer = ML_MLSE("epochs_tr",200,"epochs_dd",1,"len_tr",2^15,...
        "mu_dd",mu(i),"mu_tr",mu(i),"order",5,"sps",1,...
        "traceback_depth",128,"L",3,"delta",0,'adaptive_mu',0);
    
    [y_ml_mlse,y_ref] = ml_mlse_equalizer.process(symbols_noi,Symbols);
    ref_bits = PAMmapper(M, 0, "eth_style", 0).demap(y_ref);
    ml_mlse_bits = PAMmapper(M, 0, "eth_style", 0).demap(y_ml_mlse);
    [~, errors, ber_ml_mlse(i), errpos] = calc_ber(ml_mlse_bits.signal, ref_bits.signal, "skip_front", 10, "skip_end", 10, "returnErrorLocation", 1);
    fprintf('ML MLSE BER: %.2e \n',ber_ml_mlse(i));
    ber_training(i,:) = ml_mlse_equalizer.ber;
    ce_training(i,:) = ml_mlse_equalizer.ce;
end

%%
symbols_noi = symbols_filt;
SNR_dB = 20;
symbols_noi.signal = awgn(symbols_filt.signal, SNR_dB, 'measured');  % AWGN with given SNR
ml_mlse_equalizer_adap = ML_MLSE("epochs_tr",200,"epochs_dd",1,"len_tr",2^16,...
    "mu_dd",1,"mu_tr",1,"order",5,"sps",1,...
    "traceback_depth",128,"L",3,"delta",0,"adaptive_mu",1);

[y_ml_mlse,y_ref] = ml_mlse_equalizer_adap.process(symbols_noi,Symbols);
ref_bits = PAMmapper(M, 0, "eth_style", 0).demap(y_ref);
ml_mlse_bits = PAMmapper(M, 0, "eth_style", 0).demap(y_ml_mlse);
[~, errors, ber_ml_mlse_, errpos] = calc_ber(ml_mlse_bits.signal, ref_bits.signal, "skip_front", 10, "skip_end", 10, "returnErrorLocation", 1);
fprintf('ML MLSE BER: %.2e \n',ber_ml_mlse_);

%% 
figure();hold on

plot(mu,ber_ml_mlse,'DisplayName','ML-based MLSE');

beautifyBERplot;
xlim([mu(1), mu(end)]);
xlabel('mu');
ylabel('BER');
title('PAM-4; M=3; AWGN Channel');
ylim([1e-5 0.1]);

%%
figure()
hold on;
cols = cbrewer2('Spectral',12);
for i = 1:12
    plot(1:200,ber_training(i,:),'DisplayName', sprintf('mu=%.3f ',mu(i)),'Color',cols(i,:));
end
set(gca,'YScale','log');
xlabel('Epoch');
ylabel('BER');
title('PAM-4; L=3; SNR=20; AWGN Channel');
plot(1:200,ml_mlse_equalizer_adap.ber,'DisplayName','Adaptive mu');

%%
figure()
hold on;
cols = cbrewer2('Spectral',12);
for i = 1:12
    plot(1:200,ce_training(i,:),'DisplayName', sprintf('mu=%.3f ',mu(i)),'Color',cols(i,:));
end
set(gca,'YScale','log');
xlabel('Epoch');
ylabel('Cross-Entropy');
title('PAM-4; L=3; SNR=20; AWGN Channel');
plot(1:200,ml_mlse_equalizer_adap.ce,'DisplayName','Adaptive mu');


%% SUPER LONG EPOCHS