#!/usr/bin/env python
import numpy as np
import datetime
import obspy
import io
import os
import sys
import glob
from obspy import read, Trace, Stream, UTCDateTime
from obspy.clients.fdsn import Client
import matplotlib.pyplot as plt

from obspy.clients.fdsn.header import URL_MAPPINGS
URL_MAPPINGS['NCEDC'] = "https://service.ncedc.org"

def stream2array(stream):
    '''
    populates a 2D np.array that is the traces as rows by the samples as cols
    '''
    import numpy as np
    import obspy
    array=np.empty((len(stream),len(stream[0].data)),dtype=float) # initialize
    for index,trace in enumerate(stream):
        array[index,:]=trace.data
    return array




def spectra_local(stream,units_input='strain',units_output='strain',kind='psd',trace=99999):
    
    
    # !!!!!original code assumes input is strain rate
    # 

    import matplotlib.pyplot as plt
    import numpy as np
    from scipy.integrate import cumtrapz
    
    print('input unit : '+units_input)
    print('output unit : '+units_output)

    
    if units_input == units_output:
        
        print('same in/output unit')
        oldstream = stream.copy()
        stream = oldstream
        
    else:
        
        if units_input=='strain-rate' and units_output=='strain':
            
            # integration
            print('Integrating')
            
            oldstream = stream.copy()
            oldstream.detrend(type='linear')
            oldstream.detrend(type='demean')
            oldstream.taper(0.05)
            oldstream.integrate()
            #oldstream.integrate(method='spline')
            oldstream.detrend(type='linear')
            oldstream.detrend(type='demean')
            stream = oldstream
            
    
    
    if trace==99999:
        A = stream2array(stream)*0
        nx,nt = np.shape(A)
        nyq = nt//2
        dt = stream[0].stats.delta
        A = A[:,slice(1,nyq)]
        for i,tr in enumerate(stream.copy()):
            fr = np.fft.fftfreq(nt,d=dt)[slice(1,nyq)]
            #A[i,:] = np.fft.fft(tr.data/1e9)[slice(1,nyq)] # convert from ne/s or ne to e/s or e
            A[i,:] = np.fft.fft(tr.data)[slice(1,nyq)]
        if kind=='as':
            sp = np.abs(A.T)
            sp = 10*np.log10(sp)
        if kind=='ps':
            sp = 2*(np.abs(A.T)**2) / (nt**2)
            sp = 10*np.log10(sp)
        if kind=='psd':
            sp = 2*(np.abs(A.T)**2) / (nt**2) * dt * nt
            sp = 10*np.log10(sp)

    else:
        nt = len(stream[trace].data)
        nyq = nt//2
        dt = stream[trace].stats.delta
        fr = np.fft.fftfreq(nt,d=dt)[slice(1,nyq)]
        #A = np.fft.fft(stream[trace].data/1e9)[slice(1,nyq)] # convert from ne/s or ne to e/s or e
        A = np.fft.fft(stream[trace].data)[slice(1,nyq)]
        if kind=='as':
            sp = np.abs(A)
            sp = 10*np.log10(sp)
        if kind=='ps':
            sp = 2*(np.abs(A)**2) / (nt**2)
            sp = 10*np.log10(sp)
        if kind=='psd':
            sp = 2*(np.abs(A)**2) / (nt**2) * dt * nt
            sp = 10*np.log10(sp)

    return fr,sp


            
#=========================================================================================
if len(sys.argv) < 5:
    print("    ")
    print("USAGE: python ms_spec_v2.py [year] [jday] [net] [station] [channel]")
    print("    ")
    sys.exit()
#=========================================================================================
year = sys.argv[1]
jday = sys.argv[2]
net = sys.argv[3]
sta  = sys.argv[4]
chan  = sys.argv[5]
loc = "00"  # to specify the instrument at the station
print(year, jday)

work_dir='/home/gcl/BR/liwei/DAS01/tmp/00.tmp/'+year+'.'+jday
start_time = datetime.datetime.strptime(year+jday, '%Y%j').date()
client = Client("NCEDC")

eventTime = UTCDateTime(str(start_time))
starttime = eventTime    
endtime = eventTime + 86400


myStream = client.get_waveforms(
net, sta, loc, chan, starttime, endtime, attach_response=True
)

st = myStream.copy()
st.detrend(type='linear')
st.detrend(type='demean')
st.taper(0.01)
pre_filt = [0.0005, 0.001, 5, 10]
st.remove_response(pre_filt=pre_filt, output="VEL")
st.decimate(factor=2,no_filter=True)
st.detrend(type='linear')
st.detrend(type='demean')
st.taper(0.01)
st.filter('lowpass', freq=9, corners=6, zerophase=True)
fr,psd = spectra_local(st,units_input='strain',units_output='strain',kind='psd',trace=0)
psdOUT = work_dir+'/'+year+'_'+jday+'_'+sta+'.'+net+'_'+chan+'_spec.npy'
np.save(psdOUT, np.c_[1/fr,psd])
sys.exit()