code/utils/datapnsnv3.py

import os 
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 

import multiprocessing 
import h5py 
import datetime 
import numpy as np 
import matplotlib.pyplot as plt 
plt.switch_backend("agg")
def _label(a=0, b=20, c=40): 
    'Used for triangolar labeling'
    
    z = np.linspace(a, c, num = 2*(b-a)+1)
    y = np.zeros(z.shape)
    y[z <= a] = 0
    y[z >= c] = 0
    first_half = np.logical_and(a < z, z <= b)
    y[first_half] = (z[first_half]-a) / (b-a)
    second_half = np.logical_and(b < z, z < c)
    y[second_half] = (c-z[second_half]) / (c-b)
    return y

class Data():
    def __init__(self, file_name="models/h5test/all-gzip4.h5", n_length=10240, stride=16, 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,
        } 
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        for year in [2009+i for i in range(11)]:
            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"/home/yuzy/machinelearning/makeh5/scdata/sc.fortest.h5"
        while True:
            h5file = h5py.File(file_name, "r")
            for ekey in h5file:
                event = h5file[ekey] 
                for skey in event:
                    station = event[skey] 
                    data = []
                    for dkey in station:
                        data.append(station[dkey][:]) 
                        btime = datetime.datetime.strptime(station[dkey].attrs['btime'], "%Y/%m/%d %H:%M:%S.%f") 
                    if len(data)!=3:continue 
                    phases = {}
                    phase_count = {"P":0, "S":0}
                    dist = -1
                    for akey in station.attrs:
                        if "dist" in akey:
                            dist = float(station.attrs[akey])
                        if "Pg" in akey:
                            pname = akey.split(".")[-1]
                            if pname in self.phase_dict:
                                phase_count["P"] += 1
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                        else:
                            if akey in self.phase_dict:
                                phases[akey] = 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 dist > 800:continue 
                    if phase_count["P"] == 0 or phase_count["S"] == 0:continue 
                    if len(phases)==0:continue 
                    fqueue.put([data, btime, phases])

    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, btime, phases = 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.max(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, 4]) 
            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 
            for pkey in pidx:
                pid = self.phase_dict[pkey] 
                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:
                    if idx < self.length:
                        phase_intv["P"] = np.max([idx, 0]) 
                    else:
                        phase_intv["P"] = self.length 
                if pid == 1:
                    if idx > 0:
                        phase_intv["S"] = np.min([idx, self.length])
                    else:
                        phase_intv["S"] = 0 
            label2[0, :, 0] = np.clip(1-label2[0, :, 1]-label2[0, :, 2], 0, 1)
            label2[0, phase_intv["P"]:phase_intv["S"], 3] = 1    
            dqueue.put([rdata, label1, label2])

            count += 1

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



class DataPnSn():
    def __init__(self, file_name="models/h5test/all-gzip4.h5", n_length=10240, stride=16, 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":2, 
            #"S":3,
            "Pn":2, 
            "Sn":3, 
        } 
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        for year in [2009+i for i in range(11)]:
            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"
        while True:
            h5file = h5py.File(file_name, "r")
            h5key = np.load(f"ayrdata/keys/{year}.npy")
            np.random.shuffle(h5key) 
            for ekey in h5key:
                event = h5file[ekey] 
                for skey in event:
                    station = event[skey] 
                    data = []
                    for dkey in station:
                        data.append(station[dkey][:]) 
                        btime = datetime.datetime.strptime(station[dkey].attrs['btime'], "%Y/%m/%d %H:%M:%S.%f") 
                    if len(data)!=3:continue 
                    phases = {}
                    phase_count = {"P":0, "S":0, "Pn":0, "Sn":0}
                    dist = -1
                    for akey in station.attrs:
                        if "dist" in akey:
                            dist = float(station.attrs[akey])
                        if "Pg" in akey:
                            pname = akey.split(".")[-1]
                            if pname in self.phase_dict:
                                phase_count["P"] += 1
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                        else:
                            if akey in self.phase_dict:
                                phases[akey] = 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 
                    if len(phases)==0:continue 
                    fqueue.put([data, btime, phases])

    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, btime, phases = 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.max(np.abs(w))+1e-6)
                w = w * np.random.uniform(0.2, 2.5)
                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]) 
            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 
            for pkey in pidx:
                pid = self.phase_dict[pkey] 
                idx = (pidx[pkey] - cidx)
                if pid in [0, 1]:
                    std = 0.1 
                else:
                    std = 0.1 
                if idx > 0 and idx < self.length:
                    label2[0, :, pid+1] = norm(t, idx*0.01, std)

            label2[0, :, 0] = np.clip(1-label2[0, :, 1]-label2[0, :, 2]-label2[0, :, 3]-label2[0, :, 4], 0, 1)   
            dqueue.put([rdata, label1, label2])

            count += 1

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


class DataPnSnFixLocation():
    def __init__(self, file_name="models/h5test/all-gzip4.h5", n_length=10240, stride=16, 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":2, 
            #"S":3,
            "Pn":2, 
            "Sn":3, 
        } 
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        for year in [2009+i for i in range(11)]:
            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"
        while True:
            h5file = h5py.File(file_name, "r")
            h5key = np.load(f"ayrdata/keys/{year}.npy")
            np.random.shuffle(h5key) 
            for ekey in h5key:
                event = h5file[ekey] 
                for skey in event:
                    station = event[skey] 
                    data = []
                    for dkey in station:
                        data.append(station[dkey][:]) 
                        btime = datetime.datetime.strptime(station[dkey].attrs['btime'], "%Y/%m/%d %H:%M:%S.%f") 
                    if len(data)!=3:continue 
                    phases = {}
                    phase_count = {"P":0, "S":0, "Pn":0, "Sn":0}
                    dist = -1
                    for akey in station.attrs:
                        if "dist" in akey:
                            dist = float(station.attrs[akey])
                        if "Pg" in akey:
                            pname = akey.split(".")[-1]
                            if pname in self.phase_dict:
                                phase_count["P"] += 1
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                        else:
                            if akey in self.phase_dict:
                                phases[akey] = 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 
                    if len(phases)==0:continue 
                    fqueue.put([data, btime, phases])

    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, btime, phases = fqueue.get()
            pidx = {}
            
            #if "Pn" not in phases:continue 
            pg_time_idx = 10000000
            for pkey in phases:
                ptime = phases[pkey]
                delta = (ptime-btime).total_seconds()
                delta_idx = int(delta * 100)
                pidx[pkey] = delta_idx
                if pg_time_idx > delta_idx:
                    pg_time_idx = delta_idx
            if pg_time_idx > self.length:
                continue
            cidx = pg_time_idx - np.random.randint(self.padlen, self.padlen+1024)
            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.max(np.abs(w))+1e-6)
                w = w * np.random.uniform(0.2, 2.5)
                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]) 
            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 
            for pkey in pidx:
                pid = self.phase_dict[pkey] 
                idx = (pidx[pkey] - cidx)
                if pid in [0, 1]:
                    std = 0.15 
                else:
                    std = 0.6 
                if idx > 0 and idx < self.length:
                    label2[0, :, pid+1] = norm(t, idx*0.01, std)

            label2[0, :, 0] = np.clip(1-label2[0, :, 1]-label2[0, :, 2]-label2[0, :, 3]-label2[0, :, 4], 0, 1)   
            dqueue.put([rdata, label1, label2])

            count += 1

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


class DataPnSnWithPolarType():
    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 [2009+i for i in range(11)]:
            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] 
                    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 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 
                    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)
                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 
            for pkey in pidx:
                pid = self.phase_dict[pkey] 
                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:
                        begin = np.clip(idx-50, 0, self.length-60)
                        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, label1, label2, [label_polar, label_quali, label_weigh], etype])

            count += 1

    def batch_data(self, batch_size=32):
        x1, x2, x3 = [], [], []
        p1, p2, p3 = [], [], [] 
        x4 = []
        for _ in range(batch_size):
            data, label1, label2, (label_polar, label_quali, label_weigh), etype = 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)
        return x1, x2, x3, p1, p2, p3, x4
    


class DataPnSnWithPolarTypeTest():
    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 [2009+i for i in range(11)]:
            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] 
                    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 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 
                    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)
                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 



class DataPnSnWithPolarTypeTele():
    def __init__(self,  
                 phase_name_file = "large/phase.type", 
                 polar_name_file = "large/polar.type", 
                 event_name_file = "large/event.type", 
                 locat_name_file = "large/station.loc", 
                 n_length=40960, 
                 padlen=4096, 
                 stride=128):
        self.n_thread = 2
        self.length = n_length 
        self.stride = stride  
        self.padlen = padlen 
        with open(phase_name_file, "r") as f:
            self.phase_dict = eval(f.read())
        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(11)]:
            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:
                if "CB." not in ekey:continue 
                #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] 
                    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 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 pname not in self.phase_dict:continue 
                        #print(pname)
                        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 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 
                    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.min(plist)-np.random.randint(2048, self.padlen)#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)
                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, 37]) 
            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 
            for pkey in pidx:
                pid = self.phase_dict[pkey] 
                idx = (pidx[pkey] - cidx)
                if idx > 0 and idx < self.length:
                    label2[0, :, (pid+1)%37] = norm(t, idx*0.01, 0.1)
                    if pid == 0 and len(polars)>0:
                        begin = np.clip(idx-50, 0, self.length-60)
                        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-np.sum(label2[0, :, 1:], axis=1), 0, 1)   
            dqueue.put([rdata, label1, label2, [label_polar, label_quali, label_weigh], etype])

            count += 1

    def batch_data(self, batch_size=32):
        x1, x2, x3 = [], [], []
        p1, p2, p3 = [], [], [] 
        x4 = []
        for _ in range(batch_size):
            data, label1, label2, (label_polar, label_quali, label_weigh), etype = 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)
        return x1, x2, x3, p1, p2, p3, x4



class DataTest():
    def __init__(self, file_name="", n_length=6144, stride=8, 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,
        } 
        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"h5data/{year}.h5"
        while True:
            h5file = h5py.File(file_name, "r")
            for ekey in h5file:
                event = h5file[ekey] 
                for skey in event:
                    station = event[skey] 
                    data = []
                    for dkey in station:
                        data.append(station[dkey][:]) 
                        btime = datetime.datetime.strptime(station[dkey].attrs['btime'], "%Y/%m/%d %H:%M:%S.%f") 
                    if len(data)!=3:continue 
                    phases = {}
                    phase_count = {"P":0, "S":0}
                    dist = -1
                    for akey in station.attrs:
                        if "dist" in akey:
                            dist = float(station.attrs[akey])
                        if "Pg" in akey:
                            pname = akey.split(".")[-1]
                            if pname in self.phase_dict:
                                phase_count["P"] += 1
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                        else:
                            if akey in self.phase_dict:
                                phases[akey] = 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 dist > 800:continue 
                    if phase_count["P"] == 0 or phase_count["S"] == 0:continue 
                    if len(phases)==0:continue 
                    fqueue.put([data, btime, phases])

    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, btime, phases = 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.max(w)+1e-6)
                rdata.append(w[np.newaxis, :, np.newaxis]) 
            if flen:
                continue
            #rdata -= np.mean(rdata, axis=0, keepdims=True) 
            #rdata /= (np.max(np.abs(rdata))+1e-6)
            rdata = np.concatenate(rdata, axis=2) 
            phase_time = {0:-1, 1:-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]]])

            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 



class DataTestPnSn():
    def __init__(self, file_name="", n_length=6144, stride=8, 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":2, 
            #"S":3,
            "Pn":2, 
            "Sn":3, 
        } 
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        for year in [2020, 2021, 2022]:
            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")
        if True:
            h5file = h5py.File(file_name, "r")
            h5file = h5py.File(file_name, "r")
            h5keys = np.load(f"ayrdata/keys/{year}.npy")
            np.random.shuffle(h5keys) 
            for ekey in h5keys:
                event = h5file[ekey] 
                for skey in event:
                    #sskey = skey.split(".")[0]
                    #print(sskey)
                    #if sskey not in ["SC", "YN", "XZ"]:continue
                    station = event[skey] 
                    data = []
                    for dkey in station:
                        data.append(station[dkey][:]) 
                        btime = datetime.datetime.strptime(station[dkey].attrs['btime'], "%Y/%m/%d %H:%M:%S.%f") 
                    if len(data)!=3:continue 
                    phases = {}
                    phase_count = {"P":0, "S":0}
                    dist = -1
                    for akey in station.attrs:
                        if "dist" in akey:
                            dist = float(station.attrs[akey])
                        if "Pg" in akey:
                            pname = akey.split(".")[-1]
                            if pname in self.phase_dict:
                                phase_count["P"] += 1
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                        else:
                            if akey in self.phase_dict:
                                phases[akey] = 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 dist > 2000:continue 
                    #if phase_count["P"] == 0 or phase_count["S"] == 0:continue 
                    if len(phases)==0:continue 
                    fqueue.put([data, btime, phases])

    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, btime, phases = 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.max(np.abs(w))+1e-6)
                rdata.append(w[np.newaxis, :, np.newaxis]) 
            if flen:
                continue
            #rdata -= np.mean(rdata, axis=0, keepdims=True) 
            #rdata /= (np.max(np.abs(rdata))+1e-6)
            rdata = np.concatenate(rdata, axis=2) 
            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 



class DataTestPnSnFixLocation():
    def __init__(self, file_name="", n_length=6144, stride=8, 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":2, 
            #"S":3,
            "Pn":2, 
            "Sn":3, 
        } 
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        for year in [2020, 2021, 2022]:
            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")
        if True:
            h5file = h5py.File(file_name, "r")
            h5file = h5py.File(file_name, "r")
            h5keys = np.load(f"ayrdata/keys/{year}.npy")
            np.random.shuffle(h5keys) 
            for ekey in h5keys:
                event = h5file[ekey] 
                for skey in event:
                    #sskey = skey.split(".")[0]
                    #print(sskey)
                    #if sskey not in ["SC", "YN", "XZ"]:continue
                    station = event[skey] 
                    data = []
                    for dkey in station:
                        data.append(station[dkey][:]) 
                        btime = datetime.datetime.strptime(station[dkey].attrs['btime'], "%Y/%m/%d %H:%M:%S.%f") 
                    if len(data)!=3:continue 
                    phases = {}
                    phase_count = {"P":0, "S":0}
                    dist = -1
                    for akey in station.attrs:
                        if "dist" in akey:
                            dist = float(station.attrs[akey])
                        if "Pg" in akey:
                            pname = akey.split(".")[-1]
                            if pname in self.phase_dict:
                                phase_count["P"] += 1
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                        else:
                            if akey in self.phase_dict:
                                phases[akey] = 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 dist > 2000:continue 
                    #if phase_count["P"] == 0 or phase_count["S"] == 0:continue 
                    if len(phases)==0:continue 
                    fqueue.put([data, btime, phases])

    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, btime, phases = fqueue.get()
            pidx = {}
            plist = []
            #if "Pg" not in phases:continue 
            pg_time_idx = 100000000
            for pkey in phases:
                ptime = phases[pkey]
                delta = (ptime-btime).total_seconds()
                delta_idx = int(delta * 100)
                pidx[pkey] = delta_idx
                if pg_time_idx > delta_idx:
                    pg_time_idx = delta_idx
            
            cidx = pg_time_idx - np.random.randint(self.padlen, self.padlen+1024)
            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.max(np.abs(w))+1e-6)
                rdata.append(w[np.newaxis, :, np.newaxis]) 
            if flen:
                continue
            #rdata -= np.mean(rdata, axis=0, keepdims=True) 
            #rdata /= (np.max(np.abs(rdata))+1e-6)
            rdata = np.concatenate(rdata, axis=2) 
            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 


class DataTestPnSnFixLocationPad():
    def __init__(self, file_name="", n_length=6144, stride=8, 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":2, 
            #"S":3,
            "Pn":2, 
            "Sn":3, 
        } 
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        for year in [2020, 2021, 2022]:
            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")
        if True:
            h5file = h5py.File(file_name, "r")
            h5file = h5py.File(file_name, "r")
            h5keys = np.load(f"ayrdata/keys/{year}.npy")
            np.random.shuffle(h5keys) 
            for ekey in h5keys:
                event = h5file[ekey] 
                for skey in event:
                    #sskey = skey.split(".")[0]
                    #print(sskey)
                    #if sskey not in ["SC", "YN", "XZ"]:continue
                    station = event[skey] 
                    data = []
                    for dkey in station:
                        data.append(station[dkey][:]) 
                        btime = datetime.datetime.strptime(station[dkey].attrs['btime'], "%Y/%m/%d %H:%M:%S.%f") 
                    if len(data)!=3:continue 
                    phases = {}
                    phase_count = {"P":0, "S":0}
                    dist = -1
                    for akey in station.attrs:
                        if "dist" in akey:
                            dist = float(station.attrs[akey])
                        if "Pg" in akey:
                            pname = akey.split(".")[-1]
                            if pname in self.phase_dict:
                                phase_count["P"] += 1
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                        else:
                            if akey in self.phase_dict:
                                phases[akey] = 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 dist > 2000:continue 
                    #if phase_count["P"] == 0 or phase_count["S"] == 0:continue 
                    if len(phases)==0:continue 
                    fqueue.put([data, btime, phases])

    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, btime, phases = fqueue.get()
            pidx = {}
            plist = []
            #if "Pg" not in phases:continue 
            #pg_time_idx = 100000000
            for pkey in phases:
                ptime = phases[pkey]
                delta = (ptime-btime).total_seconds()
                delta_idx = int(delta * 100)
                pidx[pkey] = delta_idx
                #if pg_time_idx > delta_idx:
                #    pg_time_idx = delta_idx
                plist.append(delta_idx)
            
            for ptt in plist:
                for kkk in range(2):
                    if kkk == 0:
                        cidx = ptt - np.random.randint(0, self.padlen)
                    else:
                        cidx = ptt - np.random.randint(self.length-self.padlen, self.length)
                    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.max(np.abs(w))+1e-6)
                        rdata.append(w[np.newaxis, :, np.newaxis]) 
                    if flen:
                        continue
                    #rdata -= np.mean(rdata, axis=0, keepdims=True) 
                    #rdata /= (np.max(np.abs(rdata))+1e-6)
                    rdata = np.concatenate(rdata, axis=2) 
                    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 


class DataTestPnSnLongData():
    def __init__(self, file_name="", n_length=6144, stride=8, 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":2, 
            #"S":3,
            "Pn":2, 
            "Sn":3, 
        } 
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        for year in [2020, 2021, 2022]:
            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")
        if True:
            h5file = h5py.File(file_name, "r")
            h5file = h5py.File(file_name, "r")
            h5keys = np.load(f"ayrdata/keys/{year}.npy")
            np.random.shuffle(h5keys) 
            for ekey in h5keys:
                event = h5file[ekey] 
                for skey in event:
                    #sskey = skey.split(".")[0]
                    #print(sskey)
                    #if sskey not in ["SC", "YN", "XZ"]:continue
                    station = event[skey] 
                    data = []
                    for dkey in station:
                        data.append(station[dkey][:]) 
                        btime = datetime.datetime.strptime(station[dkey].attrs['btime'], "%Y/%m/%d %H:%M:%S.%f") 
                    if len(data)!=3:continue 
                    phases = {}
                    phase_count = {"P":0, "S":0}
                    dist = -1
                    for akey in station.attrs:
                        if "dist" in akey:
                            dist = float(station.attrs[akey])
                        if "Pg" in akey:
                            pname = akey.split(".")[-1]
                            if pname in self.phase_dict:
                                phase_count["P"] += 1
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                        else:
                            if akey in self.phase_dict:
                                phases[akey] = 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 dist > 2000:continue 
                    #if phase_count["P"] == 0 or phase_count["S"] == 0:continue 
                    if len(phases)==0:continue 
                    fqueue.put([data, btime, phases])

    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, btime, phases = fqueue.get()
            pidx = {}
            plist = []
            #if "Pg" not in phases:continue 
            #pg_time_idx = 100000000
            for pkey in phases:
                ptime = phases[pkey]
                delta = (ptime-btime).total_seconds()
                delta_idx = int(delta * 100)
                pidx[pkey] = delta_idx
                #if pg_time_idx > delta_idx:
                #    pg_time_idx = delta_idx
                plist.append(delta_idx)
            
            if True:
                if True:
                    
                    cidx = 10000
                    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.max(np.abs(w))+1e-6)
                        rdata.append(w[np.newaxis, :, np.newaxis]) 
                    if flen:
                        continue
                    #rdata -= np.mean(rdata, axis=0, keepdims=True) 
                    #rdata /= (np.max(np.abs(rdata))+1e-6)
                    rdata = np.concatenate(rdata, axis=2) 
                    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 

class DataEQT2():
    def __init__(self, file_name="models/h5test/all-gzip4.h5", n_length=10240, stride=16, 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,
        } 
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        for year in [2009+i for i in range(11)]:
            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"h5data/{year}.h5"
        while True:
            h5file = h5py.File(file_name, "r")
            for ekey in h5file:
                event = h5file[ekey] 
                for skey in event:
                    station = event[skey] 
                    data = []
                    for dkey in station:
                        data.append(station[dkey][:]) 
                        btime = datetime.datetime.strptime(station[dkey].attrs['btime'], "%Y/%m/%d %H:%M:%S.%f") 
                    if len(data)!=3:continue 
                    phases = {}
                    phase_count = {"P":0, "S":0}
                    dist = -1
                    for akey in station.attrs:
                        if "dist" in akey:
                            dist = float(station.attrs[akey])
                        if "Pg" in akey:
                            pname = akey.split(".")[-1]
                            if pname in self.phase_dict:
                                phase_count["P"] += 1
                                phases[pname] = datetime.datetime.strptime(station.attrs[akey], "%Y/%m/%d %H:%M:%S.%f")
                        else:
                            if akey in self.phase_dict:
                                phases[akey] = 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 dist > 800:continue 
                    if phase_count["P"] == 0 or phase_count["S"] == 0:continue 
                    if len(phases)==0:continue 
                    fqueue.put([data, btime, phases])

    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, btime, phases = 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.max(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, 4]) 
            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):
                midx = int(mu*100) 
                p = np.zeros_like(t) 
                bidx = np.max([0, midx-20])
                eidx = np.min([self.length, midx+21])
                lent = np.abs(eidx - bidx) 
                p[bidx:eidx] = _label()[:lent]
                return p 
            t = np.arange(self.length) * 0.01 
            for pkey in pidx:
                pid = self.phase_dict[pkey] 
                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:
                    if idx < self.length:
                        phase_intv["P"] = np.max([idx, 0]) 
                    else:
                        phase_intv["P"] = self.length 
                if pid == 1:
                    if idx > 0:
                        phase_intv["S"] = np.min([idx, self.length])
                    else:
                        phase_intv["S"] = 0 
            label2[0, :, 0] = np.clip(1-label2[0, :, 1]-label2[0, :, 2], 0, 1)
            label2[0, phase_intv["P"]:phase_intv["S"], 3] = 1    
            dqueue.put([rdata, label1, label2])

            count += 1

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



class DitingData():
    def __init__(self, file_name="h5data/diting/DiTing.v2.0.h5", n_length=10240, stride=16, 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,
        } 
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        self.epoch = multiprocessing.Value("d", 0.0)
        self.file_name = file_name 
        self.p1 = multiprocessing.Process(target=self.feed_data, args=(fqueue, self.epoch))
        self.p1.start() 
        self.p2s = []
        for _ in range(self.n_thread):
            p = multiprocessing.Process(target=self.process, args=(fqueue, self.dqueue))
            p.start() 
            self.p2s.append(p)
    def get_epoch(self):
        return self.epoch.value 
    def kill_all(self):
        self.p1.terminate() 
        for p in self.p2s:
            p.terminate() 
    def feed_data(self, fqueue, epoch):
        while True:
            h5file = h5py.File(self.file_name, "r")
            train = h5file["train"]
            for ekey in train:
                event = train[ekey] 
                for skey in event:
                    station = event[skey] 
                    data = station[:]
                    pt, st = station.attrs["p_pick"], station.attrs["s_pick"]
                    fqueue.put([data, [int(pt), int(st)]])
            h5file.close()
            epoch.value += 1 # 计算epoch 
    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, pidx = fqueue.get()
            pdic = {"P":pidx[0], "S":pidx[1]}
            # 做增强的，随机一个位置
            bidx = np.random.choice(pidx) - np.random.randint(self.padlen, self.length-self.padlen)
            eidx = bidx + self.length 
            rdata = np.zeros([self.length, 3])
            len_data = len(data)
            if bidx >= 0 and eidx < len_data:
                rdata = data[bidx:eidx, :] 
            if bidx < 0 and eidx < len_data:
                before = -bidx 
                rdata = np.pad(data[:eidx], ((before, 0), (0, 0)))
            if bidx > 0 and eidx >= len_data:
                after = eidx - len_data 
                rdata = np.pad(data[bidx:], ((0, after), (0, 0))) 
            if bidx < 0 and eidx >= len_data:
                after = eidx - len_data 
                before = -bidx 
                rdata = np.pad(data, ((before, after), (0, 0))) 
            rdata = rdata.astype(np.float32) 
            rdata -= np.mean(rdata, axis=0, keepdims=True) 
            rdata /= (np.max(np.abs(rdata))+1e-6)
            rdata *= np.random.uniform(0.5, 2)
            if len(rdata) != self.length:continue 
            label1 = np.zeros([1, llen, 2]) # LPPN 标签
            for pkey in pdic:
                pid = self.phase_dict[pkey] 
                idx = (pdic[pkey] - bidx)//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] = (pdic[pkey] - bidx)%self.stride
            
            def tri(t, mu, std=0.1):
                midx = int(mu*100) 
                p = np.zeros_like(t) 
                bidx = np.max([0, midx-20])
                eidx = np.min([self.length, midx+21])
                lent = np.abs(eidx - bidx) 
                p[bidx:eidx] = _label()[:lent]
                return p 
            def norm(t, mu, std=0.1):
                p = np.exp(-(t-mu)**2/std**2/2)
                p /= (np.max(p)+1e-9) 
                return p 
            t = np.arange(self.length) * 0.01 
            label2 = np.zeros([1, self.length, 4]) # 其他类型标签
            phase_intv = {"P":0, "S":0} 
            for pkey in pdic:
                pid = self.phase_dict[pkey] 
                idx = (pdic[pkey] - bidx)
                if idx > 0 and idx < self.length:
                    label2[0, :, pid+1] = norm(t, idx*0.01, 0.1)
                if pid == 0:
                    if idx < self.length:
                        phase_intv["P"] = np.max([idx, 0]) 
                    else:
                        phase_intv["P"] = self.length 
                if pid == 1:
                    if idx > 0:
                        idx = int(idx + (pdic["S"]-pdic["P"]) * 1.4) 
                        phase_intv["S"] = np.min([idx, self.length])
                    else:
                        phase_intv["S"] = 0 
            label2[0, :, 0] = np.clip(1-label2[0, :, 1]-label2[0, :, 2], 0, 1)
            label2[0, phase_intv["P"]:phase_intv["S"], 3] = 1    
            dqueue.put([rdata.astype(np.float32), label1, label2])
            count += 1

    def batch_data(self, batch_size=32):
        x1, x2, x3 = [], [], []
        for _ in range(batch_size):
            data, label1, label2 = self.dqueue.get() 
            x1.append(data) 
            x2.append(label1) 
            x3.append(label2)
        x1 = np.stack(x1, axis=0) 
        x2 = np.concatenate(x2, axis=0) 
        x3 = np.concatenate(x3, axis=0) 
        return x1, x2, x3
class DitingDataTest():
    def __init__(self, file_name="h5data/diting/DiTing.v2.0.h5", n_length=10240, stride=16, 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,
        } 
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        self.epoch = multiprocessing.Value("d", 0.0)
        self.file_name = file_name 
        self.p1 = multiprocessing.Process(target=self.feed_data, args=(fqueue, self.epoch))
        self.p1.start() 
        self.p2s = []
        for _ in range(self.n_thread):
            p = multiprocessing.Process(target=self.process, args=(fqueue, self.dqueue))
            p.start() 
            self.p2s.append(p)
    def get_epoch(self):
        return self.epoch.value 
    def kill_all(self):
        self.p1.terminate() 
        for p in self.p2s:
            p.terminate() 
    def feed_data(self, fqueue, epoch):
        while True:
            h5file = h5py.File(self.file_name, "r")
            train = h5file["valid"]
            for ekey in train:
                event = train[ekey] 
                for skey in event:
                    station = event[skey] 
                    data = station[:]
                    pt, st = station.attrs["p_pick"], station.attrs["s_pick"]
                    fqueue.put([data, [int(pt), int(st)]])
            epoch.value += 1 # 计算epoch 
    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, pidx = fqueue.get()
            pdic = {"P":pidx[0], "S":pidx[1]}
            # 做增强的，随机一个位置
            bidx = np.random.choice(pidx) - np.random.randint(self.padlen, self.length-self.padlen)
            eidx = bidx + self.length 
            rdata = np.zeros([self.length, 3])
            len_data = len(data)
            if bidx >= 0 and eidx < len_data:
                rdata = data[bidx:eidx, :] 
            if bidx < 0 and eidx < len_data:
                before = -bidx 
                rdata = np.pad(data[:eidx], ((before, 0), (0, 0)))
            if bidx > 0 and eidx >= len_data:
                after = eidx - len_data 
                rdata = np.pad(data[bidx:], ((0, after), (0, 0))) 
            if bidx < 0 and eidx >= len_data:
                after = eidx - len_data 
                before = -bidx 
                rdata = np.pad(data, ((before, after), (0, 0))) 
            rdata = rdata.astype(np.float32) 
            rdata -= np.mean(rdata, axis=0, keepdims=True) 
            rdata /= (np.max(np.abs(rdata))+1e-6)
            rdata *= np.random.uniform(0.5, 2)
            if len(rdata) != self.length:continue 
            phase_time = {0:-1, 1:-1}
            for pkey in pdic:
                pid = self.phase_dict[pkey] 
                idx = (pdic[pkey] - bidx)
                if idx > 0 and idx < self.length:
                    phase_time[pid] = idx    
            dqueue.put([rdata, [phase_time[0], phase_time[1]]])

            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.stack(x1, axis=0) 
        return x1, x2 


class DitingDataTestSNR():
    def __init__(self, file_name="h5data/diting/DiTing.v2.0.h5", n_length=10240, stride=16, 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,
        } 
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        self.epoch = multiprocessing.Value("d", 0.0)
        self.file_name = file_name 
        self.p1 = multiprocessing.Process(target=self.feed_data, args=(fqueue, self.epoch))
        self.p1.start() 
        self.p2s = []
        for _ in range(self.n_thread):
            p = multiprocessing.Process(target=self.process, args=(fqueue, self.dqueue))
            p.start() 
            self.p2s.append(p)
    def get_epoch(self):
        return self.epoch.value 
    def kill_all(self):
        self.p1.terminate() 
        for p in self.p2s:
            p.terminate() 
    def feed_data(self, fqueue, epoch):
        while True:
            h5file = h5py.File(self.file_name, "r")
            train = h5file["valid"]
            for ekey in train:
                event = train[ekey] 
                
                for skey in event:
                    station = event[skey] 
                    data = station[:]
                    pt, st = station.attrs["p_pick"], station.attrs["s_pick"]
                    mag = station.attrs["evmag"]
                    dis = station.attrs["dis"]
                    fqueue.put([data, [int(pt), int(st), mag, dis]])
            epoch.value += 1 # 计算epoch 
    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, pidx = fqueue.get()
            pt, st, mag, dis = pidx 
            pdic = {"P":pidx[0], "S":pidx[1]}
            # 做增强的，随机一个位置
            bidx = 1#np.random.choice(pidx) - np.random.randint(self.padlen, self.length-self.padlen)
            eidx = bidx + self.length 
            rdata = np.zeros([self.length, 3])
            len_data = len(data)
            if bidx >= 0 and eidx < len_data:
                rdata = data[bidx:eidx, :] 
            if bidx < 0 and eidx < len_data:
                before = -bidx 
                rdata = np.pad(data[:eidx], ((before, 0), (0, 0)))
            if bidx > 0 and eidx >= len_data:
                after = eidx - len_data 
                rdata = np.pad(data[bidx:], ((0, after), (0, 0))) 
            if bidx < 0 and eidx >= len_data:
                after = eidx - len_data 
                before = -bidx 
                rdata = np.pad(data, ((before, after), (0, 0))) 
            rdata = rdata.astype(np.float32) 
            rdata -= np.mean(rdata, axis=0, keepdims=True) 
            rdata /= (np.max(np.abs(rdata))+1e-6)
            rdata *= np.random.uniform(0.5, 2)
            if len(rdata) != self.length:continue 
            phase_time = {0:-1, 1:-1}
            phase_snr = {0:-1, 1:-1}
            for pkey in pdic:
                pid = self.phase_dict[pkey] 
                idx = (pdic[pkey] - bidx)
                if idx > 0 and idx < self.length:
                    phase_time[pid] = idx   
                    idx1 = np.clip(idx-50, 0, self.length) 
                    idx2 = np.clip(idx+50, 0, self.length)
                    pre = rdata[idx1:idx] 
                    aft = rdata[idx:idx2]
                    phase_snr[pid] = np.std(aft)/np.std(pre)
            dqueue.put([rdata, [phase_time[0], phase_time[1]], [phase_snr[0], phase_snr[1], mag, dis]])

            count += 1

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



class DitingDataForPlot():
    def __init__(self, file_name="h5data/diting/DiTing.v2.0.h5", n_length=10240, stride=16, 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,
        } 
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        self.epoch = multiprocessing.Value("d", 0.0)
        self.file_name = file_name 
        self.p1 = multiprocessing.Process(target=self.feed_data, args=(fqueue, self.epoch))
        self.p1.start() 
        self.p2s = []
        for _ in range(self.n_thread):
            p = multiprocessing.Process(target=self.process, args=(fqueue, self.dqueue))
            p.start() 
            self.p2s.append(p)
    def get_epoch(self):
        return self.epoch.value 
    def kill_all(self):
        self.p1.terminate() 
        for p in self.p2s:
            p.terminate() 
    def feed_data(self, fqueue, epoch):
        while True:
            h5file = h5py.File(self.file_name, "r")
            train = h5file["train"]
            for ekey in train:
                event = train[ekey] 
                for skey in event:
                    station = event[skey] 
                    data = station[:]
                    snr = station.attrs["Z_P_amplitude_snr"]
                    if snr < 10:continue 
                    pt, st = station.attrs["p_pick"], station.attrs["s_pick"]
                    fqueue.put([data, [int(pt), int(st)]])
            h5file.close()
            epoch.value += 1 # 计算epoch 
    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, pidx = fqueue.get()
            pdic = {"P":pidx[0], "S":pidx[1]}
            # 做增强的，随机一个位置
            bidx = 0#np.random.choice(pidx) - np.random.randint(self.padlen, self.length-self.padlen)
            eidx = bidx + self.length 
            rdata = np.zeros([self.length, 3])
            len_data = len(data)
            if bidx >= 0 and eidx < len_data:
                rdata = data[bidx:eidx, :] 
            if bidx < 0 and eidx < len_data:
                before = -bidx 
                rdata = np.pad(data[:eidx], ((before, 0), (0, 0)))
            if bidx > 0 and eidx >= len_data:
                after = eidx - len_data 
                rdata = np.pad(data[bidx:], ((0, after), (0, 0))) 
            if bidx < 0 and eidx >= len_data:
                after = eidx - len_data 
                before = -bidx 
                rdata = np.pad(data, ((before, after), (0, 0))) 
            rdata = rdata.astype(np.float32) 
            rdata -= np.mean(rdata, axis=0, keepdims=True) 
            rdata /= (np.max(np.abs(rdata))+1e-6)
            rdata *= np.random.uniform(0.5, 2)
            if len(rdata) != self.length:continue 
            label1 = np.zeros([1, llen, 2]) # LPPN 标签
            for pkey in pdic:
                pid = self.phase_dict[pkey] 
                idx = (pdic[pkey] - bidx)//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] = (pdic[pkey] - bidx)%self.stride# 回归
            
            def tri(t, mu, std=0.1):
                midx = int(mu*100) 
                p = np.zeros_like(t) 
                bidx = np.max([0, midx-20])
                eidx = np.min([self.length, midx+21])
                lent = np.abs(eidx - bidx) 
                p[bidx:eidx] = _label()[:lent]
                return p 
            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 
            label2 = np.zeros([1, self.length, 4]) # 其他类型标签
            phase_intv = {"P":0, "S":0} 
            for pkey in pdic:
                pid = self.phase_dict[pkey] 
                idx = (pdic[pkey] - bidx)
                if idx > 0 and idx < self.length:
                    label2[0, :, pid+1] = tri(t, idx*0.01, 0.1)
                if pid == 0:
                    if idx < self.length:
                        phase_intv["P"] = np.max([idx, 0]) 
                    else:
                        phase_intv["P"] = self.length 
                if pid == 1:
                    if idx > 0:
                        idx = int(idx + (pdic["S"]-pdic["P"]) * 1.4) 
                        phase_intv["S"] = np.min([idx, self.length])
                    else:
                        phase_intv["S"] = 0 
            label2[0, :, 0] = np.clip(1-label2[0, :, 1]-label2[0, :, 2], 0, 1)
            label2[0, phase_intv["P"]:phase_intv["S"], 3] = 1    
            dqueue.put([rdata.astype(np.float32), label1, label2])
            count += 1

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

class DitingDataTestForPlot():
    def __init__(self, file_name="h5data/diting/DiTing.v2.0.h5", n_length=10240, stride=16, 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,
        } 
        fqueue = multiprocessing.Queue(100)
        self.dqueue = multiprocessing.Queue(100)
        self.epoch = multiprocessing.Value("d", 0.0)
        self.file_name = file_name 
        self.p1 = multiprocessing.Process(target=self.feed_data, args=(fqueue, self.epoch))
        self.p1.start() 
        self.p2s = []
        for _ in range(self.n_thread):
            p = multiprocessing.Process(target=self.process, args=(fqueue, self.dqueue))
            p.start() 
            self.p2s.append(p)
    def get_epoch(self):
        return self.epoch.value 
    def kill_all(self):
        self.p1.terminate() 
        for p in self.p2s:
            p.terminate() 
    def feed_data(self, fqueue, epoch):
        while True:
            h5file = h5py.File(self.file_name, "r")
            train = h5file["valid"]
            count = 0 
            for ekey in train:
                event = train[ekey] 
                for skey in event:
                    station = event[skey] 
                    data = station[:]
                    snr = station.attrs["Z_P_amplitude_snr"]
                    if count !=0:
                        if snr < 10:continue 
                    pt, st = station.attrs["p_pick"], station.attrs["s_pick"]
                    fqueue.put([data, [int(pt), int(st)]])
                    count += 1 
            epoch.value += 1 # 计算epoch 
    def process(self, fqueue, dqueue):
        count = 0
        llen = self.length//self.stride
        while True:
            data, pidx = fqueue.get()
            pdic = {"P":pidx[0], "S":pidx[1]}
            # 做增强的，随机一个位置
            bidx = 0#np.random.choice(pidx) - np.random.randint(self.padlen, self.length-self.padlen)
            eidx = bidx + self.length 
            rdata = np.zeros([self.length, 3])
            len_data = len(data)
            if bidx >= 0 and eidx < len_data:
                rdata = data[bidx:eidx, :] 
            if bidx < 0 and eidx < len_data:
                before = -bidx 
                rdata = np.pad(data[:eidx], ((before, 0), (0, 0)))
            if bidx > 0 and eidx >= len_data:
                after = eidx - len_data 
                rdata = np.pad(data[bidx:], ((0, after), (0, 0))) 
            if bidx < 0 and eidx >= len_data:
                after = eidx - len_data 
                before = -bidx 
                rdata = np.pad(data, ((before, after), (0, 0))) 
            rdata = rdata.astype(np.float32) 
            rdata -= np.mean(rdata, axis=0, keepdims=True) 
            rdata /= (np.max(np.abs(rdata))+1e-6)
            rdata *= np.random.uniform(0.5, 2)
            if len(rdata) != self.length:continue 
            phase_time = {0:-1, 1:-1}
            for pkey in pdic:
                pid = self.phase_dict[pkey] 
                idx = (pdic[pkey] - bidx)
                if idx > 0 and idx < self.length:
                    phase_time[pid] = idx    
            dqueue.put([rdata, [phase_time[0], phase_time[1]]])

            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.stack(x1, axis=0) 
        return x1, x2 



import matplotlib.pyplot as plt 
if __name__ == "__main__":
    data = DitingData(n_length=6144, stride=8)
    
    for e in range(3):
        a1, a2, a3 = data.batch_data() 
        print(a1.shape, a2.shape, a3.shape, data.get_epoch()) 
    w = a1[0, :, 0] 
    #w /= np.max(w)
    plt.plot(w, c="k") 
    plt.plot(np.repeat(a2[0, :, 0], 8), c="r") 
    #plt.plot(a3[0, :, 0], c="r")
    plt.plot(a3[0, :, 1], c="g")
    plt.plot(a3[0, :, 2], c="b")
    #plt.plot(a3[0, :, 3], c="c")
    plt.savefig("temp/demo.jpg")
    data.kill_all()