numSymbs = 2^14;
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/2 / Rsym - 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, 0, 'gray');
for i = 2:numSymbs
  modData(i) = modData(i) * modData(i-1);
end


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

linewidthTx = 0; % Hz
linewidthLO = 1e6; % Hz

TsampOrig = Tsamp;

sig_x = txFilter(modData(1:numSymbs/2), rolloff, span, sps);
sig_y = txFilter(modData(numSymbs/2+1:end), rolloff, span, sps);

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;
  txx = sig_x * sqrt(power(i));
  txy = sig_y * sqrt(power(i));

  rot_omega = 1e3; % rad/s
  rot_phi = 2;
  rot_x = txx .* cos(rot_omega * t) + ...
          txy .* sin(rot_omega * t) * exp(-1j * rot_phi);
  rot_y = txx .* -sin(rot_omega * t) * exp(1j * rot_phi) + ...
          txy .* cos(rot_omega * t);

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

  [xCDKerr, yCDKerr] = ssf_pdm(rot_x, rot_y, ...
                               D, lambda, z, Tsamp, gamma);

  xpn = phaseNoise(xCDKerr, linewidthTx, linewidthLO, Tsamp);
  ypn = phaseNoise(yCDKerr, linewidthTx, linewidthLO, Tsamp);

  xout = awgn(xpn, snr_dB, 'measured', 'db');
  yout = awgn(ypn, snr_dB, 'measured', 'db');

  rx = rxFilter(xout, rolloff, span, sps);
  ry = rxFilter(yout, rolloff, span, sps);
  sps = 2;
  Tsamp = Tsamp * 4;

  rxCDComp = CDCompensation(rx, D, lambda, z, Tsamp);
  ryCDComp = CDCompensation(ry, D, lambda, z, Tsamp);

  rxSampled = rxCDComp(1:2:end);
  rySampled = ryCDComp(1:2:end);

  %% adaptive filter
  [xCMA, yCMA] = pdm_adaptiveCMA(rxSampled, rySampled);

  xpncorr = phaseNoiseCorr(xCMA, M, 0, 40).';
  ypncorr = phaseNoiseCorr(yCMA, M, 0, 40).';

  demodx = pskdemod(xpncorr, M, 0, 'gray');
  remodx = pskmod(demodx, M, 0, 'gray');
  delayx = [1; remodx(1:end-1)];
  demodx = pskdemod(remodx .* conj(delayx), M, 0, 'gray');
  clear remodx
  clear delayx

  demody = pskdemod(ypncorr, M, 0, 'gray');
  remody = pskmod(demody, M, 0, 'gray');
  delayy = [1; remody(1:end-1)];
  demody = pskdemod(remody .* conj(delayy), M, 0, 'gray');
  clear remody
  clear delayy


  [~, ber(i)] = biterr(data, [demodx; demody]);
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;