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

import obspy
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.mlab import detrend

def stringToBool(s):
    """
    Special function for parsing S and T modes of eiglst
    """
    if s == 'S':
        return True
    return False



# Load data
tr = obspy.read("data/GR.FUR..LHE.D.2008.133")[0]
df, npts = (tr.stats.sampling_rate, tr.stats.npts)
tr.data = tr.data.astype('float64') #convert to double

# Do the fourier transformation
tr.data = detrend(tr.data, 'linear') #detrend
tr.data *= np.hanning(npts)          #window
# for the real fft we pad n=4*npts points, MODIFY HERE
fdat = np.fft.rfft(tr.data, n=4*npts)

# Load the eigenmodes
eigen = np.loadtxt("eiglst.txt", usecols=[0,1,2,3], converters={1:stringToBool})
# load only the S part
ind1 = eigen[:,1].astype(bool)
ind2 = eigen[:,0]
ind = ((ind2 == 0) & ind1) #bitwise comparing for bool arrays
modes = eigen[ind,3]/1000  #normalize, freq given in mHz

# Take the first 2% only
N = len(fdat)*0.02
freq = np.linspace(0,df/2,len(fdat))
freq = freq[1:N+1] #zero frequency is offset
fdat = fdat[1:N+1]
fdat /= abs(fdat).max() #normalize to 1

# Plot the results
plt.clf()
plt.plot(freq,abs(fdat))
plt.vlines(modes[0:len(modes)/2],0.001,1)
plt.xlabel('Freq [HZ]')
plt.show()