#!/usr/bin/python3

import os
from obspy import Stream,UTCDateTime,read_events,read_inventory
from obspy.clients.fdsn import Client
from obspy.geodetics.base import locations2degrees
from obspy.taup import TauPyModel

model = TauPyModel(model="iasp91") #1D velocity model to predict arrivals, list of all options here: https://docs.obspy.org/packages/obspy.taup.html

auspass = Client('AUSPASS') #establish our connection to AusPass FDSN server


################################## enter parameters here
starttime=UTCDateTime(2016,11,9)
endtime=UTCDateTime(2016,11,15)

minmag=6.5
maxmag=10

mindist = 30 #degrees
maxdist = 90

network='1Q'
station='AQTG*'
location='*'
channel='*Z'

time_before_p = 60 #seconds of data to download before predicted p arrival
time_after_p  = 120

###################################

#STEP 1 > get station metadata from auspass (make sure to delete old file if changing parameters!)
if not os.path.exists('%s.xml' % network):
	print("downloading stationXML inventory...")
	inventory = auspass.get_stations(network=network,station=station,level='response',channel=channel,starttime=starttime,endtime=endtime)
	inventory.write('%s.xml' % network,format='STATIONXML')
else:
	print("loading pre-existing stationXML file %s.xml" % network)
	inventory = read_inventory('%s.xml' % network)


#STEP 2 > get an events catalog within the required magnitude and time limits, and save as a QuakeML file (make sure to delete old file if changing parameters!)
if not os.path.exists('events.%s.qml' % network):
	print("downloading events...")
	eventserver = Client('ISC') # shortcut to http://www.isc.ac.uk/iscbulletin/search/catalogue/. other options are 'IRIS', 'USGS' etc
	catalog = eventserver.get_events(minmagnitude=minmag, maxmagnitude=maxmag,starttime=starttime,endtime=endtime)
	catalog.write('events.%s.qml' % network,format='QUAKEML'); print('events.%s.qml written' % network)
else:
	print("loading pre-existing qml file events.%s.qml" % network)
	catalog = read_events('events.%s.qml' % network)

if not catalog: print("no events found! exiting"); raise SystemExit

#STEP 3 > iterate through each station / network in inventory, 
#  then iterate through each eq in catalog, testing if they're with the desired distance range
#    Then calculate the predicted p-arrival time at the station from each event
#       THEN download the waveform data within the desired window around this predicted arrival  

for net in inventory:
	for eq in catalog:
		eq_lat,eq_lon,eq_time = eq.origins[0].latitude, eq.origins[0].longitude, eq.origins[0].time
		print("finding data in network %s for event at %s %.4f lat %.4f long " % (network,eq.origins[0].time,eq.origins[0].latitude, eq.origins[0].longitude))

		output_stream = Stream() #we're going to write files per EVENT that include all stations. so place this within event loop, but outside of station (sta) loop
		
		for sta in net:	
			lat,lon = sta.latitude,sta.longitude
			dist_station_to_eq = locations2degrees(lat,lon,eq_lat,eq_lon)

			if mindist <= dist_station_to_eq <= maxdist:
				#calculate theoretical p arrival
				arrivals = model.get_travel_times(source_depth_in_km=eq.origins[0].depth/1000,distance_in_degree=dist_station_to_eq,phase_list='P')
				p_duration = arrivals[0].time #seconds it takes for p-wave to reach station
				#so, estimated time p-wave arrival time at station is the eq origin time + the time it takes for the wave to travel to station 
				p_arrival_time = eq.origins[0].time + p_duration

				#now that we know when the p-wave should arrive, we can download a window of data around it and add it to our output stream
				#we will try/except this in case the request returns no data (or some other error occurs!)
				try:
					st = auspass.get_waveforms(network=net.code,station=sta.code,location=location,channel=channel,
						                       starttime=p_arrival_time-time_before_p,endtime=p_arrival_time+time_after_p)
					output_stream += st #add to our output!
				except:
					print("no data for %s.%s.%s.%s between %s-%s...skipping" % (net.code,sta.code,location,channel,p_arrival_time-time_before_p,p_arrival_time+time_after_p))
					continue
				
		#At this point we can write out our collected data PER EVENT. you could also produce a file per station per event, or combine all data in to one (etc), with the proper edits
		if len(output_stream) > 0: 
			#if you'd like to remove the response before writing, it's as easy as: 
			output_stream.remove_response(inventory)

			output_stream.write("%.4d.%.2d.%.2d.%.2d.%.2d.%.2d.ms" % (eq_time.year,eq_time.month,eq_time.day, \
															eq_time.hour,eq_time.minute,eq_time.second),format='MSEED') #could also write as SAC files, or many other formats... 
			print("wrote %.4d.%.2d.%.2d.%.2d.%.2d.%.2d.ms" % (eq_time.year,eq_time.month,eq_time.day, \
															eq_time.hour,eq_time.minute,eq_time.second))