#!/usr/bin/env python
import sys
import numpy as np
import matplotlib.pyplot as plt
import glob
import datetime
import pytz
import obspy
import io
from obspy import read, Trace, Stream, UTCDateTime, read_inventory
from datetime import datetime
import pandas as pd

#=========================================================================================
if len(sys.argv) < 4:
    print("    ")
    print("USAGE: python ms_spec.py [year] [jday] [station] [channel]")
    print("    ")
    sys.exit()
#=========================================================================================


year = sys.argv[1]
jday = sys.argv[2]
stn  = sys.argv[3]
chn  = sys.argv[4]
data_path='/data/dc15/BSL_TEST_DATA/BK/'+year+'/'+year+'.'+jday+'/'+stn+'.'+chn+'.00.D.'+year+'.'+jday
work_dir='/home/gcl/BR/liwei/DAS01/tmp/00.tmp/'+year+'.'+jday
print("# data_path = ", data_path)



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


st = read(data_path)
#inv = read_inventory('/home/gcl/BR/liwei/DAS01/DATA/BK.MBARI/station_MBARI_BHZ.xml')
inv = read_inventory('/work/das01/BR/liwei/DATA/BDSN_station.xml')
st.detrend(type='linear')
st.detrend(type='demean')
st.taper(0.01)
pre_filt = [0.0005, 0.001, 5, 10]
st.remove_response(inventory=inv, 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+'_'+stn+'_'+chn+'_spec.npy'



#np.savetxt(psdOUT, np.c_[1/fr,psd])
np.save(psdOUT, np.c_[1/fr,psd])


sys.exit()