%%%%%%% MATLAB (R) script to compute TWDP %%%%%%%%%%%%%%%%%%%%%%
%% Editor:  This block of comments to be removed before including in Draft 2.3 (if accepted).
%% Modified from Draft 2.2 version to include:
%%  a) Automated OMA estimation
%%  b) Optimal offset calculation
%%  c) Finite 14,5 equalizer
%% Submitted by N. Swenson, T. Lindsay, P. Voois -- ClariPhy

%% TP2 test inputs
%% The values given below for TxDataFile and MeasuredWaveformFile are examples and should be
%% replaced by actual path\filenames for each waveform tested.
%% Transmit data file: The transmit data sequence is one of the TWDP test patterns defined in
%% Table 68-6. The file format is ASCII with a single column of chronological ones and zeros
%% with no headers or footers.
TxDataFile = 'prbs9_950.txt';
%% Measured waveform: The waveform consists of exactly N samples per unit interval T, where N
%% is the oversampling rate. The waveform must be circularly shifted to align with the data
%% sequence. The file format for the measured waveform is ASCII with a single column of
%% chronological numerical samples, in optical power, with no headers or footers.
MeasuredWaveformFile = 'preproc-1207-01.txt';
OverSampleRate = 16; % Oversampling rate, must be even
%% Simulated fiber responses, modeled as a set of ideal delta functions with specified
%% amplitudes in optical power and delays in nanoseconds, in rows. The three cases specified in
%% Table 68-5 for the comprehensive stressed receiver tests are used. The vector 'PCoefs'
%% contains the amplitudes, and the vector 'Delays' contains the delays.
FiberResp = [...
    0.000000 0.072727 0.145455 0.218182
    0.158 0.176 0.499 0.167
    0.000 0.513 0.000 0.487
    0.254 0.453 0.155 0.138];
Delays = FiberResp(1,:)';

%% Program constants %%
SymbolPeriod = 1/10.3125; % Symbol period (ns)
EqNf = 14; EqNb = 5; % Number of feedforward and feedback equalizer taps
% Set search range for equalizer delay, specified in symbol periods. Lower end of range is
% minimum channel delay. Upper end of range is the sum of the lengths of the FFE and channel.
% Round up and add 5 to account for the antialiasing filter.
EqDelMin = floor(min(Delays)/SymbolPeriod);
EqDelMax = ceil(EqNf/2 + max(Delays)/SymbolPeriod)+5;
EqDelVec = [EqDelMin:EqDelMax];
PAlloc = 6.5; % Total allocated dispersion penalty (dBo)
Q0 = 7.03; % BER = 10^(-12)
N0 = SymbolPeriod/2 / (Q0 * 10^(PAlloc/10))^2;
EFilterBW = 7.5; % Front end filter bandwidth (GHz)

%% Load input waveform and data sequence, generate filter and other matrices
yout0 = load(MeasuredWaveformFile);
XmitData = load(TxDataFile);
PtrnLength = length(XmitData);
TotLen = PtrnLength*OverSampleRate;
Fgrid = [-TotLen/2:TotLen/2-1].'/(PtrnLength*SymbolPeriod);
[b,a] = butter(4, 2*pi*EFilterBW,'s'); H_r = freqs(b,a,2*pi*Fgrid);% antialiasing filter
ExpArg = -j*2*pi*Fgrid;
ONE=ones(PtrnLength,1);

%% Normalize the received OMA to 1.  Estimate the OMA of the captured waveform
%% by using a linear fit to estimate a pulse response, sythesize a square wave,
%% and calculate the OMA of the sythesized square wave per Clause 52.9.5
ant=4; mem=40; %Anticipation and memory parameters for linear fit
X=zeros(ant+mem+1,PtrnLength); %Size data matrix for linear fit
Y=zeros(OverSampleRate,PtrnLength); %Size observation matrix for linear fit
for ind=1:ant+mem+1
    X(ind,:)=circshift(XmitData,ind-ant-1)';%Wrap appropriately for lin fit
end
X=[X;ones(1,PtrnLength)]; %The all-ones row is included to compute the bias
for ind=1:OverSampleRate
    Y(ind,:)=yout0([0:PtrnLength-1]*OverSampleRate+ind)'; %Each column is one bit period
end
Qmat=Y*X'*(X*X')^(-1); %Coefficient matrix resulting from linear fit. Each column (except
%the last) is one bit period of the pulse response.  The last column is the bias.
SqWvPer=16; %Even number; sets the period of the square wave used to compute the OMA
SqWv=[zeros(SqWvPer/2,1);ones(SqWvPer/2,1)]; %One period of square wave (column)
X=zeros(ant+mem+1,SqWvPer); %Size data matrix for synthesis
for ind=1:ant+mem+1
    X(ind,:)=circshift(SqWv,ind-ant-1)'; %Wrap appropriately for synthesis
end
X=[X;ones(1,SqWvPer)]; %Include the bias
Y=Qmat*X;Y=Y(:); %Synthesize the modulated square wave, put into one column
avgpos=[.4*SqWvPer/2*OverSampleRate:.6*SqWvPer/2*OverSampleRate]; %samples to average over
ZeroLevel=mean(Y(round(avgpos),:)); %Average over middle 20% of "zero" run
%Average over middle 20% of "one" run, compute OMA
MeasuredOMA=mean(Y(round(SqWvPer/2*OverSampleRate+avgpos),:))-ZeroLevel;
%% Subtract zero level and normalize OMA 
yout0 = (yout0-ZeroLevel)/MeasuredOMA;

%% Compute the noise autocorrelation sequence at the output of the front-end 
%% antialiasing filter and rate-2/T sampler.
Snn = N0/2 * fftshift(abs(H_r).^2) * 1/SymbolPeriod * OverSampleRate;
Rnn = real(ifft(Snn));
Corr = Rnn(1:OverSampleRate/2:end);
C = toeplitz(Corr(1:EqNf));

%% Compute the minimum slicer MSE and corresponding TWDP for the
%% three stressor fibers
X = toeplitz(XmitData, [XmitData(1); XmitData(end:-1:2)]);%Used in MSE calculation
Rxx = X'*X; %Used in MSE calculation
TrialTWDP = [];
for ii=1:3 %index for stressor fiber
    %% Propagate the waveform through fiber ii.
    %% The DC response of each fiber is normalized to 1.
    PCoefs = FiberResp(ii+1,:)';
    Hsys = exp(ExpArg * Delays') * PCoefs; Hx = fftshift(Hsys/sum(PCoefs));
    yout = real(ifft(fft(yout0).*Hx));
    %% Process signal through front-end antialiasing filter %%%%%%%%%%%%%%%%%%
    yout = real(ifft(fft(yout) .* fftshift(H_r)));
    %% Compute MMSE-DFE %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    %% The MMSE-DFE filter coefficients computed below minimize mean-squared error at the
    %% slicer input. The derivation follows from the fact that the slicer input over the period
    %% of the data sequence can be expressed as Z = (R+N)*W - X*[0 B]', where R and N are
    %% Toeplitz matrices constructed from the signal and noise components, respectively, at the
    %% sampled output of the antialiasing filter, W is the feedforward filter, X is a
    %% Toeplitz matrix constructed from the input data sequence, and B is the feedback filter.
    %% The computed W and B minimize the mean square error between the input to the slicer and
    %% the transmitted sequence due to residual ISI and Gaussian noise.
    %% Minimize MSE over 2/T sampling phase and FFE delay and determine BER
    MseOpt = Inf;
    for jj= [0:OverSampleRate-1]-OverSampleRate/2 % sampling phase
        %% Sample at rate 2/T with new phase (wrap around as required)
        yout_2overT = yout(mod([1:OverSampleRate/2:TotLen]+jj-1,TotLen)+1);
        Rout = toeplitz(yout_2overT, [yout_2overT(1); yout_2overT(end:-1:end-EqNf+2)]);
        R = Rout(1:2:end, :);
        RINV = inv([R'*R+PtrnLength*C R'*ONE;ONE'*R PtrnLength]);
        R=[R ONE]; % Add all-ones column to compute optimal offset
        Rxr = X'*R; Px_r = Rxx - Rxr*RINV*Rxr';
        %% Minimize MSE over equalizer delay
        for kk = 1:length(EqDelVec)
            EqDel = EqDelVec(kk);
            SubRange = [EqDel+1:EqDel+EqNb+1];
            SubRange = mod(SubRange-1,PtrnLength)+1;
            P = Px_r(SubRange,SubRange);
            P00 = P(1,1); P01 = P(1,2:end); P11 = P(2:end,2:end);
            Mse = P00 - P01*inv(P11)*P01';
            if (Mse<MseOpt)
                MseOpt = Mse;
                B = -inv(P11)*P01'; % Feedback filter
                XSel = X(:,SubRange);
                W = RINV*R'*XSel*[1;B]; % Feedforward filter
                Z = R*W - XSel*[0;B]; % Input to slicer
                %% STEP 6 - Compute BER using semi-analytic method %%%%%%%%%%%%%%%%%%%%%%
                MseGaussian = W(1:end-1)'*C*W(1:end-1);
                Ber = mean(0.5*erfc((abs(Z-0.5)/sqrt(MseGaussian))/sqrt(2)));
            end
        end
    end

    %% Compute equivalent SNR %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    %% This function computes the inverse of the Gaussian error probability function. The
    %% built-in function erfcinv() is not sensitive enough for low probability of error cases.
    if Ber>10^(-12) Q = sqrt(2)*erfinv(1-2*Ber);
    elseif Ber>10^(-323) Q = 2.1143*(-1.0658-log10(Ber)).^0.5024;
    else Q = inf;
    end

    %% Compute penalty %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    RefSNR = 10 * log10(Q0) + PAlloc;
    TrialTWDP(ii) = RefSNR-10*log10(Q);
end

%% Pick highest value due to the multiple fiber responses from TrialTWDP.
TWDP = max(TrialTWDP)
%% End of program