code/utils/dataAMAG.py

from cmath import polar
import os
from telnetlib import PRAGMA_HEARTBEAT 
import obspy 
import pickle 
import datetime 
import h5py 
import numpy as np 
from obspy.clients.fdsn.header import FDSNNoDataException
#from obspy.clients.fdsn import Client
from obspy import UTCDateTime
import time 
import multiprocessing 
from obspy.signal.filter import bandpass 



class DataForAMAGTest():
    def __init__(self, file_name="models/h5test/all-gzip4.h5", n_length=10240, stride=1, padlen=256, max_dist=2000):
        self.file_name = file_name
        self.length = n_length 
        self.stride = stride  
        self.max_dist = max_dist 
        self.padlen = padlen 
        self.n_thread = 2
        self.phase_dict = {
            "Pg":0, 
            "Sg":1,
            #"P":0, 
            #"S":1,
            "Pn":2, 
            "Sn":3, 
        } 
        self.ploar1 = {
            "C":0, 
            "U":0, 
            "R":1, 
            "D":1, 
        }
        self.ploar2 = {
            "I":0, 
            "M":1, 
            "E":2, 
        }
        self.etype_dict = {
            "eq":0, 
            "ve":1, 
            "ss":2, 
            "sp":3, 
            "ep":4, 
            "ot":5, 
            "se":6
        }
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        for year in [2009+i for i in range(10)]:
            multiprocessing.Process(target=self.feed_data, args=(fqueue, year)).start() 
        for _ in range(self.n_thread):
            multiprocessing.Process(target=self.process, args=(fqueue, self.dqueue)).start()
        #multiprocessing.Process(target=self.batch_data, args=(dqueue, )).start()
    def feed_data(self, fqueue, year):
        file_name = f"ayrdata/csndata/{year}.h5"
        h5keys = np.load(f"ayrdata/keys/{year}.npy")
        with open("large/distaz.pkl", "rb") as f:
            distaz = pickle.load(f) 
        while True:
            h5file = h5py.File(file_name, "r")
            print(f"{file_name}数据加载完成")
            for ekey in h5keys:
                #print(ekey)
                event = h5file[ekey]  #事件ID
                etype = event.attrs["type"] #事件类型
                mag = event.attrs["mag"]
                if etype in self.etype_dict:
                    typeid = self.etype_dict[etype] #是否为可识别的事件
                else:
                    typeid = 7  #不是可识别的6种，则ID设为7
                for skey in event:
                    station = event[skey] #获取事件的台站信息
                    data = [0, 0, 0] #初始化一个长度为3的列表，用于存储台站的数据
                    #print(skey)
                    
                    for dkey in station:
                        if "HZ" in dkey:
                            idx = 2 
                        elif "HE" in dkey:
                            idx = 0 
                        elif "HN" in dkey:
                            idx = 1 
                        else:
                            print(dkey, "Not exist")
                            continue 
                        data[idx] = station[dkey][:] 
                        #data.append(station[dkey][:]) 
                        btime = datetime.datetime.strptime(station[dkey].attrs['btime'], "%Y/%m/%d %H:%M:%S.%f")  #当前数据通道的属性中提取台站数据的起始时间
                    if len(data)!=3:continue  #data[HE数据，HN数据，HZ数据]
                    cnt = 0 
                    for d in data:
                        if type(d)==int:
                            cnt+= 1
                    if cnt!=0:continue 
                    phases = {}
                    phase_count = {"P":0, "S":0, "Pn":0, "Sn":0}
                    dist = -1
                    if "POLARITY.Pg.UPDOWN" in station.attrs and "POLARITY.Pg.CLARITY" in station.attrs:
                        ptype1 = station.attrs["POLARITY.Pg.UPDOWN"] #获取台站极性
                        ptype2 = station.attrs["POLARITY.Pg.CLARITY" ] #获取台站极性的清晰度
                        if ptype1 not in self.ploar1 or ptype2 not in self.ploar2:
                            polars = []
                        else:
                            polars = [self.ploar1[ptype1], self.ploar2[ptype2]]#获取极性信息和清晰度对应的数值
                    else:
                        polars = []
                    ptypes = [i.split(".")[-1] for i in station.attrs["types"].split(",")]
                    #print(ptypes)
                    if "Pg" not in ptypes :continue
                    #if "Sg" not in ptypes :continue
                    #if "Pn" not in ptypes :continue
                    #if "Sn" not in ptypes :continue

                    for akey in station.attrs:
                        pkey = akey.split(".")[-1]
                        if f"{pkey}.dist" in station.attrs:dist = float(station.attrs[f"{pkey}.dist"])
                        if pkey not in ptypes:continue 
                        pname = pkey.split("+i")[-1].split("-i")[-1].split("+")[-1].split("-")[-1].split("i")[-1].split("2")[0].split("*")[-1]
                        if "Pg" in pname:
                            pname = akey.split(".")[-1]
                            if pname in self.phase_dict:
                                phase_count["P"] += 1
                                #print(akey, ptypes)
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                        else:
                            if pname in self.phase_dict:
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                                if akey == "P" or akey == "Pg":
                                    phase_count["P"] += 1
                                if akey == "S" or akey == "Sg":
                                    phase_count["S"] += 1
                                if akey in ["Pn", "Sn"]:
                                    phase_count[akey] += 1
                    if dist > self.max_dist:continue 
                    #if phase_count["Pn"] !=0:
                    #    if phase_count["P"] ==0:
                    #        continue 
                    #if phase_count["Sn"] !=0:
                    #    if phase_count["S"] ==0:
                    #        continue 
                    #print(data)
                    #print(phases)
                    if len(phases)==0:continue 
                    fqueue.put([data, btime, phases, typeid, [dist, ekey, skey, mag]])

    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, btime, phases, etype, infos = fqueue.get()
            pidx = {}
            plist = []
            for pkey in phases:
                if pkey != "Pg": continue
                ptime = phases[pkey]
                delta = (ptime-btime).total_seconds()
                delta_idx = int(delta * 100)
                pidx[pkey] = delta_idx
                plist.append(delta_idx) 
            #print(plist)

            wz = data[-1]
            we = data[-2]
            wn = data[-3]
            wz_filter = bandpass(wz, 0.5, 20, 100)
            we_filter = bandpass(we, 0.5, 20, 100)
            wn_filter = bandpass(wn, 0.5, 20, 100)
            snrs = [-10000.0, -10000.0, -10000.0, -10000.0]
            for pkey, pdet in pidx.items():
                pi = self.phase_dict[pkey] 
                if pi == 0 or pi==2:# Pg or Pn
                    pre = wz_filter[pdet-50:pdet]
                    aft = wz_filter[pdet:pdet+50]
                    if len(pre)==0 or len(aft)==0:
                        continue 
                    snrs[pi] = 10 * np.log10((np.std(aft)+1e-6)/(np.std(pre)+1e-6))
                else:
                    pre1 = we_filter[pdet-150:pdet]
                    aft1 = we_filter[pdet:pdet+150]
                    pre2 = wn_filter[pdet-150:pdet]
                    aft2 = wn_filter[pdet:pdet+150]     
                    if len(pre1)==0 or len(pre2)==0 or len(aft1)==0 or len(aft2)==0:
                        continue 
                    snr1 = 10 * np.log10((np.std(aft1)+1e-6)/(np.std(pre1)+1e-6))
                    snr2 = 10 * np.log10((np.std(aft2)+1e-6)/(np.std(pre2)+1e-6))
                    snrs[pi] = snr1 * 0.5 + snr2 * 0.5 

            #fixplist = [plist[0]]
            #print(plist)
            cidx = np.random.choice(plist) - np.random.randint(self.padlen, self.length-self.padlen)
            rdata = [] 
            flen = False
            for d in data:
                w = d[cidx:cidx+self.length] 
                
                wlen = len(w) 
                if wlen!=self.length:
                    flen = True
                    break 
                #w = bandpass(w, 1, 20, 100)
                #w = w - np.mean(w)
                #w = w / (np.std(w)+1e-6)
                #w = w / (np.max(np.abs(w))+1e-6)
                rdata.append(w[np.newaxis, :, np.newaxis]) 
            if flen:
                continue
            


            rdata = np.concatenate(rdata, axis=2) 
            label1 = np.zeros([1, llen, 2])
            label2 = np.zeros([1, self.length, 5]) 
            label_polar = np.zeros([1, self.length]) 
            label_quali = np.zeros([1, self.length]) 
            label_weigh = np.zeros([1, self.length])
            for pkey in pidx:
                pid = self.phase_dict[pkey] 
                idx = (pidx[pkey] - cidx)//self.stride 
                
                if idx-1>0:
                    label1[0, idx-1:idx+2] = -1
                if idx > 0 and idx < llen:
                    label1[0, idx, 0] = pid + 1
                    label1[0, idx, 1] = (pidx[pkey] - cidx)%self.stride
            phase_intv = {"P":0, "S":0} 
            def norm(t, mu, std=0.1):
                p = np.exp(-(t-mu)**2/std**2/2)
                p /= (np.max(p)+1e-6) 
                return p 
            t = np.arange(self.length) * 0.01 
            phase_time = {0:-1, 1:-1, 2:-1, 3:-1}
            for pkey in pidx:
                pid = self.phase_dict[pkey] 
                idx = (pidx[pkey] - cidx)
                if idx > 0 and idx < self.length:
                    phase_time[pid] = idx    
            dqueue.put([rdata, [phase_time[0], phase_time[1], phase_time[2], phase_time[3]], snrs, infos])

            count += 1

    def batch_data(self, batch_size=50):
        x1, x2, x3, x4 = [], [], [], []
        for _ in range(batch_size):
            data, label1, snrs, infos = self.dqueue.get() 
            x1.append(data) 
            x2.append(label1) 
            x3.append(snrs)
            x4.append(infos)
            
        x1 = np.concatenate(x1, axis=0) 
        return x1, x2, x3, x4 




class DataTestLargeForIPolar():
    def __init__(self, file_name="models/h5test/all-gzip4.h5", n_length=10240, stride=1, padlen=256):
        self.file_name = file_name
        self.length = n_length 
        self.stride = stride  
        self.padlen = padlen 
        self.n_thread = 2
        self.phase_dict = {
            "Pg":0, 
            "Sg":1,
            "P":0, 
            "S":1,
            "Pn":2, 
            "Sn":3, 
        } 
        self.ploar1 = {
            "C":0, 
            "U":0, 
            "R":1, 
            "D":1, 
        }
        self.ploar2 = {
            #I平缓，M锐利， E既不平缓也不锐利，无标签
            "I":0, 
            "M":1, 
            "E":2, 
        }
        self.etype_dict = {
            "eq":0, 
            "ve":1, 
            "ss":2, 
            "sp":3, 
            "ep":4, 
            "ot":5, 
            "se":6
        }
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        for year in [2020+i for i in range(3)]:
            multiprocessing.Process(target=self.feed_data, args=(fqueue, year)).start() 
        for _ in range(self.n_thread):
            multiprocessing.Process(target=self.process, args=(fqueue, self.dqueue)).start()
        #multiprocessing.Process(target=self.batch_data, args=(dqueue, )).start()
    def feed_data(self, fqueue, year):
        file_name = f"ayrdata/csndata/{year}.h5"
        h5keys = np.load(f"ayrdata/keys/{year}.npy")
        with open("large/distaz.pkl", "rb") as f:
            distaz = pickle.load(f) 
        while True:
            h5file = h5py.File(file_name, "r")
            print(f"{file_name}数据加载完成")
            for ekey in h5keys:
                #print(ekey)
                event = h5file[ekey] 
                etype = event.attrs["type"]
                if etype in self.etype_dict:
                    typeid = self.etype_dict[etype] 
                else:
                    typeid = 7 
                for skey in event:
                    station = event[skey] 
                    #print(station)
                    data = [0, 0, 0]
                    #print(skey)
                    
                    for dkey in station:
                        if "HZ" in dkey:
                            idx = 2 
                        elif "HE" in dkey:
                            idx = 0 
                        elif "HN" in dkey:
                            idx = 1 
                        else:
                            print(dkey, "Not exist")
                            continue 
                        data[idx] = station[dkey][:] 
                        #data.append(station[dkey][:]) 
                        btime = datetime.datetime.strptime(station[dkey].attrs['btime'], "%Y/%m/%d %H:%M:%S.%f") 
                    if len(data)!=3:continue 
                    cnt = 0 
                    for d in data:
                        if type(d)==int:
                            cnt+= 1
                    if cnt!=0:continue 
                    phases = {}
                    phase_count = {"P":0, "S":0, "Pn":0, "Sn":0}
                    dist = -1
                    if "POLARITY.Pg.UPDOWN" in station.attrs and "POLARITY.Pg.CLARITY" not in station.attrs:print("有初动方向，无清晰度标识")
                    if "POLARITY.Pg.UPDOWN" not in station.attrs and "POLARITY.Pg.CLARITY" in station.attrs:print("无初动方向，有清晰度标识")
                    
                    if "POLARITY.Pg.UPDOWN" in station.attrs and "POLARITY.Pg.CLARITY" in station.attrs:
                        ptype1 = station.attrs["POLARITY.Pg.UPDOWN"]
                        ptype2 = station.attrs["POLARITY.Pg.CLARITY"]
                        if ptype2 != "I": continue
                        #print(ptype1,ptype2)
                        if ptype1 not in self.ploar1 or ptype2 not in self.ploar2:
                            polars = []
                        else:
                            polars = [self.ploar1[ptype1], self.ploar2[ptype2]]
                    else:
                        polars = []
                    ptypes = [i.split(".")[-1] for i in station.attrs["types"].split(",")]
                    

                    for akey in station.attrs:
                        pkey = akey.split(".")[-1]
                        if pkey not in ptypes:continue 
                        pname = pkey.split("+i")[-1].split("-i")[-1].split("+")[-1].split("-")[-1].split("i")[-1].split("2")[0].split("*")[-1]
                        if "Pg" in pname:
                            pname = akey.split(".")[-1]
                            if pname in self.phase_dict:
                                phase_count["P"] += 1
                                #print(akey, ptypes)
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                        else:
                            if pname in self.phase_dict:
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                                if akey == "P" or akey == "Pg":
                                    phase_count["P"] += 1
                                if akey == "S" or akey == "Sg":
                                    phase_count["S"] += 1
                                if akey in ["Pn", "Sn"]:
                                    phase_count[akey] += 1
                    #if dist > 2000:continue 
                    #if phase_count["Pn"] !=0:
                    #    if phase_count["P"] ==0:
                    #        continue 
                    #if phase_count["Sn"] !=0:
                    #    if phase_count["S"] ==0:
                    #        continue 
                    #print(data)
                    if len(phases)==0:continue 
                    if len(polars)==0:continue
                    #if polars[1] == 2:
                        #print(polars)
                        #print("无标签")
                        #continue
                    #print(polars)
                    #print(ptype1,ptype2)
                    fqueue.put([data, btime, phases, polars, typeid])

    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, btime, phases, polars, etype = fqueue.get()
            #print(polars)
            pidx = {}
            plist = []
            for pkey in phases:
                ptime = phases[pkey]
                delta = (ptime-btime).total_seconds()
                delta_idx = int(delta * 100)
                pidx[pkey] = delta_idx
                plist.append(delta_idx) 
            #print(pidx)
            #cidx设置数据的随机起始点    
            cidx = np.random.choice(plist) - np.random.randint(self.padlen, self.length-self.padlen)
            rdata = [] 
            flen = False
            for d in data:
                w = d[cidx:cidx+self.length] 
                wlen = len(w) 
                if wlen!=self.length:
                    flen = True
                    break 
                #mean去均值
                #std归一化
                w = w - np.mean(w)
                w = w / (np.std(w)+1e-6)
                rdata.append(w[np.newaxis, :, np.newaxis]) 
            if flen:
                continue
            rdata = np.concatenate(rdata, axis=2) 
            label1 = np.zeros([1, llen, 2])
            label2 = np.zeros([1, self.length, 5]) 
            label_polar = np.zeros([1, self.length]) 
            label_quali = np.zeros([1, self.length]) 
            label_weigh = np.zeros([1, self.length])
            #pidx 各个震相在rdata里面的时间点
            for pkey in pidx:
                pid = self.phase_dict[pkey] 
                idx = (pidx[pkey] - cidx)//self.stride 
                
                if idx-1>0:
                    label1[0, idx-1:idx+2] = -1
                if idx > 0 and idx < llen:
                    label1[0, idx, 0] = pid + 1
                    label1[0, idx, 1] = (pidx[pkey] - cidx)%self.stride
            phase_intv = {"P":0, "S":0} 
            def norm(t, mu, std=0.1):
                p = np.exp(-(t-mu)**2/std**2/2)
                p /= (np.max(p)+1e-6) 
                return p 
            t = np.arange(self.length) * 0.01 
            #print(pidx)  {'Pg': 50000, 'Sg': 51760, 'Sn': 51557}
            
            #print(polars)
            for pkey in pidx:
                
                pid = self.phase_dict[pkey] 
                # idx 为rdata的起始事件到该震相的时间长度
                idx = (pidx[pkey] - cidx)
                if idx > 0 and idx < self.length:
                    label2[0, :, pid+1] = norm(t, idx*0.01, 0.1)
                    if pid == 0 and len(polars)>0:  #只取pg波的时间
                        begin = np.clip(idx-50, 0, self.length-60)
                        #print(begin)
                        label_polar[0, begin:begin+100] = polars[0] 
                        label_quali[0, begin:begin+100] = polars[1]
                        label_weigh[0, begin:begin+100] = 1
            label2[0, :, 0] = np.clip(1-label2[0, :, 1]-label2[0, :, 2]-label2[0, :, 3]-label2[0, :, 4], 0, 1)   
            dqueue.put([rdata, [label_polar, label_quali, label_weigh]])

            count += 1

    def batch_data(self, batch_size=32):
        x1, x2, x3 = [], [], []
        p1, p2, p3 = [], [], [] 
        x4 = []
        for _ in range(batch_size):
            data,(label_polar, label_quali, label_weigh) = self.dqueue.get() 
            x1.append(data) 
            #x2.append(label1) 
            #x3.append(label2)
            #x4.append(etype) 
            p1.append(label_polar) 
            p2.append(label_quali) 
            p3.append(label_weigh)
        x1 = np.concatenate(x1, axis=0) 
        #x2 = np.concatenate(x2, axis=0) 
        #x3 = np.concatenate(x3, axis=0) 
        p1 = np.concatenate(p1, axis=0) 
        p2 = np.concatenate(p2, axis=0) 
        p3 = np.concatenate(p3, axis=0) 
        #x4 = np.array(x4)
        #print(p1)
        return  x1, p1, p2, p3
    


class DataTestLargeForAllPolar():
    def __init__(self, file_name="models/h5test/all-gzip4.h5", n_length=10240, stride=1, padlen=256):
        self.file_name = file_name
        self.length = n_length 
        self.stride = stride  
        self.padlen = padlen 
        self.n_thread = 2
        self.phase_dict = {
            "Pg":0, 
            "Sg":1,
            #"P":0, 
            #"S":1,
            "Pn":2, 
            "Sn":3, 
        } 
        self.ploar1 = {
            "C":0, 
            "U":0, 
            "R":1, 
            "D":1, 
        }
        self.ploar2 = {
            #I平缓，M锐利， E既不平缓也不锐利，无标签
            "I":0, 
            "M":1, 
            "E":2, 
        }
        self.etype_dict = {
            "eq":0, 
            "ve":1, 
            "ss":2, 
            "sp":3, 
            "ep":4, 
            "ot":5, 
            "se":6
        }
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        for year in [2020+i for i in range(3)]:
            multiprocessing.Process(target=self.feed_data, args=(fqueue, year)).start() 
        for _ in range(self.n_thread):
            multiprocessing.Process(target=self.process, args=(fqueue, self.dqueue)).start()
        #multiprocessing.Process(target=self.batch_data, args=(dqueue, )).start()
    def feed_data(self, fqueue, year):
        file_name = f"ayrdata/csndata/{year}.h5"
        h5keys = np.load(f"ayrdata/keys/{year}.npy")
        with open("large/distaz.pkl", "rb") as f:
            distaz = pickle.load(f) 
        while True:
            h5file = h5py.File(file_name, "r")
            print(f"{file_name}数据加载完成")
            for ekey in h5keys:
                #print(ekey)
                event = h5file[ekey] 
                etype = event.attrs["type"]
                if etype in self.etype_dict:
                    typeid = self.etype_dict[etype] 
                else:
                    typeid = 7 
                for skey in event:
                    station = event[skey] 
                    #print(station)
                    data = [0, 0, 0]
                    #print(skey)
                    
                    for dkey in station:
                        if "HZ" in dkey:
                            idx = 2 
                        elif "HE" in dkey:
                            idx = 0 
                        elif "HN" in dkey:
                            idx = 1 
                        else:
                            print(dkey, "Not exist")
                            continue 
                        data[idx] = station[dkey][:] 
                        #data.append(station[dkey][:]) 
                        btime = datetime.datetime.strptime(station[dkey].attrs['btime'], "%Y/%m/%d %H:%M:%S.%f") 
                    if len(data)!=3:continue 
                    cnt = 0 
                    for d in data:
                        if type(d)==int:
                            cnt+= 1
                    if cnt!=0:continue 
                    phases = {}
                    phase_count = {"P":0, "S":0, "Pn":0, "Sn":0}
                    dist = -1
                    if "POLARITY.Pg.UPDOWN" in station.attrs and "POLARITY.Pg.CLARITY" in station.attrs:
                        ptype1 = station.attrs["POLARITY.Pg.UPDOWN"]
                        ptype2 = station.attrs["POLARITY.Pg.CLARITY" ]
                        #if ptype2 != "M": continue
                        if ptype1 not in self.ploar1 or ptype2 not in self.ploar2:
                            polars = []
                        else:
                            polars = [self.ploar1[ptype1], self.ploar2[ptype2]]
                    else:
                        polars = []
                    
                    ptypes = [i.split(".")[-1] for i in station.attrs["types"].split(",")]


                    for akey in station.attrs:
                        pkey = akey.split(".")[-1]
                        if pkey not in ptypes:continue 
                        pname = pkey.split("+i")[-1].split("-i")[-1].split("+")[-1].split("-")[-1].split("i")[-1].split("2")[0].split("*")[-1]
                        if "Pg" in pname:
                            pname = akey.split(".")[-1]
                            if pname in self.phase_dict:
                                phase_count["P"] += 1
                                #print(akey, ptypes)
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                        else:
                            if pname in self.phase_dict:
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                                if akey == "P" or akey == "Pg":
                                    phase_count["P"] += 1
                                if akey == "S" or akey == "Sg":
                                    phase_count["S"] += 1
                                if akey in ["Pn", "Sn"]:
                                    phase_count[akey] += 1
                    #if dist > 2000:continue 
                    #if phase_count["Pn"] !=0:
                    #    if phase_count["P"] ==0:
                    #        continue 
                    #if phase_count["Sn"] !=0:
                    #    if phase_count["S"] ==0:
                    #        continue 
                    #print(data)
                    if len(phases)==0:continue 
                    if len(polars)==0:continue

                    #if polars[1] == 2:
                    #    print("无标签")
                    #    continue
                    #print(polars)
                    fqueue.put([data, btime, phases, polars, typeid])

    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, btime, phases, polars, etype = fqueue.get()
            #print(polars)
            pidx = {}
            plist = []
            for pkey in phases:
                ptime = phases[pkey]
                delta = (ptime-btime).total_seconds()
                delta_idx = int(delta * 100)
                pidx[pkey] = delta_idx
                plist.append(delta_idx) 
            #print(pidx)
            #cidx设置数据的随机起始点    
            cidx = np.random.choice(plist) - np.random.randint(self.padlen, self.length-self.padlen)
            rdata = [] 
            flen = False
            for d in data:
                w = d[cidx:cidx+self.length] 
                wlen = len(w) 
                if wlen!=self.length:
                    flen = True
                    break 
                #mean去均值
                #std归一化
                w = w - np.mean(w)
                w = w / (np.std(w)+1e-6)
                rdata.append(w[np.newaxis, :, np.newaxis]) 
            if flen:
                continue
            rdata = np.concatenate(rdata, axis=2) 
            label1 = np.zeros([1, llen, 2])
            label2 = np.zeros([1, self.length, 5]) 
            label_polar = np.zeros([1, self.length]) 
            label_quali = np.zeros([1, self.length]) 
            label_weigh = np.zeros([1, self.length])
            #pidx 各个震相在rdata里面的时间点
            for pkey in pidx:
                pid = self.phase_dict[pkey] 
                idx = (pidx[pkey] - cidx)//self.stride 
                
                if idx-1>0:
                    label1[0, idx-1:idx+2] = -1
                if idx > 0 and idx < llen:
                    label1[0, idx, 0] = pid + 1
                    label1[0, idx, 1] = (pidx[pkey] - cidx)%self.stride
            phase_intv = {"P":0, "S":0} 
            def norm(t, mu, std=0.1):
                p = np.exp(-(t-mu)**2/std**2/2)
                p /= (np.max(p)+1e-6) 
                return p 
            t = np.arange(self.length) * 0.01 
            #print(pidx)  {'Pg': 50000, 'Sg': 51760, 'Sn': 51557}
            
            #print(polars)
            for pkey in pidx:
                
                pid = self.phase_dict[pkey] 
                # idx 为rdata的起始事件到该震相的时间长度
                idx = (pidx[pkey] - cidx)
                if idx > 0 and idx < self.length:
                    label2[0, :, pid+1] = norm(t, idx*0.01, 0.1)
                    if pid == 0 and len(polars)>0:  #只取pg波的时间
                        begin = np.clip(idx-50, 0, self.length-60)
                        #print(begin)
                        label_polar[0, begin:begin+100] = polars[0] 
                        label_quali[0, begin:begin+100] = polars[1]
                        label_weigh[0, begin:begin+100] = 1
            label2[0, :, 0] = np.clip(1-label2[0, :, 1]-label2[0, :, 2]-label2[0, :, 3]-label2[0, :, 4], 0, 1)   
            dqueue.put([rdata, [label_polar, label_quali, label_weigh]])

            count += 1

    def batch_data(self, batch_size=32):
        x1, x2, x3 = [], [], []
        p1, p2, p3 = [], [], [] 
        x4 = []
        for _ in range(batch_size):
            data,(label_polar, label_quali, label_weigh) = self.dqueue.get() 
            x1.append(data) 
            #x2.append(label1) 
            #x3.append(label2)
            #x4.append(etype) 
            p1.append(label_polar) 
            p2.append(label_quali) 
            p3.append(label_weigh)
        x1 = np.concatenate(x1, axis=0) 
        #x2 = np.concatenate(x2, axis=0) 
        #x3 = np.concatenate(x3, axis=0) 
        p1 = np.concatenate(p1, axis=0) 
        p2 = np.concatenate(p2, axis=0) 
        p3 = np.concatenate(p3, axis=0) 
        #x4 = np.array(x4)
        #print(p1)
        return  x1, p1, p2, p3










class DataTestLargeForEPolar():
    def __init__(self, file_name="models/h5test/all-gzip4.h5", n_length=10240, stride=1, padlen=256):
        self.file_name = file_name
        self.length = n_length 
        self.stride = stride  
        self.padlen = padlen 
        self.n_thread = 2
        self.phase_dict = {
            "Pg":0, 
            "Sg":1,
            "P":0, 
            "S":1,
            "Pn":2, 
            "Sn":3, 
        } 
        self.ploar1 = {
            "C":0, 
            "U":0, 
            "R":1, 
            "D":1, 
        }
        self.ploar2 = {
            #I平缓，M锐利， E既不平缓也不锐利，无标签
            "I":0, 
            "M":1, 
            "E":2, 
        }
        self.etype_dict = {
            "eq":0, 
            "ve":1, 
            "ss":2, 
            "sp":3, 
            "ep":4, 
            "ot":5, 
            "se":6
        }
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        for year in [2020]:
            multiprocessing.Process(target=self.feed_data, args=(fqueue, year)).start() 
        for _ in range(self.n_thread):
            multiprocessing.Process(target=self.process, args=(fqueue, self.dqueue)).start()
        #multiprocessing.Process(target=self.batch_data, args=(dqueue, )).start()
    def feed_data(self, fqueue, year):
        file_name = f"ayrdata/csndata/{year}.h5"
        h5keys = np.load(f"ayrdata/keys/{year}.npy")
        with open("large/distaz.pkl", "rb") as f:
            distaz = pickle.load(f) 
        while True:
            h5file = h5py.File(file_name, "r")
            print(f"{file_name}数据加载完成")
            for ekey in h5keys:
                #print(ekey)
                event = h5file[ekey] 
                etype = event.attrs["type"]
                if etype in self.etype_dict:
                    typeid = self.etype_dict[etype] 
                else:
                    typeid = 7 
                for skey in event:
                    station = event[skey] 
                    #print(station)
                    data = [0, 0, 0]
                    #print(skey)
                    
                    for dkey in station:
                        if "HZ" in dkey:
                            idx = 2 
                        elif "HE" in dkey:
                            idx = 0 
                        elif "HN" in dkey:
                            idx = 1 
                        else:
                            print(dkey, "Not exist")
                            continue 
                        data[idx] = station[dkey][:] 
                        #data.append(station[dkey][:]) 
                        btime = datetime.datetime.strptime(station[dkey].attrs['btime'], "%Y/%m/%d %H:%M:%S.%f") 
                    if len(data)!=3:continue 
                    cnt = 0 
                    for d in data:
                        if type(d)==int:
                            cnt+= 1
                    if cnt!=0:continue 
                    phases = {}
                    phase_count = {"P":0, "S":0, "Pn":0, "Sn":0}
                    dist = -1
                    if "POLARITY.Pg.UPDOWN" in station.attrs and "POLARITY.Pg.CLARITY" not in station.attrs:print("有初动方向，无清晰度标识")
                    if "POLARITY.Pg.UPDOWN" not in station.attrs and "POLARITY.Pg.CLARITY" in station.attrs:print("无初动方向，有清晰度标识")
                    if "POLARITY.Pg.UPDOWN" in station.attrs and "POLARITY.Pg.CLARITY" in station.attrs:
                        ptype1 = station.attrs["POLARITY.Pg.UPDOWN"]
                        ptype2 = station.attrs["POLARITY.Pg.CLARITY"]
                        if ptype2 != "E": continue
                        #if ptype2 == "E": print(ptype1)
                        if ptype1 not in self.ploar1 or ptype2 not in self.ploar2:
                            polars = []
                        else:
                            polars = [self.ploar1[ptype1], self.ploar2[ptype2]]
                    else:
                        polars = []
                    ptypes = [i.split(".")[-1] for i in station.attrs["types"].split(",")]
                    #if station.attrs["POLARITY.Pg.CLARITY" ] != "E":continue

                    for akey in station.attrs:
                        pkey = akey.split(".")[-1]
                        if pkey not in ptypes:continue 
                        pname = pkey.split("+i")[-1].split("-i")[-1].split("+")[-1].split("-")[-1].split("i")[-1].split("2")[0].split("*")[-1]
                        if "Pg" in pname:
                            pname = akey.split(".")[-1]
                            if pname in self.phase_dict:
                                phase_count["P"] += 1
                                #print(akey, ptypes)
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                        else:
                            if pname in self.phase_dict:
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                                if akey == "P" or akey == "Pg":
                                    phase_count["P"] += 1
                                if akey == "S" or akey == "Sg":
                                    phase_count["S"] += 1
                                if akey in ["Pn", "Sn"]:
                                    phase_count[akey] += 1
                    #if dist > 2000:continue 
                    #if phase_count["Pn"] !=0:
                    #    if phase_count["P"] ==0:
                    #        continue 
                    #if phase_count["Sn"] !=0:
                    #    if phase_count["S"] ==0:
                    #        continue 
                    #print(data)
                    if len(phases)==0:continue 
                    if len(polars)==0:continue
                    #if polars[1] == 2:
                    #    print(polars)
                    #    #print("无标签")
                    #    continue
                    #print(polars)
                    fqueue.put([data, btime, phases, polars, typeid])

    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, btime, phases, polars, etype = fqueue.get()
            print(polars)
            pidx = {}
            plist = []
            for pkey in phases:
                ptime = phases[pkey]
                delta = (ptime-btime).total_seconds()
                delta_idx = int(delta * 100)
                pidx[pkey] = delta_idx
                plist.append(delta_idx) 
            #print(pidx)
            #cidx设置数据的随机起始点    
            cidx = np.random.choice(plist) - np.random.randint(self.padlen, self.length-self.padlen)
            rdata = [] 
            flen = False
            for d in data:
                w = d[cidx:cidx+self.length] 
                wlen = len(w) 
                if wlen!=self.length:
                    flen = True
                    break 
                #mean去均值
                #std归一化
                w = w - np.mean(w)
                w = w / (np.std(w)+1e-6)
                rdata.append(w[np.newaxis, :, np.newaxis]) 
            if flen:
                continue
            rdata = np.concatenate(rdata, axis=2) 
            label1 = np.zeros([1, llen, 2])
            label2 = np.zeros([1, self.length, 5]) 
            label_polar = np.zeros([1, self.length]) 
            label_quali = np.zeros([1, self.length]) 
            label_weigh = np.zeros([1, self.length])
            #pidx 各个震相在rdata里面的时间点
            for pkey in pidx:
                pid = self.phase_dict[pkey] 
                idx = (pidx[pkey] - cidx)//self.stride 
                
                if idx-1>0:
                    label1[0, idx-1:idx+2] = -1
                if idx > 0 and idx < llen:
                    label1[0, idx, 0] = pid + 1
                    label1[0, idx, 1] = (pidx[pkey] - cidx)%self.stride
            phase_intv = {"P":0, "S":0} 
            def norm(t, mu, std=0.1):
                p = np.exp(-(t-mu)**2/std**2/2)
                p /= (np.max(p)+1e-6) 
                return p 
            t = np.arange(self.length) * 0.01 
            #print(pidx)  {'Pg': 50000, 'Sg': 51760, 'Sn': 51557}
            
            #print(polars)
            for pkey in pidx:
                
                pid = self.phase_dict[pkey] 
                # idx 为rdata的起始事件到该震相的时间长度
                idx = (pidx[pkey] - cidx)
                if idx > 0 and idx < self.length:
                    label2[0, :, pid+1] = norm(t, idx*0.01, 0.1)
                    if pid == 0 and len(polars)>0:  #只取pg波的时间
                        begin = np.clip(idx-50, 0, self.length-60)
                        #print(begin)
                        label_polar[0, begin:begin+100] = polars[0] 
                        label_quali[0, begin:begin+100] = polars[1]
                        label_weigh[0, begin:begin+100] = 1
            label2[0, :, 0] = np.clip(1-label2[0, :, 1]-label2[0, :, 2]-label2[0, :, 3]-label2[0, :, 4], 0, 1)   
            dqueue.put([rdata, [label_polar, label_quali, label_weigh]])

            count += 1

    def batch_data(self, batch_size=32):
        x1, x2, x3 = [], [], []
        p1, p2, p3 = [], [], [] 
        x4 = []
        for _ in range(batch_size):
            data,(label_polar, label_quali, label_weigh) = self.dqueue.get() 
            x1.append(data) 
            #x2.append(label1) 
            #x3.append(label2)
            #x4.append(etype) 
            p1.append(label_polar) 
            p2.append(label_quali) 
            p3.append(label_weigh)
        x1 = np.concatenate(x1, axis=0) 
        #x2 = np.concatenate(x2, axis=0) 
        #x3 = np.concatenate(x3, axis=0) 
        p1 = np.concatenate(p1, axis=0) 
        p2 = np.concatenate(p2, axis=0) 
        p3 = np.concatenate(p3, axis=0) 
        #x4 = np.array(x4)
        #print(p1)
        return  x1, p1, p2, p3










class DataLargeSNRTest():
    def __init__(self, file_name="models/h5test/all-gzip4.h5", n_length=10240, stride=1, padlen=256):
        self.file_name = file_name
        self.length = n_length 
        self.stride = stride  
        self.padlen = padlen 
        self.n_thread = 2
        self.phase_dict = {
            "Pg":0, 
            "Sg":1,
            #"P":0, 
            #"S":1,
            "Pn":2, 
            "Sn":3, 
        } 
        self.ploar1 = {
            "C":0, 
            "U":0, 
            "R":1, 
            "D":1, 
        }
        self.ploar2 = {
            "I":0, 
            "M":1, 
            "E":2, 
        }
        self.etype_dict = {
            "eq":0, 
            "ve":1, 
            "ss":2, 
            "sp":3, 
            "ep":4, 
            "ot":5, 
            "se":6
        }
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        #for year in [2020+i for i in range(3)]:
        for year in [2020]:
            multiprocessing.Process(target=self.feed_data, args=(fqueue, year)).start() 
        for _ in range(self.n_thread):
            multiprocessing.Process(target=self.process, args=(fqueue, self.dqueue)).start()
        #multiprocessing.Process(target=self.batch_data, args=(dqueue, )).start()
    def feed_data(self, fqueue, year):
        file_name = f"ayrdata/csndata/{year}.h5"
        h5keys = np.load(f"ayrdata/keys/{year}.npy")
        with open("large/distaz.pkl", "rb") as f:
            distaz = pickle.load(f) 
        while True:
            h5file = h5py.File(file_name, "r")
            print(f"{file_name}数据加载完成")
            for ekey in h5keys:
                #print(ekey)
                event = h5file[ekey]  #事件ID
                
                etype = event.attrs["type"] #事件类型
                if etype in self.etype_dict:
                    typeid = self.etype_dict[etype] #是否为可识别的事件
                else:
                    typeid = 7  #不是可识别的6种，则ID设为7
                for skey in event:
                    station = event[skey] #获取事件的台站信息
                    
                    data = [0, 0, 0] #初始化一个长度为3的列表，用于存储台站的数据
                    #print(skey)
                    
                    for dkey in station:
                        if "HZ" in dkey:
                            idx = 2 
                        elif "HE" in dkey:
                            idx = 0 
                        elif "HN" in dkey:
                            idx = 1 
                        else:
                            print(dkey, "Not exist")
                            continue 
                        data[idx] = station[dkey][:] 
                        #data.append(station[dkey][:]) 
                        btime = datetime.datetime.strptime(station[dkey].attrs['btime'], "%Y/%m/%d %H:%M:%S.%f")  #当前数据通道的属性中提取台站数据的起始时间
                    if len(data)!=3:continue  #data[HE数据，HN数据，HZ数据]
                    cnt = 0 
                    for d in data:
                        if type(d)==int:
                            cnt+= 1
                    if cnt!=0:continue 
                    phases = {}
                    phase_count = {"P":0, "S":0, "Pn":0, "Sn":0}
                    dist = -1
                    if "POLARITY.Pg.UPDOWN" in station.attrs and "POLARITY.Pg.CLARITY" in station.attrs:
                        ptype1 = station.attrs["POLARITY.Pg.UPDOWN"] #获取台站极性
                        ptype2 = station.attrs["POLARITY.Pg.CLARITY" ] #获取台站极性的清晰度
                        if ptype1 not in self.ploar1 or ptype2 not in self.ploar2:
                            polars = []
                        else:
                            polars = [self.ploar1[ptype1], self.ploar2[ptype2]]#获取极性信息和清晰度对应的数值
                    else:
                        polars = []
                    ptypes = [i.split(".")[-1] for i in station.attrs["types"].split(",")]
                    #print(ptypes)
                    if "Pg" not in ptypes :continue
                    if "Sg" not in ptypes :continue
                    if "Pn" not in ptypes :continue
                    if "Sn" not in ptypes :continue

                    for akey in station.attrs:
                        pkey = akey.split(".")[-1]
                        if f"{pkey}.dist" in station.attrs:dist = float(station.attrs[f"{pkey}.dist"])
                        if pkey not in ptypes:continue 
                        pname = pkey.split("+i")[-1].split("-i")[-1].split("+")[-1].split("-")[-1].split("i")[-1].split("2")[0].split("*")[-1]
                        if "Pg" in pname:
                            pname = akey.split(".")[-1]
                            if pname in self.phase_dict:
                                phase_count["P"] += 1
                                #print(akey, ptypes)
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                        else:
                            if pname in self.phase_dict:
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                                if akey == "P" or akey == "Pg":
                                    phase_count["P"] += 1
                                if akey == "S" or akey == "Sg":
                                    phase_count["S"] += 1
                                if akey in ["Pn", "Sn"]:
                                    phase_count[akey] += 1
                    if dist > 1000:continue 
                    #if phase_count["Pn"] !=0:
                    #    if phase_count["P"] ==0:
                    #        continue 
                    #if phase_count["Sn"] !=0:
                    #    if phase_count["S"] ==0:
                    #        continue 
                    #print(data)
                    if len(phases)==0:continue 
                    fqueue.put([data, btime, phases, polars, typeid])

    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, btime, phases, polars, etype = fqueue.get()
            pidx = {}
            plist = []
            for pkey in phases:
                ptime = phases[pkey]
                delta = (ptime-btime).total_seconds()
                delta_idx = int(delta * 100)
                pidx[pkey] = delta_idx
                plist.append(delta_idx) 
            cidx = np.random.choice(plist) - np.random.randint(self.padlen, self.length-self.padlen)
            rdata = [] 
            flen = False
            for d in data:
                w = d[cidx:cidx+self.length] 
                wlen = len(w) 
                if wlen!=self.length:
                    flen = True
                    break 
                w = w - np.mean(w)
                w = w / (np.std(w)+1e-6)
                w = w / (np.max(np.abs(w))+1e-6)
                rdata.append(w[np.newaxis, :, np.newaxis]) 
            if flen:
                continue
            rdata = np.concatenate(rdata, axis=2) 
            label1 = np.zeros([1, llen, 2])
            label2 = np.zeros([1, self.length, 5]) 
            label_polar = np.zeros([1, self.length]) 
            label_quali = np.zeros([1, self.length]) 
            label_weigh = np.zeros([1, self.length])
            for pkey in pidx:
                pid = self.phase_dict[pkey] 
                idx = (pidx[pkey] - cidx)//self.stride 
                
                if idx-1>0:
                    label1[0, idx-1:idx+2] = -1
                if idx > 0 and idx < llen:
                    label1[0, idx, 0] = pid + 1
                    label1[0, idx, 1] = (pidx[pkey] - cidx)%self.stride
            phase_intv = {"P":0, "S":0} 
            def norm(t, mu, std=0.1):
                p = np.exp(-(t-mu)**2/std**2/2)
                p /= (np.max(p)+1e-6) 
                return p 
            t = np.arange(self.length) * 0.01 
            phase_time = {0:-1, 1:-1, 2:-1, 3:-1}
            for pkey in pidx:
                pid = self.phase_dict[pkey] 
                idx = (pidx[pkey] - cidx)
                if idx > 0 and idx < self.length:
                    phase_time[pid] = idx    
            dqueue.put([rdata, [phase_time[0], phase_time[1], phase_time[2], phase_time[3]]])

            count += 1

    def batch_data(self, batch_size=50):
        x1, x2 = [], []
        for _ in range(batch_size):
            data, label1 = self.dqueue.get() 
            x1.append(data) 
            x2.append(label1) 
            
        x1 = np.concatenate(x1, axis=0) 
        return x1, x2