% (1) Read and plot whitened LIGO data provided by Neil Cornish:
data = dlmread('LIGO_event_white.dat',' ');
t = data(:,1);
hanford = data(:,2);
livingston = data(:,3);
event_times = [0.9,1.1]; % for setting up time axis
data_range = 1.1*[ min([hanford;livingston]),max([hanford;livingston]) ];

figure(1)
hold off
subplot(2,1,1)
plot(t,hanford,'b.-')
axis([event_times, data_range])
title('LIGO Hanford')
xlabel('time (s)')
ylabel('strain')
subplot(2,1,2)
plot(t,livingston,'r.-')
axis([event_times, data_range])
title('LIGO Livingston')
xlabel('time (s)')
ylabel('strain')
print('event_white.pdf','-dpdfwrite')



% (2) Calculate wavelet power spectrum and overplot significance contour:
dt = t(2)-t(1); % sampling period
pad = 1;      % pad the time series with zeroes (recommended)
dj = 1/12;    % this will do 12 sub-octaves (semitones in music lingo) per octave.
   % NOTE: I find that oversampling in frequency produces a less "glitchy" reconstruction.
s0 = 2*dt;    % set minimum scale at Nyquist period
j1 = 4.5/dj;    % this says do 7 powers-of-two with dj sub-octaves each
mother = 'Morlet';
Cdelta = 0.776; % Morlet m=6 reconstruction factor.
Psi0_0 = pi^(-1/4); % Morlet mother wavelet amplitude at origin
signif_level = 0.98

% Some basic signal stats
N = length(t)
mean_hanford = mean(hanford)
var_hanford = var(hanford)
lag1_hanford = mean( (hanford(2:N)-mean_hanford) .* (hanford(1:N-1)-mean_hanford) ) ...
   / var_hanford
mean_livingston = mean(livingston)
var_livingston = var(livingston)
lag1_livingston = mean( (livingston(2:N)-mean_livingston) .* (livingston(1:N-1)-mean_livingston) ) ...
   / var_livingston

% Wavelet transform and significance (Hanford):
addpath('../wave_matlab'); % Provide access to the C&T wavelet functions.
[hanford_w,period,scale,coi] = wavelet(hanford,dt,pad,dj,s0,j1,mother);
frequency = 1./period;
power_hanford = (abs(hanford_w)).^2 ;        % compute wavelet power spectrum
% Significance levels (Hanford):
[signif,fft_theor] = wave_signif(var_hanford,dt,scale,0,lag1_hanford,signif_level,-1,mother);
sig = (signif')*(ones(1,N));  % expand signif --> (J+1)x(N) array
sig_hanford = power_hanford ./ sig;         % where ratio > 1, power is significant

% Wavelet transform and significance (Livingston):
addpath('../wave_matlab'); % Provide access to the C&T wavelet functions.
[livingston_w,period,scale,coi] = wavelet(livingston,dt,pad,dj,s0,j1,mother);
power_livingston = (abs(livingston_w)).^2 ;        % compute wavelet power spectrum
% Significance levels (Livingston):
[signif,fft_theor] = wave_signif(var_livingston,dt,scale,0,lag1_livingston,signif_level,-1,mother);
sig = (signif')*(ones(1,N));  % expand signif --> (J+1)x(N) array
sig_livingston = power_livingston ./ sig;         % where ratio > 1, power is significant

figure(2)
period_range = [min(log2(period)), max(log2(period))]
hold off
subplot(2,1,1) % Hanford
imagesc(t,log2(period),log2(power_hanford));
axis([event_times, period_range])
hold on
xlabel('time (s)')
ylabel('log_2(period) (s)')
title('Wavelet Power Spectrum (Hanford)')
set(gca,'YDir','reverse')
% significance contour, levels at -99 (fake) and 1 (signif_level)
contour(t,log2(period),sig_hanford,[-99, 1],'k');
% cone-of-influence, anything "below" is dubious
plot(t,log2(coi),'k')

subplot(2,1,2) % Livingston
imagesc(t,log2(period),log2(power_livingston));
axis([event_times, period_range])
hold on
xlabel('time (s)')
ylabel('log_2(period) (s)')
title('Wavelet Power Spectrum (Livingston)')
set(gca,'YDir','reverse')
% significance contour, levels at -99 (fake) and 1 (signif_level)
contour(t,log2(period),sig_livingston,[-99, 1],'k');
% cone-of-influence, anything "below" is dubious
plot(t,log2(coi),'k')

print("event_wps.pdf","-dpdfwrite")


% (3) Filter and reconstruct the event.

% Hanford
unreal_stuff = find(sig_hanford < 1);
hanford_w_filt = hanford_w;
hanford_w_filt(unreal_stuff) = 0;
hanford_filt = (dj*sqrt(dt)/Cdelta/Psi0_0) .* sum( real(hanford_w_filt) ...
   ./ (sqrt(scale')*ones(1,N)), 1 ) ;

% Livingston
unreal_stuff = find(sig_livingston < 1);
livingston_w_filt = livingston_w;
livingston_w_filt(unreal_stuff) = 0;
livingston_filt = (dj*sqrt(dt)/Cdelta/Psi0_0) .* sum( real(livingston_w_filt) ...
   ./ (sqrt(scale')*ones(1,N)), 1 ) ;

figure(3)
hold off
subplot(2,1,1) % Hanford
plot(t, hanford_filt, 'b', 'linewidth', 2.0)
axis([event_times, data_range])
title('Reconstructed event (Hanford)')
xlabel('time (s)')
ylabel('strain')

hold off
subplot(2,1,2) % Livingston
plot(t, livingston_filt, 'r', 'linewidth', 2.0)
axis([event_times, data_range])
title('Reconstructed event (Livingston)')
xlabel('time (s)')
ylabel('strain')
print("event_recon.pdf","-dpdfwrite")


% (4) Find the time delay.
a = fft(hanford_filt);
b = fft(livingston_filt);
u = a .* b;
v = a .* conj(b);
U = ifft(u);
V = ifft(v);
foo = (U.^2) + (V.^2);
   % Maximum of this function is the time delay between the waveforms (Neil taught me this trick).

figure(4)
plot(t,foo);
axis([0,0.015, min(foo), max(foo)])


% Save filtered results.
data_out = [t,hanford_filt',livingston_filt'];
dlmwrite('test1_filtered_waveforms.dat',data_out,' ');