%-- SOLUTIONS to Day 5 Exercises
%-- Peter S. Shawhan

samprate = 16384;
nyquist = samprate/2;

%-------------------------------
% Problem 1

%----- Part 1a
%-- Read in the data
dataA = load('day5dataA.txt');
dataB = load('day5dataB.txt');
%-- Construct array with times of samples
tdata = (0:length(dataA)-1)/samprate;

%----- Part 1b
%-- Create an 8th-order bandpass filter (N=4 creates 8th-order)
[b,a] = butter(4,[64/nyquist 128/nyquist]);
filtdata = filter(b,a,dataA);
%-- Plot and report
plot(tdata,filtdata);
title('Data stream A, filtered from 64 to 128 Hz');
input('(1b) Data stream A plot; Press RETURN to continue')

%----- Part 1c
filtdata = filter(b,a,dataB);
%-- Plot and report
plot(tdata,filtdata);
title('Data stream B, filtered from 64 to 128 Hz');
input('(1c) Data stream B plot; Press RETURN to continue')

%----- Part 1d
fprintf('(1d) Sequence of plots...\n');
%-- Apply Butterworth filters for each 64-Hz band from 128 Hz to 1024 Hz
for flow=128:64:(1024-64)
  fhigh = flow+64;
  [b,a] = butter(4,[flow/nyquist fhigh/nyquist]);
  filtdataA = filter(b,a,dataA);
  filtdataB = filter(b,a,dataB);
  %-- Plot and report
  subplot(2,1,1); plot(tdata,filtdataA);
  title(sprintf('Data stream A, filtered from %d to %d Hz',flow,fhigh));
  subplot(2,1,2); plot(tdata,filtdataB);
  title(sprintf('Data stream B, filtered from %d to %d Hz',flow,fhigh));
  s=input('  Press RETURN to continue, or enter "1" to jump out of loop ');
  if length(s)>0
    subplot(1,1,1);
    break
  end
end
subplot(1,1,1);

%-------------------------------------
% Problem 2

%----- Part 2a
%-- (I already loaded both data streams at the beginning of the file)

%----- Parts 2b/2c/2d
%-- Loop over window sizes
for winsize=[0.020 0.100 0.500]

  winbins = floor(winsize*samprate);
  stepbins = floor(0.010*samprate);

  icorr = 0;
  for ioff=0:stepbins:(length(dataA)-winbins)
    icorr = icorr+1;
    tcorr(icorr) = ioff/samprate;
    corr(icorr) = sum( dataA(ioff+1:ioff+winbins).*dataB(ioff+1:ioff+winbins));
  end
  plot(tcorr,corr);
  title(sprintf('Cross-corr with window = %5.3f sec, unfiltered',winsize));
  input(sprintf('(2cd) Plot for %.3f sec window; Press RETURN to continue',winsize));
end

%----- Part 2e
%-- Apply a Butterworth band-pass filter to each data stream
[b,a] = butter(4,[64/nyquist 2048/nyquist]);
filtdataA = filter(b,a,dataA);
filtdataB = filter(b,a,dataB);

%-- Loop over window sizes
for winsize=[0.020 0.100 0.500]

  winbins = floor(winsize*samprate);
  stepbins = floor(0.010*samprate);

  icorr = 0;
  for ioff=0:stepbins:(length(filtdataA)-winbins)
    icorr = icorr+1;
    tcorr(icorr) = ioff/samprate;
    corr(icorr) = sum( filtdataA(ioff+1:ioff+winbins).*filtdataB(ioff+1:ioff+winbins));
  end
  plot(tcorr,corr);
  title(sprintf('Cross-corr with window = %5.3f sec, filtered',winsize));
  input(sprintf('(2e) Plot for %.3f sec window; Press RETURN to continue',winsize));
end

%-------------------------
% Problem 3

%----- Part 3a
dataC = load('day5dataC.txt');
dataD = load('day5dataD.txt');

%----- Part 3b
%-- The cross-correlation is simple:
cross_corr = sum(dataC.*dataD)
%-- Another way: matrix product
cross_corr = dataC*dataD';

%-- Now look at the cross-correlation using several time-shifts
for ishift = 1:10
  shiftbins = ishift*8192;
  %-- Do a circular shift of the second data stream
  shiftedD = [ dataD(shiftbins+1:length(dataD)) dataD(1:shiftbins) ];
  shiftcorr(ishift) = sum(dataC.*shiftedD);
end
rms_shiftcorr = sqrt(mean(shiftcorr.^2))
SNR = cross_corr / rms_shiftcorr
input('End of part 3b; Press RETURN to continue');

%----- Part 3c
[b,a] = butter(4,[40/nyquist 200/nyquist]);
filtdataC = filter(b,a,dataC);
filtdataD = filter(b,a,dataD);

%-- The cross-correlation is simple:
filtered_cross_corr = sum(filtdataC.*filtdataD)

%-- Now look at the cross-correlation using several time-shifts
for ishift = 1:10
  shiftbins = ishift*8192;
  %-- Do a circular shift of the second data stream
  shiftedD = [ filtdataD(shiftbins+1:length(filtdataD)) filtdataD(1:shiftbins) ];
  shiftcorr(ishift) = sum(filtdataC.*shiftedD);
end
filtered_rms_shiftcorr = sqrt(mean(shiftcorr.^2))
SNR = filtered_cross_corr / filtered_rms_shiftcorr
%-- The SNR does improve, by a significant factor

input('End of part 3c; Press RETURN to continue');

%----- Part 3d
%-- Take the Fourier transforms in preparation for correlating in freq domain
fftC = fft(dataC);
fftD = fft(dataD);

%-- Do the freq-domain cross-correlation, check the SNR
%-- Take the real part to ignore the small imaginary part from round-off
f_cross_corr = real(fftC*fftD');

%-- Now look at the cross-correlation using several time-shifts
for ishift = 1:10
  shiftbins = ishift*8192;
  %-- Do a circular shift of the second data stream
  shiftedD = [ dataD(shiftbins+1:length(dataD)) dataD(1:shiftbins) ];
  shiftcorr(ishift) = real(fftC*fft(shiftedD)');
end
f_rms_shiftcorr = sqrt(mean(shiftcorr.^2));
nonoptimal_freq_SNR = f_cross_corr / f_rms_shiftcorr

%-- To do optimal filtering, we need the power spectrum of the noise.
%-- It's the same as we used on day 3 for inpiral filtering
noisespectrum = load('day3asd.txt');
noisepsd = noisespectrum.^2;

%-- We also need that 1/f^3 factor, so we have to know the frequency of each
%-- bin in Matlab's fft representation.  (Absolute value)
ndata = length(dataC);
nyquist = samprate/2;
nsec = ndata/samprate;
freqarr = (0:ndata-1) / nsec;
toohigh = find(freqarr>=nyquist);
freqarr(toohigh) = abs( freqarr(toohigh) - samprate );
%-- Don't allow the frequency to be zero
freqarr = max(freqarr,1);

%-- OK, now we can assemble our optimal filter kernel.
%-- Note that we're assuming that Omega_GW is flat (constant) over the
%-- relevant frequency band, and that the overlap reduction function is unity
%-- for these two detectors -- i.e. they are co-located.
kernel = freqarr.^-3 ./ (noisepsd.*noisepsd) ;

%-- Do the freq-domain cross-correlation with that kernel
optf_cross_corr = real(sum(fftC.*conj(fftD).*kernel));

%-- Now look at the cross-correlation using several time-shifts
for ishift = 1:10
  shiftbins = ishift*8192;
  %-- Do a circular shift of the second data stream
  shiftedD = [ dataD(shiftbins+1:length(dataD)) dataD(1:shiftbins) ];
  shiftcorr(ishift) = real(sum(fftC.*conj(fft(shiftedD)).*kernel));
end
optf_rms_shiftcorr = sqrt(mean(shiftcorr.^2));
optimal_freq_SNR = optf_cross_corr / optf_rms_shiftcorr

input('End of part 3d; Press RETURN to continue');
