# -*- coding: utf-8 -*-
# 2009-12-07 Moritz

from obspy.arclink import Client
from obspy.core import UTCDateTime
from obspy.signal.util import xcorr
import obspy, os
import matplotlib.pyplot as plt
from obspy.signal import bandpass
import numpy as np
plt.close('all')


# Download data, if not already there
file = "BW.RTBE..EHZ.194"
if not os.path.exists(file):
    t = UTCDateTime("2009-07-13T05:30:00")
    cl = Client()
    cl.saveWaveform(file, "BW", "RTBE", "", "EHZ", t, t+(3600*4.5))

# Read the data
tr = obspy.read(file)[0]
df, npts = tr.stats.sampling_rate, tr.stats.npts

# Bandpass to interest region
tr.data = tr.data.astype('float32')
tr.data = bandpass(tr.data, 2, 15, df = df)

# Cut out/zoom into earthquake prototype
data = tr.data[478500:479500].copy()

# Calculate shifts, overlap..
winlen = len(data)*3
ovlap = .5
shilen = len(data)*ovlap
nshifts = int(float(npts - 2*winlen)/(shilen))

# Correlate in a loop
data_xcorr = np.empty(nshifts)
time_xcorr = np.empty(nshifts)
for i in xrange(nshifts):
    if i % 20 == 0: # some verbose output
        print i, nshifts
    _i, _o = xcorr(data, tr.data[i*shilen:i*shilen+winlen], len(data))
    data_xcorr[i] = _o
    time_xcorr[i] = ((i*shilen+winlen/2.0) + _i)/df


#
# The plotting part
#
time = np.arange(0, npts)/df
plt.figure(1)
ax = plt.subplot(211)
plt.plot(time, tr.data, 'k', label='Original Data')
plt.legend()
plt.subplot(212, sharex=ax)
plt.plot(time_xcorr, data_xcorr, 'b', label='Max Xcorr Value')
plt.legend()
plt.xlabel("Time [s]")
plt.show()