# %%
import numpy as np
import datetime
import obspy
import io
import os
import sys
import glob
from obspy import read, Trace, Stream, UTCDateTime,read_inventory
import matplotlib.pyplot as plt
import pandas as pd



#=========================================================================================
if len(sys.argv) < 5:
    print("    ")
    print("USAGE: python BDSN_get_sac.py [net] [sta] [year] [start jday] [duration (day) \"30\"]")
    print("    ")
    sys.exit()
#=========================================================================================

net        = sys.argv[1]
stn        = sys.argv[2]
year       = sys.argv[3]
jday       = sys.argv[4]
duration   = sys.argv[5]


# %%
#year = int(yr)
#jday = "001"
#duration = int(dur) # unit in day
#net = "BK"
#sta  = "BKS" # HOPS, KCC, ORV, CMB
loc = "00"  # to specify the instrument at the station
channels = ['LHE', 'LHN', 'LHZ']
datetime_obj = pd.to_datetime(f'{year}-{jday}', format='%Y-%j')



# Define the duration of each segment (7200 seconds)
segment_duration = 7200

# Calculate the number of segments
num_segments = 12
num_segments = int(np.floor(int(duration)*86400 / (segment_duration)))

start_time = datetime.datetime.strptime(year+jday, '%Y%j').date()

eventTime = UTCDateTime(str(start_time))
starttime = eventTime
endtime = eventTime + 86400 * float(duration)

for channel in channels:
    data_path=f"/data/dc15/BSL_TEST_DATA/BK/{year}/{year}.{jday}/{stn}.{net}.{channel}.00.D.{year}.{jday}"
    #print(data_path)
    st = read(data_path)
    for i in range(int(jday),endtime.julday,1):
        data_path2=f"/data/dc15/BSL_TEST_DATA/BK/{year}/{year}.{i:03d}/{stn}.{net}.{channel}.00.D.{year}.{i:03d}"
        st += read(data_path2) 

    st.merge(method=1)    
    inv = read_inventory('/home/gcl/BR/liwei/DAS01/DATA/BK.MBARI/station_MBARI.xml')
    st.detrend(type='linear')
    st.detrend(type='demean')
    pre_filt = [0.0005, 0.001, 5, 10]
    st.remove_response(inventory=inv, pre_filt=pre_filt, output="VEL")
    #st.decimate(4, no_filter=False)
    st.detrend(type='linear')
    st.detrend(type='demean')


    # Loop over each segment and save it as a new SAC file
    for i in range(num_segments):
        # Calculate the start and end times of the current segment
        segment_start = starttime + i * segment_duration
        segment_end = segment_start + segment_duration
        
        # Cut the Trace to the current segment
        segment = st.slice(segment_start, segment_end)
        
        # Construct the output file name
        hour=segment_start.hour
        year=segment_start.year
        jjday=segment_start.julday
        output_dir = f"SAC_TRACES/{stn}/{year}.{jjday:03d}"

        # Create the output directory if it doesn't exist
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        output_file = f"{output_dir}/{year:04d}.{jjday:03d}.{hour:02d}.00.{stn}.{channel}.sac"
        
        # Save the segment as a new SAC file
        segment.write(output_file, format='SAC')