import os
import obspy 
import pickle 
import datetime 
import h5py 
from obspy.clients.fdsn.header import FDSNNoDataException
from obspy import UTCDateTime
import time 
import pandas as pd
import multiprocessing 
import numpy as np 
import matplotlib.pyplot as plt
from obspy.signal.filter import bandpass 
from scipy import signal
from obspy.geodetics.base import gps2dist_azimuth 

class DataForSeiPnTest():
    def __init__(self, file_name="models/h5test/all-gzip4.h5", maxdist = 2000, slatrang = [-90, 90], slonrang= [-180, 180]):
        self.file_name = file_name
        self.n_thread = 2
        self.maxdist = maxdist
        self.slatrang = slatrang
        self.slonrang = slonrang
        self.dqueue = multiprocessing.Queue(100)
        for h5index in range(3):
            multiprocessing.Process(target=self.feed_data, args=(self.dqueue, h5index)).start() 
        #multiprocessing.Process(target=self.batch_data, args=(dqueue, )).start()
    def feed_data(self, dqueue, h5index):
        file_name = f"SeisPnSn/Pn/SeisPnSn{h5index+1}.hdf5"
        
        csvfile = pd.read_csv(f"seipnsn/pn_part_metadata/pn{h5index+1}.csv")
        h5keys = csvfile["key"]
        h5file = h5py.File(file_name, "r")
        h5fileyears = h5file["pn"]
        ##pn/2000/2000.10.30.22.40.20.KN.TKM2
        #print("Actual key types:", [type(k) for k in h5file["pn"].keys()])
        print(f"{file_name}数据加载完成")
        years = list(h5fileyears.keys()) 
        #print(years)
        for i in  range(5):
            year = years[i]
            h5fileyear = h5fileyears[year]
            print(f"当前数据年份为{year}")
            for ekey in h5keys:
                #print(ekey)
                if ekey.startswith(year):
                    #print(ekey)
                    event = h5fileyear[ekey]
                    elat =  event.attrs["event_la"]
                    elon =  event.attrs["event_lo"]
                    slat =  event.attrs["station_la"]
                    slon =  event.attrs["station_lo"]
                    if slat < self.slatrang[0] or slat > self.slatrang[1]:continue
                    if slon < self.slonrang[0] or slon > self.slonrang[1]:continue
                    dist, abz, baz = gps2dist_azimuth(elat, elon, slat, slon)
                    #print(elat, elon, slat, slon, dist)
                    dist = dist/1000 
                    if dist > self.maxdist:continue
                    #print(dist)
                    rdata = [] 
                    data = [0, 0, 0]
                    snr = event.attrs["snr_z_dB"]
                   
                    skey = event.attrs["station"]
                    data[0] = event[:, 0]
                    data[1] = event[:, 1]
                    data[2] = event[:, 2]
                    data[0] =  signal.resample(data[0], 18000)
                    data[1] =  signal.resample(data[1], 18000)
                    data[2] =  signal.resample(data[2], 18000)
                    if len(data)!=3:continue  #data[HE数据，HN数据，HZ数据]
                    for d in data:
                        w = d[0:10240]
                        w = w - np.mean(w)
                        w = w / (np.max(np.abs(w))+1e-6)
                        rdata.append(w[np.newaxis, :, np.newaxis])
                    rdata = np.concatenate(rdata, axis=2) 
                    dqueue.put([rdata, [dist, ekey, skey, snr]])
    def batch_data(self, batch_size=50):
        x1, x2 = [], []
        for _ in range(batch_size):
            data, infos = self.dqueue.get() 
            x1.append(data) 
            x2.append(infos)
        x1 = np.concatenate(x1, axis=0) 
        return x1, x2