import h5py 
import matplotlib.pyplot as plt 
import datetime 
import numpy as np 
from obspy.geodetics import locations2degrees, degrees2kilometers 
def cal_km(loc1, loc2):
    return degrees2kilometers(locations2degrees(loc1[1], loc1[0], loc2[1], loc2[0]))

file_name = "/home/lilu/work/beijingju/h5data/publish/credit-x1.h5"
h5file = h5py.File(file_name, "r") 

for ekey in h5file:
    event = h5file[ekey]
    print(f"event:{ekey}")
    for key in event.attrs:
        if "strs" in key:continue 
        print(f"|-{key}:{event.attrs[key]}")
    for skey in event:
        station = event[skey]
        print(f"|-station:{skey}")
        for key2 in station:
            wave = station[key2]
            if "strs" in key:continue 
            print(f"  |-Data:{key2}:{wave[:].shape}")
            for key3 in wave.attrs:
                print(f"    |-{key3}:{wave.attrs[key3]}")
        for key in station.attrs:
            if "strs" in key:continue 
            print(f"  |-{key}:{station.attrs[key]}")
        dist = cal_km(
            [event.attrs["event_longitude"], event.attrs["event_latitude"]], 
            [station.attrs["station_longitude"], station.attrs["station_latitude"]]
        )
        print(dist)
        break 
    break 
#绘图
for ekey in h5file:
    event = h5file[ekey]
    count = 0 
    etime = datetime.datetime.strptime(event.attrs["event_origin_time"], "%Y/%m/%d %H:%M:%S.%f")
    for skey in event:
        station = event[skey]
        for key2 in station:
            wave = station[key2]
            for key3 in wave.attrs:
                if "start_time" in key3:
                    btime = datetime.datetime.strptime(wave.attrs[key3], "%Y/%m/%d %H:%M:%S.%f")
        w = wave[:] 
        w = w.astype(np.float32)
        w -= np.mean(w) 
        w /= np.max(np.abs(w)) 
        dist = cal_km(
            [event.attrs["event_longitude"], event.attrs["event_latitude"]], 
            [station.attrs["station_longitude"], station.attrs["station_latitude"]]
            )
        t = np.arange(len(w)) * 0.01 + (btime-etime).total_seconds()
        plt.plot(t, w*5+dist, c="k", alpha=0.5, lw=1)
        for key in station.attrs:
            if "strs" in key:continue 
            if key.endswith("Pg") or key.endswith("Sg"):
                #if "/" not in wave.attrs[key3]:continue 
                ptime = datetime.datetime.strptime(station.attrs[key], "%Y/%m/%d %H:%M:%S.%f")
                delta = (ptime-etime).total_seconds()
                if "Pg" in key:
                    c = "r"
                else:
                    c = "b"
                #print(delta)
                plt.plot([delta, delta], [count-5 + dist, count+5+dist], c=c)
        count += 1
    #plt.xlim((0, 60))
    plt.savefig(f"odata/figs/{ekey}.m65.png")
    plt.cla()
    plt.clf()