numSymbs = 2^16;
M = 4;

Rsym = 2.5e10; % symbol rate (sym/sec)

rolloff = 0.25;
span = 6; % filter span
sps = 8; % samples per symbol

fs = Rsym * sps; % sampling freq (Hz)
Tsamp = 1 / fs;

t = (0 : 1 / fs : numSymbs / Rsym + (1.5 * span * sps - 1) / fs).';

EbN0_db = 0:0.5:14;
EbN0 = 10 .^ (EbN0_db ./ 10);

Es = 1;
Eb = Es / log2(M);
N0 = Eb ./ EbN0;

EsN0 = EbN0 .* log2(M);
EsN0_db = 10 .* log10(EsN0);

plotlen = length(EbN0);

ber = zeros(1, plotlen);
berNoComp = zeros(1, plotlen);
berAdapt = zeros(1, plotlen);
berMatlabAdapt = zeros(1, plotlen);

data = randi([0 M - 1], numSymbs, 1);
modData = pskmod(data, M, pi / M, 'gray');
x = txFilter(modData, rolloff, span, sps);

%% Simulate chromatic dispersion
D = 17; % ps / (nm km)
lambda = 1550; % nm
z = 3000; % km


[xCD, xCDkstart] = chromaticDispersion(x, D, lambda, z, Tsamp);

TsampOrig = Tsamp;

for i = 1:plotlen
  sps = 8;

  snr = EbN0_db(i) + 10 * log10(log2(M)) - 10 * log10(sps);

  y = awgn(xCD, snr, 'measured');

  r = rxFilter(y, rolloff, span, sps);

  sps = 2;
  Tsamp = TsampOrig * 4;

  [rCDComp, CDCompkstart] = CDCompensation(r, D, lambda, z, Tsamp);
  rCDComp = normalizeEnergy(rCDComp, numSymbs*sps, 1);

  rSampled = rCDComp(2:2:end);
  rNoCompSampled = r(2:2:end);

  %% rotate rNoCompSampled to match original data
  theta = angle(-sum(rNoCompSampled .^ M)) / M;
  %% if theta approx +pi/M, wrap to -pi/M
  if abs(theta - pi / M) / (pi / M) < 0.1
    theta = -pi / M;
  end
  rNoCompSampled = rNoCompSampled .* exp(-j * theta);


  %% Not entirely sure why, but after using FFT instead of time-domain
  %% convolution for simulating CD, we now need to do the same rotation
  %% for rSampled as well, but this time with a positive rotation.
  theta = angle(-sum(rSampled .^ M)) / M;
  if abs(theta + pi / M) / (pi / M) < 0.1
    theta = +pi / M;
  end
  rSampled = rSampled .* exp(-1j * theta);



  %% adaptive filter
  adaptFilterOut = adaptiveCMA(rSampled);

  demodData = pskdemod(rSampled, M, pi / M, 'gray');
  demodNoComp = pskdemod(rNoCompSampled, M, pi / M, 'gray');
  demodAdapt = pskdemod(adaptFilterOut, M, pi / M, 'gray');
  %%demodMatlabAdapt = pskdemod(matlabEq, M, pi / M, 'gray');

  [bitErrors, ber(i)] = biterr(data, demodData);
  [bitErrors, berNoComp(i)] = biterr(data, demodNoComp);
  [~, berAdapt(i)] = biterr(data, demodAdapt);
  %%[~, berMatlabAdapt(i)] = biterr(data, demodMatlabAdapt);

%{
  if EbN0_db(i) == 14
    figure(1);
    scatterplot(normalizeEnergy(rSampled, numSymbs, 1));
    formatFigure;
    title('Constellation after CD comp.', 'interpreter', 'latex');
    xlabel('In-Phase', 'interpreter', 'latex');
    ylabel('Quadrature', 'interpreter', 'latex');
    set(gca, 'FontSize', 18);
    %%scatterplot(modData);
    %%title('Original constellation');
    scatterplot(normalizeEnergy(rNoCompSampled, numSymbs, 1));
    formatFigure;
    title('Constellation without CD comp.', 'interpreter', 'latex');
    xlabel('In-Phase', 'interpreter', 'latex');
    ylabel('Quadrature', 'interpreter', 'latex');
    set(gca, 'FontSize', 18);
    %scatterplot(adaptFilterOut);
    %title('Constellation with CD compensation and adaptive filter');
    %scatterplot(matlabEq);
    %title('Matlab equalizer');
    ber(i)
    %berNoComp(i)
    %berAdapt(i)
    %berMatlabAdapt(i)
  end
%}
end

figure(1);
clf;

%% Plot simulated results
semilogy(EbN0_db, ber, 'r', 'LineWidth', 2);
hold on;
semilogy(EbN0_db, berNoComp, 'm', 'LineWidth', 2);
semilogy(EbN0_db, berAdapt, 'Color', [0, 0.6, 0], 'LineWidth', 2);
%%%semilogy(EbN0_db, berMatlabAdapt, 'c', 'LineWidth', 1.4);

theoreticalPSK(EbN0_db, M, 'b', 'LineWidth', 1);
%%legend({'CD + AWGN + CD comp.', 'CD + AWGN + CD comp.~+ CMA', ...
%%        'Theoretical AWGN'}, 'Location', 'southwest');
%%legend({'CD + AWGN + CD comp.', 'CD + AWGN', 'Theoretical AWGN'}, ...
%%       'Location', 'southwest');
legend({'CD + AWGN + CD comp.', 'CD + AWGN', ...
        'CD + AWGN + CD comp.~+ CMA', 'Theoretical AWGN'}, 'Location', ...
       'Southwest');

%%title(strcat(num2str(M), '-PSK with chromatic dispersion and compensation'));
title({'QPSK with chromatic dispersion and compensation', ...
       strcat(['$D = 17$ ps/(nm km), $z = ', num2str(z), '$ km'])});
grid on;
xlabel('$E_b/N_0$ (dB)');
ylabel('BER');

formatFigure;