numSymbs = 2^16;
M = 4;

Rsym = 2.5e10; % symbol rate (sym/sec)
Tsym = 1 / Rsym; % symbol period (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).';

power_dBm = -6:1:4;
%%power_dBm = 0;
power = 10 .^ (power_dBm / 10) * 1e-3; % watts

Es = power * Tsym; % joules
Eb = Es / log2(M); % joules

N0ref_db = 10; % Eb/N0 at power = 1mW
%% Fix N0, such that Eb/N0 = N0ref_db at power = 1mW
N0 = 1e-3 * Tsym / (log2(M) * 10 ^ (N0ref_db / 10)); % joules
%% At current settings, N0 = 0.002 pJ

plotlen = length(power);

ber = zeros(1, plotlen);

data = randi([0 M - 1], numSymbs, 1);
modData = dpskmod(data, M, 0, 'gray');
%%modData = pskmod(data, M, pi/4, 'gray');


%% Chromatic dispersion
D = 17; % ps / (nm km)
lambda = 1550; % nm
z = 100; % km


TsampOrig = Tsamp;

for i = 1:plotlen
  sps = 8;
  Tsamp = TsampOrig;

  snr = Es(i) / sps / N0;
  snr_dB = 10 * log10(snr);

  x = txFilter(modData, rolloff, span, sps);
  %% Now, sum(abs(x) .^ 2) / length(x) should be 1.
  %% We can set its power simply by multiplying.
  x = sqrt(power(i)) * x;

  %% We can now do split-step Fourier.
  gamma = 1.2; % watt^-1 / km


  xCDKerr = splitstepfourier(x, D, lambda, z, Tsamp, gamma);

  y = awgn(xCDKerr, snr_dB, 'measured', 'db');
  %y = xCDKerr;

  r = rxFilter(y, rolloff, span, sps);
  sps = 2;
  Tsamp = Tsamp * 4;

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

  rSampled = rCDComp(2:2:end);

  %% adaptive filter
  [adaptFilterOut, convergeIdx] = adaptiveCMA(rSampled);

  demod = dpskdemod(adaptFilterOut, M, 0, 'gray');
  %%demod = pskdemod(adaptFilterOut, M, pi/4, 'gray');

  if convergeIdx < Inf
    [~, ber(i)] = biterr(data(convergeIdx:end), demod(convergeIdx:end));
  else
    [~, ber(i)] = biterr...
                    (data(ceil(0.8*numSymbs):end), ...
                     demod(ceil(0.8*numSymbs):end));
  end
end

ber


figure;
clf;

%% Plot simulated results
qp = 20 * log10(erfcinv(2*ber)*sqrt(2));
plot(power_dBm, qp, 'Color', [0, 0.6, 0], 'LineWidth', 2);
hold on;

title({'CD + Kerr + CD compensation', ...
       strcat(['$D = 17$ ps/(nm km), $z = ', num2str(z), '$ km'])});
grid on;
xlabel('Optical power (dBm)');
ylabel('$20 \log_{10}\left(\sqrt{2}\mathrm{erfc}^{-1}(2 BER)\right)$');

formatFigure;