% A simple demonstration of confidence-based wavelet filtering.
% This is crude mimicry of the assumption-free wavelet reconstruction of GW150914.
%
% Some key differences from the actual GW work:
%     - Noise is Gaussian
%     - Method is 'frequentist' rather than Bayesian
%     - Reconstructed from a single detector!
%
% CCK 2016-Feb

addpath('../wave_matlab'); % Provide access to the C&T wavelet functions.

stdev = 0.15; % standard deviation for red noise background
lag1 = 0.9;   % lag-1 autocorrelation for red noise background
[t, data, waveform, freq, noise] = chirp(stdev,lag1); % Generate synthetic data
dt = t(2) - t(1);
N = length(data);

%%%%%%% Wavelet transform and plot (cribbed from wavetest.m) %%%%%%%%%

pad = 1;      % pad the time series with zeroes (recommended)
dj = 0.25;    % this will do 4 sub-octaves per octave
s0 = 1e-3;    % set minimum scale at 1ms.
j1 = 7/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

% Wavelet transform:
[data_w,period,scale,coi] = wavelet(data,dt,pad,dj,s0,j1,mother);
frequency = 1./period;
power = (abs(data_w)).^2 ;        % compute wavelet power spectrum

% Significance levels:
signif_level = 0.98
[signif,fft_theor] = wave_signif(var(data),dt,scale,0,lag1,signif_level,-1,mother);
   % It would be more honest to compute lag1 from the data.
sig = (signif')*(ones(1,N));  % expand signif --> (J+1)x(N) array
sig = power ./ sig;         % where ratio > 1, power is significant

figure(3)
hold off
imagesc(t,log2(period),log2(power));
axis([0.2,0.5, min(log2(period)), max(log2(period))])
hold on
xlabel('time (s)')
ylabel('log_2(period) (s)')
title('Wavelet Power Spectrum')
set(gca,'YDir','reverse')
% significance contour, levels at -99 (fake) and 1 (signif_level)
contour(t,log2(period),sig,[-99, 1],'k');
% cone-of-influence, anything "below" is dubious
plot(t,log2(coi),'k')
print("filter3.pdf","-dpdfwrite")

%%%%%%%%% Filter! %%%%%%%%
unreal_stuff = find(sig < 1);
data_w_filt = data_w;
data_w_filt(unreal_stuff) = 0;
data_filt = (dj*sqrt(dt)/Cdelta/Psi0_0) .* sum( real(data_w_filt) ./ (sqrt(scale')*ones(1,N)), 1 ) ;

figure(4)
hold off
plot(t, waveform, 'g', 'linewidth', 2.0)
axis([0.25,0.47,-2,2])
hold on
plot(t, data_filt, 'r', 'linewidth', 2.0)
xlabel('time (s)')
legend("true waveform","wavelet filtered data")
legend("boxoff")
legend("location","northwest")

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