#!/usr/bin/env python

## Alright people, prepare yourself.

## I am going to:

## Populate my mseed files with event and station information

## Write them to SaC files

## Time shift the SAC files (I KNOW RIGHT)

## Remove the instrument response, detrend, remove mean, filter etc...

## ROTATE THE TRACES

## Then leave them all in the "processed" directory

## Good stuff

#### -----This needs to be run from the "raw" directory as given in the ObspyDMT directory structure ----- ####

## import useful packages:

import obspy

import os

from glob import *

## To take in arguments from the terminal

import argparse

# To get relative path

from pathlib import Path

## to only read certain lines in file:

from itertools import islice

##distance baz calculator

from utilities_Array import deg_km_az_baz

# To dictate headers

from obspy.core.util import AttribDict

## Web service packages

from obspy.clients.iris import Client

# TauP

from obspy.taup import TauPyModel

# my model of choice is prem :D

model = TauPyModel(model="prem")

# Make client - just for giggles

client = Client()

import shutil

import numpy

import matplotlib.pyplot as plt

import argparse

parser = argparse.ArgumentParser(description='Preprocessing script for data retrieved from obspy DMT')

parser.add_argument("-fmin","--minimuum_frequency", help="Enter the lowest frequency to be bandpass filtered.", type=float, required=False, action="store", default = 0.05)

parser.add_argument("-fmax","--maximuum_frequency", help="Enter the highest frequency to be bandpass filtered.", type=float, required=False, action="store", default = 1.0)

parser.add_argument("-p","--phase_list", help="Enter the phases you want the travel time predictions for.", type=str, nargs='+', required=True, action="store", default=[])

parser.add_argument("-tap","--taper_percentage", help="Enter the percentage of the trace at each end you want to taper as a decimal.", type=float, required=False, action="store", default=0.05)

parser.add_argument("-rot","--rotate", help="Enter whether you want the traces to be rotated or not.", required=False, action="store_true", default=False)

parser.add_argument("-split","--splitting", help="Enter whether you want the traces stored in a file and be compativle with SHEBA.", type=bool, required=False, action="store", default=False)

parser.add_argument("-detrend","--detrend_T_F", help="Enter whether you want to detrend the trace or not.", type=bool, required=False, action="store", default=True)

parser.add_argument("-v", "--verbose", help="Increase verbosity of output, just --verbose is enough to turn on.",action="store_true")

args = parser.parse_args()

min_freq = args.minimuum_frequency

max_freq = args.maximuum_frequency

phases = args.phase_list[:]

taper_percent = args.taper_percentage

Q_rotate = args.rotate

Q_detrend = args.detrend_T_F

Q_splitting = args.splitting

major=True

print(min_freq, max_freq,phases,taper_percent,Q_rotate,Q_detrend)

### Find the files with the station and event information in them

## Event dir name:

ev_relpath = "EVENTS-INFO/catalog.txt"

## full path of current directory

pwd = os.getcwd()

## path to directory with the Event directory

base_path = str(Path(pwd).parents[1]) + "/"

ev_path = os.path.join(base_path, ev_relpath)

## path to directory with station XML files :D

## gets the dierctory outside the current one:

base_dir = os.path.split(pwd)[0]

station_dir = os.path.join(base_dir,"resp/")

#print(station_dir)

## get event id

ev_name = os.path.basename(base_dir)

#print(ev_name)

## make a list for the stations

station_temp = []

### BEGIN THE LOOP

for M_file in glob('*BH*'):

print(M_file)

bits = M_file.split('.')

print('bits', len(bits))

station_name=bits[1]

network_name=bits[0]

### stream the file

stream = obspy.read(M_file)

## extract the trace

trace = stream[0]

## copy the trace for no real reason.

s_time = trace.stats.starttime

e_time = trace.stats.endtime

trace2= trace.trim(starttime=s_time)

## open the xml file as an inventory

inv = obspy.read_inventory("%sSTXML.%s" %(station_dir, M_file))

network=inv[0]

station=network[0]

channel=station[0]

## For each of the channels the station has, if the code (BHZ BHE BHN) is the same

# take the information

# this was the bane of my existance for an hour

for chnl in station:

if chnl.code==trace.stats.channel:

channel_name = chnl.code

response=channel.response

dip= chnl.dip

inclination = chnl.dip + 90

## take the coordinates of the station

stla=station.latitude

stlo=station.longitude

stel=station.elevation

station_temp.append(station.code)

## MORE SANITY CHECKS

with open(ev_path) as f:

## only read lines from line 6

## This will remove all of the nonesense at the start of the file

for line in islice(f,5,None):

if not line.startswith('-') and ev_name in line:

number,event_id,datetime,latitude,longitude,depth,magnitude,magnitude_type,author,flynn_region,mrr,mtt,mpp,mrt,mrp,mtp,stf_func,stf_duration,t1,t2 = line.split(',')

evla = latitude

evlo = longitude

evdp = depth

mag = magnitude

event_time = obspy.UTCDateTime(datetime)

#distances = client.distaz(stalat=stla, stalon=stlo, evtlat=evla, evtlon=evlo)

dist_deg,dist_km,az,baz = deg_km_az_baz(lat1=float(evla),lon1=float(evlo),lat2=float(stla),lon2=float(stlo))

#### Want to find the origin time relative to the trace period

### i.e. we want to find the difference between the event and

## the start of the trace

O = obspy.UTCDateTime(event_time) - obspy.UTCDateTime(trace.stats.starttime)

### I think O is equal to zero...

### Going to need to loop over all of the phases ###

## I don't think this is the most efficient way of doing things but I'm ##

## going to make a list of lists, where i have [phase name, travel time, number (to be put on t) ##

# good luck my friend #

ph_time_list = []

#### Populate information into SAC part ####

trace2.stats.sac = AttribDict({'baz': float(baz), 'az': float(az), 'kstnm': str(station),'cmpinc': inclination, 'gcarc': str(dist_deg),

'dist': str(dist_km), 'evla': evla, 'evlo': evlo, 'evdp': evdp, 'stla': stla, 'stlo': stlo, 'stel': stel,'o': O,

'stdp': dip, 'iztype':9, 'mag': mag})

for t,phase in enumerate(phases):

time_prediction = model.get_travel_times(source_depth_in_km=float(evdp), distance_in_degree=float(dist_deg), receiver_depth_in_km=0.0, phase_list=[phase])

## if the phase is present add the origin time to it

if time_prediction:

time_pred_total = O + time_prediction[0].time

else:

time_pred_total = None

print(phase, time_pred_total)

### add these values to a list ###

ph_time_list.append([phase, time_pred_total])

tn = str('t'+str(t+1))

ktn = str('k'+tn)

print(tn,ktn)

# i.e. for SKKS major arc.

if major == True:

if len(time_prediction) > 1:

# only use the first major arc arrival

time_pred_total_2 = O + time_prediction[1].time

phase_label2 = "%s_Major" %phase

ph_time_list.append([phase_label2, time_pred_total_2])

### for each loop se tthe attributes that I want ###

# set predicted travel time

for t,pt_list in enumerate(ph_time_list):

time_pred = pt_list[1]

phase_label = pt_list[0]

tn = str('t'+str(t+1))

ktn = str('k'+tn)

setattr(trace2.stats.sac, tn, time_pred)

setattr(trace2.stats.sac,ktn, phase_label)

if trace2.stats.channel == "BHZ":

trace2.stats.sac.cmpaz = 0.0

print(trace2.stats)

#### process the data, filter etc...

trace_rr = trace2.copy()

#question=input("Velocity(v) or Displacement(d)?")

#trace_rr.remove_response(inventory=inv, output='VEL')

trace_rr.remove_response(inventory=inv, output='DISP')

trace_detrend = trace_rr.copy()

if Q_detrend:

#### detrend/demean stuff

trace_detrend.detrend(type='demean')

#### SO, SO, SO, SO. I'll now filter the trace with the response removed! - Cheers Josh

#### taper it

trace_taper = trace_detrend.copy()

# This gets the same results as the SAC function taper.

trace_taper.taper(max_percentage=taper_percent) ## Defaults to both sides and 'hann' type.

trace_filt = trace_taper.copy()

#trace_filt.filter('bandpass', freqmin=0.05, freqmax=0.3, corners=4, zerophase=True)

trace_filt.filter('bandpass', freqmin=min_freq, freqmax=max_freq, corners=4, zerophase=True)

trace_filt.write("%s.%s.%s.SAC" %(network_name,station_name,channel_name), format="SAC")

print("------------Next file -------------")

## Time shift the SAC file

import obspy.io.sac.sactrace

for sac_file in glob('*SAC'):

sac=obspy.io.sac.sactrace.SACTrace.read(sac_file)

FILE_NAME = os.path.splitext(sac_file)

print(FILE_NAME[0])

sac.reftime= event_time

print(sac.t1, sac.b, sac.o)

tr2 = sac.to_obspy_trace()

print(tr2.stats)

tr2.write("%s.SAC" %FILE_NAME[0], format="SAC")

print("YAY - It worked, I now have SAC files... WITH POPULATED INFORMATION")

print("I hope this is living up to your expectations")

print("If not, you can leave and write your own script")

print("Then send me it.")

print("It'll be better and more efficient than mine...")

## make a list with only no repeated station name

stations = list(set(station_temp))

print(stations)

# Q_rotate = True

#### ROTATE! ####

print("rotate:", Q_rotate)

if Q_rotate == True:

## Now we rotate ## if the user wants to ##

## use the lists made at the start of the code

for station in stations:

## print station, because, why not?

print(station)

## make temporary list for the sac files

temp_list = []

## loop over all the sac files, because efficiency is for losers

for sac_file in glob('*SAC'):

## print the name, because, see above comment about efficiency and losers

## take the BHN and BHE files and shove them in the temporary list

if "%s." %station in sac_file and 'BHN' in sac_file:

temp_list.append(sac_file)

if "%s." %station in sac_file and 'BHE' in sac_file:

temp_list.append(sac_file)

#print(temp_list)

## make a stream of those two lists only

stream = obspy.read(temp_list[0])

stream += obspy.read(temp_list[1])

## find back azimuth and inclination

baz = stream[0].stats.sac.baz

inc = stream[0].stats.sac.cmpinc

### check if they are the same length, if not, trim the longer one by the difference in data points

npts1 = stream[0].stats.npts

npts2 = stream[1].stats.npts

samp_rate = stream[0].stats.delta

s_time1 = stream[0].stats.starttime

s_time2 = stream[1].stats.starttime

e_time1 = stream[0].stats.endtime

e_time2 = stream[1].stats.endtime

# print(stream)

## correct for different trace start times

if s_time1 != s_time2:

if s_time1 > s_time2:

stream[0].trim(starttime=s_time1)

stream[1].trim(starttime=s_time1)

if s_time1 < s_time2:

stream[0].trim(starttime=s_time2)

stream[1].trim(starttime=s_time2)

if e_time1 != e_time2:

if e_time1 > e_time2:

stream[0].trim(endtime=e_time2)

stream[1].trim(endtime=e_time2)

if e_time1 < e_time2:

stream[0].trim(endtime=e_time1)

stream[1].trim(endtime=e_time1)

stream.rotate(method = 'NE->RT', back_azimuth = baz, inclination = inc)

### Get the traces from the stream ###

trace1 = stream[0]

trace2 = stream[1]

s_time3 = trace1.stats.starttime

s_time4 = trace2.stats.starttime

e_time3 = trace1.stats.endtime

e_time4 = trace2.stats.endtime

## correct for different trace start times

if s_time3 != s_time4:

if s_time3 > s_time4:

trace1.trim(starttime=s_time3)

trace2.trim(starttime=s_time3)

if s_time3 < s_time4:

trace1.trim(starttime=s_time4)

trace2.trim(starttime=s_time4)

if e_time3 != e_time4:

if e_time3 > e_time4:

trace1.trim(endtime=e_time4)

trace2.trim(endtime=e_time4)

if e_time1 < e_time2:

trace1.trim(endtime=e_time3)

trace2.trim(endtime=e_time3)

print(stream)

ch1 = trace1.stats.channel

ch2 = trace2.stats.channel

print(ch1,ch2)

if ch1[-1] == "R":

cmpaz_1 = float(trace1.stats.sac.baz) + 180

elif ch1[-1] == "T":

cmpaz_1 = float(trace1.stats.sac.baz) + 270

else:

pass

if ch2[-1] == "R":

cmpaz_2 = float(trace1.stats.sac.baz) + 180

elif ch2[-1] == "T":

cmpaz_2 = float(trace1.stats.sac.baz) + 270

else:

pass

### write component information into the file

trace1.stats.sac['cmpaz'] = float(cmpaz_1)

trace2.stats.sac['cmpaz'] = float(cmpaz_2)

print("CMPAZ: ", trace1.stats.sac.cmpaz)

###### get info from the name of the first trace #####

network_name, station_name,channel_name,SAC = temp_list[0].split(".")

## write the dude ##

trace1.write("%s.%s.%s.SAC" %(str(trace1.stats.network), station_name,str(trace1.stats.channel)), format = "SAC")

trace2.write("%s.%s.%s.SAC" %(str(trace2.stats.network), station_name,str(trace2.stats.channel)), format = "SAC")

## Store the BHN and BHE in case I want them for whatever reason.

# make directory

H_E_Dir="../processed/BHN_BHE_processed"

if not os.path.exists(H_E_Dir):

os.makedirs(H_E_Dir)

# move files

for old_file in glob('*BHN*.SAC'):

if os.path.exists(str(H_E_Dir) + "/" + str(old_file)):

os.remove(str(H_E_Dir) + "/" + str(old_file))

shutil.move(old_file,H_E_Dir)

for old_file in glob('*BHE*.SAC'):

if os.path.exists(str(H_E_Dir) + "/" + str(old_file)):

os.remove(str(H_E_Dir) + "/" + str(old_file))

shutil.move(old_file,H_E_Dir)

## should be left with only BHR/BHT/BHZ files

### We are, in fact, not done...

## need to move the files to the "processed" directory

if Q_splitting:

### Going to do my splitting work within a seperate directory

Splitting_Dir = "../Splitting"

if not os.path.exists(Splitting_Dir):

os.makedirs(Splitting_Dir)

else:

pass

for pro_file in glob('*SAC'):

pro_file_First = pro_file.split(".")[0]

pro_file_Second = pro_file.split(".")[1]

pro_file_no_SAC = pro_file_First + "." + pro_file_Second

if os.path.exists("../processed/%s" %pro_file):

os.remove("../processed/%s" %pro_file)

shutil.copy(pro_file, "../processed")

if Q_splitting:

if os.path.exists("%s/%s" %(Splitting_Dir,pro_file)):

os.remove("%s/%s" %(Splitting_Dir,pro_file))

shutil.move(pro_file, Splitting_Dir)

os.rename("%s/%s" %(Splitting_Dir, pro_file), "%s/%s" %(Splitting_Dir,pro_file_no_SAC))

else:

pass