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

import obspy, pickle
from obspy.signal import seisSim, pazToFreqResp
import matplotlib.pyplot as plt
import numpy as np
plt.close('all')

# Set water level
wlev = 65.0

# Read in the data and poles
tr = obspy.read("data/GR.FUR..LHN.D.2004.361")[0]
paz = pickle.load(open("data/GR.FUR..LHN.D.paz",'rb'))

# plot freq response
h, f = pazToFreqResp(paz['poles'], paz['zeros'], paz['gain'],
        1./tr.stats.sampling_rate, tr.stats.npts, freq=True)
plt.figure(1)
plt.loglog(f,abs(h), 'k', label='GR.FUR STS2')
swamp = np.abs(h).max() * 10.0 ** (-wlev / 20.0)
plt.axhline(swamp, color='b', label='Water-Level')
plt.legend(loc='lower right')
plt.xlabel('Freq [Hz]')
plt.grid()

# Some Preprocessing
tr.data = tr.data.astype('float64')
tr.data -= tr.data.mean()

# Do loop over different corrections
container = []
for label in ['m/s', 'm']:

    # Do the instrument correction
    # Correct for paz frequency response
    container.append(seisSim(tr.data, tr.stats.sampling_rate, paz, 
                             inst_sim=None, water_level=wlev))
    # Change from count to meter, last index (-1) is current
    container[-1] /= paz['sensitivity']
    # Add zero to integrate
    paz['zeros'].append(0j)

#
# The plotting part
#
time = np.arange(0,tr.stats.npts)/tr.stats.sampling_rate
plt.figure(2)
ax = plt.subplot(311)
plt.plot(time, tr.data, color='k')
plt.ylabel("Raw Counts")
plt.subplot(312,sharex=ax)
plt.plot(time, container[0], color='k')
plt.ylabel("Velocity [m/s]")
plt.subplot(313,sharex=ax)
plt.plot(time, container[1], color='k')
plt.ylabel("Displacement [m]")
plt.xlabel("Time [s]")
plt.suptitle("%s %s" % (tr.id, tr.stats.starttime))
plt.show()