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

import numpy as np
import matplotlib
matplotlib.rc('figure.subplot', hspace=.35, wspace=.35) #adjust subplot layout
matplotlib.rc('font', size=8) # adjust font size of plot
import matplotlib.pyplot as plt
plt.close('all')
import obspy
from obspy.signal import lowpass, lowpassZPHSH
from scipy.io import loadmat

myfilter = lowpass
# uncomment for zero phase lowpass
myfilter = lowpassZPHSH

#
# f-domain representation of lowpass filter filter
#
f = np.arange(0,100,dtype='float64')
f0 = 20 #cut off frequency f0
fs = 15

plt.figure(1)
for i in xrange(1,5):
    plt.subplot(2,2,i)
    # n=1
    n=i**2
    Fb = 1/( 1+(f/f0)**(2*n))
    plt.plot(f,Fb)
    plt.title('Lowpass n=%i, f0=%.2f Hz'% (n,f0))
    plt.xlabel('Frequency [Hz]')
    plt.ylabel('Filter amplitude')
    plt.ylim([plt.ylim()[0]-.1, plt.ylim()[1]+.1]) 


#
# effect on a spike
#
N = 1000
dt = .005;
t = np.linspace(dt,N*dt,N)
s = np.zeros(N)
s[99] = 1

plt.figure(2)
ax = plt.subplot(2,2,1)
plt.plot(t,s)
plt.title('Original function')
plt.xlabel('Time [s]')
plt.ylabel('Amplitude')
plt.ylim([plt.ylim()[0]-.1, plt.ylim()[1]+.1]) 


for i in xrange(2,5):
    n=(i-1)**2
    plt.subplot(2,2,i, sharex=ax )
    Fs=myfilter(s,f0,1./dt,n);
    plt.plot(t,Fs)
    plt.title('Filtered with n=%i, f0=%.2f Hz'% (n,f0))
    plt.xlabel('Time [s]')
    plt.ylabel('Amplitude')
    plt.ylim([plt.ylim()[0]-.01, plt.ylim()[1]+.01]) 


#
# effect on a seismogram
#
tr = obspy.read("data/GR.FUR..LHN.D.2003.268")[0]
t = tr.stats.starttime
tr.trim(t+4600,t+7500) #zoom window
dt = 1./tr.stats.sampling_rate
f0 = .1

s = tr.data
t = np.arange(0,tr.stats.npts,dtype='f')/tr.stats.sampling_rate
t1 = 200
t2 = 380

plt.figure(3)
ax = plt.subplot(2,2,1)
plt.plot(t,s)
plt.title('Original function')
plt.xlabel('Time [s]')
plt.ylabel('Amplitude')
rmax= float(max(abs(s[np.arange(int(t1/dt),int(t2/dt))])))
plt.axis([t1, t2, -rmax, rmax])

for i in xrange(2,5):
    n=(i)**2
    plt.subplot(2,2,i, sharex=ax)
    Fs=myfilter(s,f0,1./dt,n);
    plt.plot(t, Fs)
    plt.title('Filtered with n=%i, f0=%.2f Hz'% (n,f0))
    plt.xlabel('Time [s]')
    plt.ylabel('Amplitude')
    rmax= float(max(abs(s[np.arange(int(t1/dt),int(t2/dt))])))
    plt.axis([ t1, t2, -rmax, rmax ])


plt.figure(4)
ax = plt.subplot(2,2,1)
plt.plot(t,s)
plt.title('Original function')
plt.xlabel('Time [s]')
plt.ylabel('Amplitude')
rmax= float(max(abs(s[np.arange(int(t1/dt),int(t2/dt))])))
plt.axis([ t1,t2, -rmax,rmax ])

for i in xrange(2,5):
    n=4
    f0i =f0/i
    plt.subplot(2,2,i, sharex=ax)
    Fs=myfilter(s,f0i,1./dt,n);
    plt.plot(t,Fs)
    plt.title('Filtered with n=%i, f0=%.3f Hz'% (n,f0i))
    plt.xlabel('Time [s]')
    plt.ylabel('Amplitude')
    rmax= float(max( abs(Fs[np.arange(int(t1/dt),int(t2/dt))])))
    plt.axis([ t1,t2, -rmax,rmax ])
plt.show()