M = 4;
numSymbs = 5e5;

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

span = 6; % Tx/Rx filter span
rolloff = 0.25; % Tx/Rx RRC rolloff
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).';

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

x = normalizeEnergy(x, numSymbs*sps, 1);

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

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

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

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

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

sps = 2;
Tsamp = Tsamp * 4;


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

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

%% if no CD comp, then rotate constellation. Use:
theta = angle(-sum(rNoCompSa .^ M)) / M;
%% if theta approx +pi/M, wrap to -pi/M
if abs(theta - pi / M) / (pi / M) < 0.1
  theta = -pi / M;
end
rNoCompSa = rNoCompSa .* 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);


%%rAdaptEq = adaptiveCMA(rSampled);
%{
%% Compare original signal and compensated signal
figure(101);
clf;
tsym = t(sps*span/2+1:sps:(numSymbs+span/2)*sps);
subplot(211);
plot(t(1:length(x)), real(normalizeEnergy(x, numSymbs*sps, 1)), 'b');
hold on
plot(t(1:length(x)), real(normalizeEnergy(yCDComp(1:length(x)), numSymbs*sps, 1)), 'r');
plot(tsym, real(rAdaptEq), 'x', 'Color', [0, 0.6, 0], 'LineWidth', 2);
hold off;
title('Real part');
legend('original', 'dispersion compensated', 'CMA equalized samples');
axis([t(6000*sps+1) t(6000*sps+150) -Inf +Inf]);
subplot(212);
plot(t(1:length(x)), imag(normalizeEnergy(x, numSymbs*sps, 1)), 'b');
hold on;
plot(t(1:length(x)), imag(normalizeEnergy(yCDComp(1:length(x)), numSymbs*sps, 1)), 'r');
plot(tsym, imag(rAdaptEq), 'x', 'Color', [0, 0.6, 0], 'LineWidth', 2);
hold off;
title('Imag part');
axis([t(6000*sps+1) t(6000*sps+150) -Inf +Inf]);

scatterplot(modData);
formatFigure;
%title('Constellation of original modulation', 'interpreter', 'latex');
xlabel('In-Phase', 'interpreter', 'latex');
%scatterplot(rSampled);
%title('Constellation of matched filter output');
scatterplot(rNoCompSa);
title('Constellation of dispersed signal', 'interpreter', 'latex');
scatterplot(rAdaptEq);
title('Constellation of adaptive filter output');
%}
demodData = pskdemod(rSampled, M, pi / M, 'gray');
%%demodAdapt = pskdemod(rAdaptEq, M, pi / M, 'gray');

[~, ber] = biterr(data, demodData)
%[~, berNoComp] = biterr(data, pskdemod(rNoCompSa, M, pi/M, 'gray'))
%[~, ber] = biterr(data, demodAdapt)