# -*- coding: utf-8 -*-
# 2009-12-05 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, highpass, bandpass
from scipy.io import loadmat

#
# 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=i**2
    Fb = 1 - 1/( 1+(f/f0)**(2*n))
    plt.plot(f,Fb)
    plt.title('highpass 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]) 

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)
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 )
    Fs=highpass(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
# signal
s = tr.data
# time
t = np.arange(0,tr.stats.npts,dtype='f')/tr.stats.sampling_rate

plt.figure(5)
plt.plot(t, s)
plt.title("Original Seismogram")
plt.ylabel("Amplitude")
plt.xlabel("Time [s]")

# zoom figures to t1 and t2
t1 = 200
t2 = 380

plt.figure(3)
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 )
    Fs=highpass(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)

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)
    Fs=highpass(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()