from obspy.clients.fdsn import Client
import numpy as np
import obspy as op
from obspy.core import UTCDateTime, read, Stream
import os
from obspy.io.sac.sactrace import SACTrace


print("# np_version = ",np.__version__)
print("# op version = ",op.__version__)


# specify the event id
eventid_ncedc = EVENTID
#eventid_ncedc = 73160480 # example is M2.5 event https://earthquake.usgs.gov/earthquakes/eventpage/nc73160480/executive
#eventid_ncedc = 71093184
#eventid_ncedc = 21377014

# select request_server
request_server = REQUEST_SERVER # 1 is ncedc 2 is iris 3 is scedc
#request_server = 1 # Hayward fault project -> 1 

if request_server == 1:
    net_listOUT = "BK,NC,PB"
    com_listOUT = "HHZ,EHZ,DP1,EP1"
    server_nameOUT = "ncedc"
    clientOUT = Client("NCEDC")


if request_server == 2:
    net_listOUT = "PB,AZ"
    com_listOUT = "EHZ,HHZ"
    server_nameOUT = "iris"
    clientOUT = Client("IRIS")


if request_server == 3:
    net_listOUT = "CI,AZ"
    com_listOUT = "EHZ,HHZ"
    server_nameOUT = "iris"
    clientOUT = Client("SCEDC")



#distdeg_from_eq = 0.4
#distdeg_from_eq = 0.2
# search distance from event in degrees
distdeg_from_eq = 0.35

# need for catalog
#client_scedc = Client("SCEDC")
# extracting event info. from datacenter catalog
client_ncedc = Client("NCEDC",timeout=600)
print(client_ncedc)  

# band pass when we remove instrument response
# f1<f2<f3<f4. effective band is between f2 and f3
f1 = 0.1 # 0.1 Hz
f2 = 0.5  # 0.5 Hz
f3 = 40 # 50 Hz
f4 = 50 # 100 Hz

# time window from origin time to extract waveform
pre_tw_sec = 90 # 90 seconds before the origin time
tw_sec = 5*60 # 5 min after the origin time


catalog_ncedc = client_ncedc.get_events(eventid=eventid_ncedc)
print("# catalog_ncedc = ",catalog_ncedc)

# event info. origin time, location, magnitude
event = catalog_ncedc[0]
origin = event.origins[0]
origin_time = origin.time
evlat = origin.latitude
evlon = origin.longitude
evdp_km = origin.depth / 1000
evmag = event.magnitudes[0].mag

evyearOUT = origin_time.year
evjdayOUT = origin_time.julday
evhourOUT = origin_time.hour
evminOUT = origin_time.minute
evsecOUT = origin_time.second


# need for file name
evyearOUT2 = (str)(evyearOUT)
evjdayOUT2 = (str)(evjdayOUT)
if evjdayOUT < 100:
    evjdayOUT2 = "0"+(str)(evjdayOUT)
        

if evjdayOUT < 10:
    evjdayOUT2 = "00"+(str)(evjdayOUT)
    

evhourOUT2 = (str)(evhourOUT)
if evhourOUT < 10:
    evhourOUT2 = "0"+(str)(evhourOUT)
        

evminOUT2 = (str)(evminOUT)
if evminOUT < 10:
    evminOUT2 = "0"+(str)(evminOUT)


evsecOUT2 = (str)(evsecOUT)
if evsecOUT < 10:
    evsecOUT2 = "0"+(str)(evsecOUT)
        


print("# evyearOUT2 = ",evyearOUT2," evjdayOUT2 = ",evjdayOUT2," evhourOUT2 = ",evhourOUT2," evminOUT2 = ",evminOUT2," evsecOUT2 = ",evsecOUT2)
    
evmseedid = evyearOUT2+"."+evjdayOUT2+"."+evhourOUT2+""+evminOUT2+""+evsecOUT2
print("# evmseedid = "+evmseedid)

# name for output directory
pwd_dir = os.getcwd() 
#sacdir= pwd_dir +"/"+ evmseedid +"_M"+(str)(evmag)+"_"+(str)(eventid_ncedc)+"_fl"+(str)(f2)+"_fh"+str(f3)+"_dist"+(str)(distkm_from_eq)+"km"
sacdir= pwd_dir +"/"+ evmseedid +"_M"+(str)(evmag)+"_"+(str)(eventid_ncedc)+"_fl"+(str)(f2)+"_fh"+str(f3)+"_dist"+(str)(distdeg_from_eq)+"deg"
print("# sacdir = ",sacdir)

# create output directory
if not os.path.exists(sacdir):
    os.makedirs(sacdir)


latitude = evlat
longitude = evlon
minradius = 0
maxradius = distdeg_from_eq
starttime = UTCDateTime(origin_time-pre_tw_sec)
endtime = UTCDateTime(origin_time+tw_sec)

#starttime2 = UTCDateTime(origin_time-86400)
#endtime2 = UTCDateTime(origin_time+86400)

starttime2 = UTCDateTime(origin_time-pre_tw_sec)
endtime2 = UTCDateTime(origin_time+tw_sec)

print("# starttime = ",starttime," endtime = ",endtime)
print("# starttime2 = ",starttime2," endtime2 = ",endtime2)

yearOUT = starttime.year
jdayOUT = starttime.julday
hourOUT = starttime.hour
minOUT = starttime.minute
secOUT = starttime.second

yearOUT2 = (str)(yearOUT)
    
jdayOUT2 = (str)(jdayOUT)
if jdayOUT < 100:
    jdayOUT2 = "0"+(str)(jdayOUT)
        

if jdayOUT < 10:
    jdayOUT2 = "00"+(str)(jdayOUT)
        

hourOUT2 = (str)(hourOUT)
if hourOUT < 10:
    hourOUT2 = "0"+(str)(hourOUT)
        

minOUT2 = (str)(minOUT)
if minOUT < 10:
    minOUT2 = "0"+(str)(minOUT)


secOUT2 = (str)(secOUT)
if secOUT < 10:
    secOUT2 = "0"+(str)(secOUT)
        


print("# yearOUT2 = ",yearOUT2," jdayOUT2 = ",jdayOUT2," hourOUT2 = ",hourOUT2," minOUT2 = ",minOUT2," secOUT2 = ",secOUT2)
    
mseedid = yearOUT2+"."+jdayOUT2+"."+hourOUT2+""+minOUT2+""+secOUT2
print("# mseedid = "+mseedid)


# this would be only need for SCEDC. SSAF project
# change all m type to this. we have error "SoCal ML"
#catalog_scedc[0].magnitudes[0].method_id = "smi:ci.anss.org/magnitude/CISNml2"

# write event information in QUAKEML format
cat_fi = sacdir + "/" + (str)(eventid_ncedc)+".eq.xml"
print("# cat_fi = "+cat_fi)
catalog_ncedc.write(cat_fi, format="QUAKEML")  

print("# invOUT server_nameOUT = ",server_nameOUT)
print("# net_listOUT = ",net_listOUT," com_listOUT = ",com_listOUT)

# extracting stations within min and max radius, also get instrument response
#invOUT = clientOUT.get_stations(network=net_listOUT, station="*", channel=com_listOUT, 
#                                starttime=starttime, endtime=endtime,
#                                latitude=latitude, longitude=longitude,
#                                minradius=minradius, maxradius=maxradius,
#                                level="response")

# use time2 +/- 1 days to fetch multiple time epochs if we have. there is a bug for FDSN code cannot handle the epoch change 05.07.19
invOUT = clientOUT.get_stations(network=net_listOUT, station="*", channel=com_listOUT, 
                                starttime=starttime2, endtime=endtime2,
                                latitude=latitude, longitude=longitude,
                                minradius=minradius, maxradius=maxradius,
                                level="response")

channelOUT = clientOUT.get_stations(network=net_listOUT, station="*", channel=com_listOUT, 
                                starttime=starttime2, endtime=endtime2,
                                latitude=latitude, longitude=longitude,
                                minradius=minradius, maxradius=maxradius,
                                level="channel")



# get station list
ch_listOUT = ""
j2=0
sncl_list = []
for i in range(len(channelOUT)):        
    netOUT = channelOUT[i].code
    #print (inv_ncedc[i].code)
    for ii in range(len(channelOUT[i])):
        #print(ii)
        #print(inv[i][ii].code)    
        staOUT = channelOUT[i][ii].code
        for iii in range(len(channelOUT[i][ii])):
            chaOUT = channelOUT[i][ii][iii].code 
            locOUT = channelOUT[i][ii][iii].location_code
            snclOUT = netOUT+"."+staOUT+"."+locOUT+"."+chaOUT
            print("# j2 = ",j2," snclOUT = ",snclOUT)
            sncl_list.append(snclOUT)
            j2+=1
           



# write station metadata in STATIONXML format
#inv_fiOUT = sacdir + "/" + (str)(eventid_ncedc)+".sta."+(str)(distdeg_from_eq)+"km."+server_nameOUT+".xml"
inv_fiOUT = sacdir + "/" + (str)(eventid_ncedc)+".sta."+(str)(distdeg_from_eq)+"deg."+server_nameOUT+".xml"
print("# inv_fiOUT = "+inv_fiOUT)
invOUT.write(inv_fiOUT, format='STATIONXML')  

inv_fiOUT = sacdir + "/" + (str)(eventid_ncedc)+".sta."+(str)(distdeg_from_eq)+"deg."+server_nameOUT+".sacpz"
print("# inv_fiOUT = "+inv_fiOUT)
invOUT.write(inv_fiOUT, format='SACPZ')  

# get station list
sta_listOUT = ""
j=0
for i in range(len(invOUT)):
    #print (inv_ncedc[i].code)
    for ii in range(len(invOUT[i])):
        #print(ii)
        #print(inv[i][ii].code)
        if j==0:
                sta_listOUT = invOUT[i][ii].code
        else:
                sta_listOUT = sta_listOUT+","+invOUT[i][ii].code
        j+=1
    


print("# sta_listOUT = ",sta_listOUT)

# get seismic waveforms nased on station & network lists
stOUT = clientOUT.get_waveforms(network=net_listOUT, station=sta_listOUT, channel=com_listOUT, location="*",
                                starttime=starttime, endtime=endtime)

stOUT2 = stOUT.copy()

#st2 = stOUT2.select(network="NC")

print(stOUT2.__str__(extended=True))

for trace in stOUT2:
    nptsOUT = trace.stats.npts
    deltaOUT = trace.stats.delta
    sampHZOUT = 1.0/deltaOUT
    #print("sampHZOUT = ", sampHZOUT, "nptsOUT = ", nptsOUT)    
    if nptsOUT < sampHZOUT:
        stOUT2.remove(trace)



print(stOUT2.__str__(extended=True))

stOUT22 = stOUT2.copy()
stOUT22.clear()  
for i in range(len(sncl_list)):
    print("# sncl_list[i] = ",sncl_list[i])
    stOUT22 += stOUT2.select(id=sncl_list[i])


stOUT2.clear() 
stOUT2 = stOUT22.copy()
print(stOUT2.__str__(extended=True))

# remove instrument response
print("# remove_response server_nameOUT = ",server_nameOUT)
stOUT2.remove_response(inventory=invOUT, 
                   pre_filt=(f1, f2, f3, f4),
                   output="VEL")

#st2 = stOUT2.select(network="NC",station="N*", location="10")
#print(st2.__str__(extended=True))

#print("# remove_response server_nameOUT = ",server_nameOUT)
#st2.remove_response(inventory=invOUT, 
#                   pre_filt=(f1, f2, f3, f4),
#                   output="VEL")



# save waveforms in SAC format
for i in range(len(stOUT2)):
    #print (st_ncedc2[i])
    sac = SACTrace.from_obspy_trace(stOUT2[i])
    sac.lcalda = True
    
    netOUT = stOUT2[i].stats.network
    comOUT = stOUT2[i].stats.channel
    staOUT = stOUT2[i].stats.station
    locOUT = stOUT2[i].stats.location

    # st_starttime is different from each sacfile, but mseedid is used origin_time - pre_tw_sec
    #st_starttime = st_ncedc2[i].stats.starttime
    #print("# st_starttime = ",st_starttime)
        
    
    stid = netOUT+"."+staOUT+"."+locOUT+"."+comOUT
    print("# staOUT = ",staOUT," netOUT = ",netOUT," comOUT = ",comOUT," locOUT = ",locOUT," stid = ",stid)
    #print("# stid = "+stid)
    sta_coordinate = invOUT.get_coordinates(stid, starttime)
    sac.stla = sta_coordinate['latitude']
    sac.stlo = sta_coordinate['longitude']
    sac.stel = sta_coordinate['elevation']
    sac.evla = evlat
    sac.evlo = evlon
    sac.evdp = evdp_km
     
    sac.kuser0 = "ncss"
    sac.kevnm = (str)(eventid_ncedc)
    sac.mag = evmag
    
    outsacfi = sacdir +"/" + staOUT+"."+netOUT+"."+comOUT+"."+locOUT+".D."+mseedid+"."+(str)(eventid_ncedc)+".sac"
    #print("# outsacfi = "+outsacfi)
    #st_ncedc[i].write(outsacfi, format="SAC")  
    sac.write(outsacfi)