numSymbs = 5e4;
M = 4;

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

rolloff = 0.25;
span = 6; % filter span
sps = 2; % 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 = -3:0.25:4;
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


plotlen = length(power);

ber = zeros(1, plotlen);

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


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


for i = 1:plotlen
  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
  %%stepnum = round(40 * z * gamma); % Nonlinear Fiber optics, App B
  stepnum = 100;
  xCD = splitstepfourier(x, D, lambda, z, Tsamp, gamma, stepnum);

  y = awgn(xCD, snr, power(i), 'linear');
  %%y = xCD;

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

  rSampled = rCDComp(sps*span/2+1:sps:(numSymbs+span/2)*sps);
  rNoCompSampled = r(sps*span/2+1:sps:(numSymbs+span/2)*sps);

  %% 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);

  demodAdapt = pskdemod(adaptFilterOut, M, pi / M, 'gray');
  [~, ber(i)] = biterr(data, demodAdapt);
end

figure(1);
clf;

%% Plot simulated results
semilogy(power_dBm, ber, 'Color', [0, 0.6, 0], 'LineWidth', 2);
hold on;

title({'CD + Kerr + CD compensation', ...
       strcat(['$D = 17$ ps/(nm km), $z = ', num2str(z), '$ km']), ...
       strcat(['$E_b/N_0 = ', num2str(N0ref_dB), '$ dB at 1 mW'])});
grid on;
xlabel('Optical power (dBm)');
ylabel('BER');

formatFigure;