M = 4;
numSymbs = 100000;

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

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

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

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

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

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

yCDComp = CDCompensation(y, D, lambda, z, Tsamp);
%%yCDComp = y;

r = rxFilter(yCDComp, rolloff, span, sps);
rSampled = r(sps*span/2+1:sps:(numSymbs + span/2) * sps);

%% if no CD comp, then rotate constellation. Use:
%{
theta = angle(-sum(rSampled .^ M)) / M;
%% if theta approx +pi/M, wrap to -pi/M
if abs(theta - pi / M) / (pi / M) < 0.1
  theta = -pi / M;
end
rSampled = rSampled .* exp(-j * 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), 'xg');
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), 'xg');
hold off;
title('Imag part');
axis([t(6000*sps+1) t(6000*sps+150) -Inf +Inf]);

scatterplot(modData);
title('Constellation of original modulation');
scatterplot(rSampled);
title('Constellation of matched filter output');
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)
[~, ber] = biterr(data, demodAdapt)