# %%
from obspy import read
from obspy import Inventory

from obspy import read_inventory
from obspy.core.inventory import Inventory, Network, Station, Channel, Site
from obspy import UTCDateTime
from obspy.clients.nrl import NRL


from obspy.signal import PPSD
from obspy.imaging.cm import pqlx
from obspy.imaging.cm import viridis_white_r




# %%
import obspy as ob
print("# obspy version = ", ob.__version__)

# %%
import matplotlib.pyplot as plt
import numpy as np


# %%
import os
import sys

# %%
SMALL_SIZE = 12
MEDIUM_SIZE = 14
BIGGER_SIZE = 16

plt.rc('font', size=SMALL_SIZE)          # controls default text sizes
plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title
plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize
plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title


# %%

# lat, lon ele from BK.BKS
#st_lat = 37.87622
#st_lon = -122.23558
#st_ele = 243.9
#st_dep = 25.6

def make_resp(sta, net, com, loc, st_lat=37.87622, st_lon=-122.23558, st_ele=243.9, st_dep=25.6, azimuth=0, dip=-90, sample_rate=250, db_resp=0):

# Select the preamplifier gain (7 items):
#   '0 dB (1)', '12 dB (4)', '18 dB (8)', '24 dB (16)', '30 dB (32)',
#   '36 dB (64)', '6 dB (2)'

    if db_resp == 0:    
        db_resp_str = '0 dB (1)'

    if db_resp == 12:    
        db_resp_str = '12 dB (4)'

    if db_resp == 18:    
        db_resp_str = '18 dB (8)'

    if db_resp == 24:    
        db_resp_str = '24 dB (16)'

    if db_resp == 30:    
        db_resp_str = '30 dB (32)'

    if db_resp == 36:    
        db_resp_str = '36 dB (64)'

    if db_resp == 6:    
        db_resp_str = '6 dB (2)'


    # We'll first create all the various objects. These strongly follow the
    # hierarchy of StationXML files.
    inv = Inventory(
        # We'll add networks later.
        networks=[],
        # The source should be the id whoever create the file.
        source="ObsPy-Tutorial")

    net_tmp = Network(
        # This is the network code according to the SEED standard.
        #code="AA",
        code=net,
        # A list of stations. We'll add one later.
        stations=[],
        description="SmartSolo test",
        # Start-and end dates are optional.
        start_date=UTCDateTime(2019, 1, 1))

    sta_tmp = Station(
        # This is the station code according to the SEED standard.
        #code="9001",
        code=sta,
        latitude=st_lat,
        longitude=st_lon,
        elevation=st_ele,
        creation_date=UTCDateTime(2019, 1, 1),
        site=Site(name="Inside the Byerly Vault "))


    cha_tmp = Channel(
        # This is the channel code according to the SEED standard.
        #code="DPZ",
        code=com,
        # This is the location code according to the SEED standard.
        location_code=loc,
        # Note that these coordinates can differ from the station coordinates.
        latitude=st_lat,
        longitude=st_lon,
        elevation=st_ele,
        depth=st_dep,
        azimuth=azimuth,
        dip=dip,
        sample_rate=sample_rate)
        
    nrl = NRL()
    response = nrl.get_response( 
        sensor_keys=['DTCC (manuafacturers of SmartSolo)', 'DT-SOLO', '5 Hz', 'Rc=1850, Rs=430000'],
        datalogger_keys=['DTCC (manufacturers of SmartSolo', 'SmartSolo IGU-16HR3C', db_resp_str, str(sample_rate), 'Minimum Phase', 'Off'])
        
    
    cha_tmp.response = response
    sta_tmp.channels.append(cha_tmp)
    net_tmp.stations.append(sta_tmp)

    inv.networks.append(net_tmp)

    return inv



# %%
plotOPT = 1

# %%
#data_dir = "/work/dc6/ftp/outgoing/taira/WQC/lvnodal/llnlset_1st"
data_dir = "/work/dc6/ftp/outgoing/taira/WQC/lvnodal/ucbset_1st"

# BDA01	x	20251208_s2	x	1185 Apache St, Livermore, CA	37.690987	-121.784102	453012161	2025-10-25T10:55	0	453016777	2025-12-06T10:11																	
#453012161.0016.2025.11.09.00.00.00.000.E.miniseed
ms_fi = "453012161*.2025.11.09.00.00.00.000.*.miniseed"

st = read(data_dir+"/"+ms_fi)

print(st)

# %%
st.merge()

# %%
if plotOPT:
    _plot = st.plot()

# %%
net = st[0].stats.network 
loc = st[0].stats.location 
com = st[0].stats.channel 
sta = st[0].stats.station 

print("# net = ", net)
print("# loc = ", loc)
print("# com = ", com)
print("# sta = ", sta)


# %%
# LVNodal
sample_rate = 500 # 250sps
db_resp = 18 # 18 db


# %%
inv = Inventory()
#for tr in st_tmp:
for tr in st:

    net, sta, loc, com = tr.id.split('.') 
    inv += make_resp(sta, net, com, loc, sample_rate=sample_rate, db_resp=db_resp )
    #tr.write(tr.id+".sac", format="SAC")


# %%
print(inv)
#inv.write("MB_Nodal_inv.xml", format="STATIONXML")
inv.write("LVNodal_inv.xml", format="STATIONXML")

# %%
plt.rcParams['figure.figsize'] = 15,10

if plotOPT:
    #_plot = inv.plot(0.001, output="VEL")  
    #_plot = inv.plot(output="VEL")          
    #_plot = inv.plot_response(min_freq=0.001)
    _plot = inv.plot_response(min_freq=0.01)


# %%
# psd estimate
min_db = -200 # min. db 
max_db = -50 # max db
ddb = 1 # 1 db increment

#ppsd_length = 4*60*60 # 4 hour unit is seconds
ppsd_length = 2*60*60 # 2 hour unit is seconds


# %%
st[0].stats

# %%

#freq_min = 1/4000 
#freq_max = 200

freq_min = 1/1000 
freq_max = 500 

period_min = 1/freq_max
period_max = 1/freq_min

cmap = viridis_white_r



# %%
plt.rcParams['figure.figsize'] = 12, 10


for tr in st:
    #ppsd = PPSD(stats=st[0].stats, metadata=inv, db_bins=(min_db, max_db, ddb), ppsd_length=ppsd_length)
    #ppsd.add(st[0])
    
    ppsd = PPSD(stats=tr.stats, metadata=inv, db_bins=(min_db, max_db, ddb), ppsd_length=ppsd_length)
    ppsd.add(tr)
    ppsd.plot(cmap=cmap, period_lim=(period_min, period_max), show_mean = True, xaxis_frequency=False)