close all

% SETUP PARAMETERS
nSegments    = 1000;              % Specify number of time iterations
distributionParameter = 0.095; %(for exponential to be similar to simulation)

%Bandwidth
BW_PD = 200;

%Photodetector Noise amplitude
noiseAmplitude = 0.0001: 0.03: 20;
noiseAmplitude = sqrt(noiseAmplitude.*1e-3);

start_windowSize = 150;

delay = 60;

showcorr = 0;

% OPTIONS
%  For correlation to use two times the size of the reference (optional)
    referenceTwiceSize = 1;
%  filter the noise to have same bandwidth as the trace
    filterNoise = 1;


%% Generate trace

coarseDelay = fix(delay);
fineDelay = delay*100 - coarseDelay*100;

save_var_corr_END = [];
save_var_LS_END = [];

% Generate the trace

rawData = detrend(normrnd(0,distributionParameter,1000,1),'constant');   % Generate 9000 point trace, force zero mean, and filter it with a gauss shape.
rawData_f = fft(rawData);

f = -length(rawData_f)/2:length(rawData_f)/2-1;
gauss_pass = exp(-f.^2/(2*BW_PD^2));
rawData = ifft(ifftshift(fftshift(rawData_f)'.*gauss_pass)).'; % Filtered
rawData=repmat(rawData.',nSegments,1);

% Delay the trace

fineDelay = fineDelay + 1;
rs_pdTrace = resample(rawData(1:end,:).',100,1).'; %Resample the trace, each point becomes 100.
rawData(1,:) = rs_pdTrace(1,1:100:end);            % Save the reference trace
rawData(2:end,:) = rs_pdTrace(2:end,fineDelay:100:end); % Generate the other traces.


% START SIMULATION
startPosTrace = 0.52 + 0.02*rand;                                      % Pick a window at random (Position is defined so that it finds a perturbation in case of generated trace w/ perturbation).

% Allocate storage for the simulated variances.
save_var_corr = zeros(1,length(noiseAmplitude));
save_var_LS = zeros(1,length(noiseAmplitude));
save_snr = zeros(1,length(noiseAmplitude));
idx = 0;

% save calculated variances
saveCalcVar_DC = [];
saveCalcVar_LS = [];

for nAmp = noiseAmplitude
    % Variables to calculate the variances
    x_dc = []; y_dc = []; z_dc = [];
    x_ls = []; y_ls = []; z_ls = [];

    % Run it for increasing photodetector noise.
    idx = idx+1;
    noise = nAmp * randn(size(rawData));   % Generate a matrix of normal distributed values with the amplitude of the noise.

    if filterNoise == 1
        for seg = 1:size(noise,1)     % Filter each of the noises using the same bandpass filter as the signal.
            noise_f = fft(noise(seg,:));
            f = -length(noise(seg,:))/2:length(noise(seg,:))/2-1;
            gauss_pass = exp(-f.^2/(2*BW_PD^2));

            noise_f_filtered = fftshift(noise_f).*gauss_pass;

            noise(seg,:) = ifft(ifftshift(noise_f_filtered));

        end
    end


    SNR = mean((var(rawData.')) ./ (var(noise.')));    % Calculate SNR as the variance of signal over variance of noise.
    save_snr(idx) = SNR;

    traceMatrix = rawData + noise;

    if plot_traces == 1
        figure(1), subplot(2,2,1),plot(rawData(1,:)) , hold on
        figure(1),subplot(2,2,1), plot(noise(1,:)),
        hold off
        figure(1),subplot(2,2,2), plot(traceMatrix(1,:))
        figure(1),subplot(2,2,3), plot(abs(fft(noise(1,:))));
        figure(1),subplot(2,2,4),plot(abs(fft(traceMatrix(1,:))));
        drawnow;
    end

    windowSize = start_windowSize;
    pd_trace = traceMatrix;
    windowPos = fix(length(pd_trace(1,:))*startPosTrace);
    timeIncrement = 1;       % Time increment (only relevant in case there is some changing perturbation to our signal)




    % Correlation
    if referenceTwiceSize
            win_t0 = detrend(pd_trace(1,windowPos-windowSize/2:windowPos+windowSize+windowSize/2),'constant');
    else
        win_t0 = detrend(pd_trace(1,windowPos:windowPos+windowSize),'constant');
    end

    save_lags = [];

    for i = 2:timeIncrement:nSegments-1                               % TDE by direct correlation
        win = detrend(pd_trace(i,windowPos+coarseDelay:windowPos+coarseDelay+windowSize),...    % Make sure that the window is zero-mean.
            'constant');
        cor = xcorr(win_t0,win); % Correlate window_t1 with window_t0

        [~,max_cor]=max(cor);            % Find position of maximum

        %Save lag values
        x_dc = [x_dc cor(max_cor+1)];
        y_dc = [y_dc cor(max_cor-1)];
        z_dc = [z_dc cor(max_cor)];

        lag = quadFit_lag(cor,max_cor); % Parabolic fit using nearest neighbours.

        save_lags = [save_lags lag];                                    % Save the lag


        if showcorr
            figure(2),
            if referenceTwiceSize
                subplot(2,1,1), plot(fix(-length(cor)/2):1:fix(length(cor)/2),cor,'o-'), hold on;
            else
                subplot(2,1,1), plot(fix(-windowSize):1:fix(windowSize),cor,'o-'), hold on;
            end
        end

    end

    cov_xy = cov(x_dc,y_dc); cov_zx = cov(z_dc,x_dc); cov_yz = cov(y_dc,z_dc);
    X = mean(x_dc);  Y = mean(y_dc); Z = mean(z_dc);

    calc_var_DC = 1.^2 * (2./(X-2*Z+Y)^4 * ((Y-Z).*(Z-X)*cov_xy(3) + ...
        (X-Y).*(Y-Z)*cov_zx(3) + (Z-X).*(X-Y)*cov_yz(3) + ...
        ((Y - Z)^2/2)*var(x_dc) + ((Z - X)^2/2)*var(y_dc) + ...
        ((X - Y)^2/2)*var(z_dc)));


    saveCalcVar_DC = [saveCalcVar_DC calc_var_DC];
    save_var_corr(idx) = var(save_lags);





% LEAST MEAN SQUARES
%
%

    if referenceTwiceSize                                   % For correlation : Try using a reference window 2x the size.
        win_t0 = detrend(pd_trace(1,windowPos-windowSize/2:windowPos+windowSize+windowSize/2),'constant');
    else
        win_t0 = detrend(pd_trace(1,windowPos:windowPos+windowSize), ...    % define a new window t0, force zero-mean (just in case.)
        'constant');
    end

    save_lags_LS = [];

    for i = 2:timeIncrement:nSegments-1                               % define a new window t0, force zero-mean (just in case.)
        win = detrend(pd_trace(i,windowPos+coarseDelay:windowPos+coarseDelay+windowSize), ...   % define a new window t1, force zero-mean (just in case.)
            'constant');

        leastSquares = ls_TDE(win_t0,win,windowSize);
        if showcorr
            figure(2);
            subplot(2,1,2), plot(fix(-windowSize/2):1:fix(windowSize/2),leastSquares,'o-'), hold on %title(strcat(num2str(100*i*shift_increment/windowSize),'% of windowSize shift'));
        end

        [~,min_LS]=min(leastSquares);    % Find the minimum.


        if min_LS == 1 || min_LS == length(leastSquares) % Just to prevent errors in very low snr
            lag = min_LS;
        else
            lag = quadFit_lag(leastSquares,min_LS);                                % Parabolic fit using nearest neighbours.
        end

        save_lags_LS = [save_lags_LS lag];

        if min_LS == 1 || min_LS == length(leastSquares)  % This is simply to prevent errors when anomalous detection occurs at very low SNR.
            x_ls = [x_ls leastSquares(min_LS)];
            y_ls = [y_ls leastSquares(min_LS)];
        else
            x_ls = [x_ls leastSquares(min_LS+1)];
            y_ls = [y_ls leastSquares(min_LS-1)];
        end
        z_ls = [z_ls leastSquares(min_LS)];

    end

    cov_xy = cov(x_ls,y_ls); cov_zx = cov(z_ls,x_ls); cov_yz = cov(y_ls,z_ls);
    X = mean(x_ls); Y = mean(y_ls); Z = mean(z_ls);

    calc_var_LS = 1.^2 * (2./(X-2*Z+Y)^4 * ((Y-Z).*(Z-X)*cov_xy(3) + ...
        (X-Y).*(Y-Z)*cov_zx(3) + (Z-X).*(X-Y)*cov_yz(3) + ...
        ((Y - Z)^2/2)*var(x_ls) + ((Z - X)^2/2)*var(y_ls) + ...
        ((X - Y)^2/2)*var(z_ls)));

    saveCalcVar_LS = [saveCalcVar_LS calc_var_LS];

    save_var_LS(idx) = var(save_lags_LS);


end


fprintf('\n\n============\n Measured delay : \n \t CORR : %i \n \t LS : %i\n', median(save_lags)-151-225*referenceTwiceSize,median(save_lags_LS)-76)


% Display Results


% Theoretical values.
% theocorr = (1/windowSize) * 3/pi.^2 * ( (1 + 2 .* save_snr)./save_snr.^2 + 1);
% theoLS = (1/windowSize) * 3/pi.^2 * ( (2 + 2 .* save_snr)./save_snr.^2);

cramer_rao = (1/windowSize) * 3/(pi.^2) * ( (1 + 2 .* save_snr)./save_snr.^2);
save_snr = 10*log10(save_snr);

figure(66), semilogy((save_snr),save_var_corr,'.-b',(save_snr),save_var_LS,'.-r',(save_snr),saveCalcVar_DC,'--b',(save_snr),saveCalcVar_LS,'--r',(save_snr),cramer_rao,'-k'),
% hold on
% semilogy((save_snr),theocorr,(save_snr),theoLS);





% Aux Functions

function [ lag ] = quadFit_lag( corr_, coarse )

    fine = (corr_(coarse+1) - corr_(coarse-1)) / (corr_(coarse+1) + corr_(coarse-1) - 2*corr_(coarse));
    lag = -(1/2)*fine + coarse;

end


function [leastSquares] = ls_TDE(win_t0,win, windowSize)
    sumResiduals = 0;                                               % Result of the sum of the residuals. Define it as zero for each iteration.
    if length(win_t0) == 2*windowSize+1  % IF WE'RE DOING TWICE THE REFERENCE SIZE
        for el = 1:windowSize                                           % For position = 0. Just the residual for every pair of points squared.
            sumResiduals = sumResiduals + (win(el)-win_t0(el+windowSize/2))^2;
        end
    else
        for el = 1:windowSize                                           % For position = 0. Just the residual for every pair of points squared.
            sumResiduals = sumResiduals + (win(el)-win_t0(el))^2;
        end
    end

    leastSquares = (1/windowSize * sumResiduals);                   % Divide it by the number of points.

    if length(win_t0) == 2*windowSize+1  % IF WE'RE DOING TWICE THE REFERENCE SIZE
        for pos = 1:(windowSize/2)
            sumResiduals = 0;
            for el = 1:windowSize
                sumResiduals = sumResiduals + (win(round(el))-win_t0(windowSize/2-pos+el))^2;
            end
            leastSquares = [(1/(windowSize))*sumResiduals leastSquares];

            sumResiduals = 0;
            for el = 1:windowSize
                sumResiduals = sumResiduals + (win(round(el))-win_t0(windowSize/2+pos+el))^2;
            end
            leastSquares = [leastSquares (1/(windowSize))*sumResiduals];

        end

    else
        for pos = 1:(windowSize/2)                                      % Shift the position of window t1 up to half the total windowSize.
            sumResiduals = 0;
            for el = 1:windowSize-pos                                   % Shift to the left, so less one point to calculate.
                sumResiduals = sumResiduals + (win(round(pos+el))-win_t0(el))^2; % Add every point from (starting on the amount of points delayed for the wt1, and on zero for wt0).
            end

            leastSquares = [(1/(windowSize-pos))*sumResiduals leastSquares];

            sumResiduals = 0;
            for el = 1:windowSize-pos                                   % Shift to the right.
                sumResiduals = sumResiduals + (win(round(el))-win_t0(pos+el))^2; % Add every point from (starting on the amount of points delayed for the wt0, and on zero for wt1).
            end
            leastSquares = [leastSquares (1/(windowSize-pos))*sumResiduals];
        end
    end
end