import os
import glob
import obspy
from obspy import UTCDateTime,read
from obspy.io.sac.sactrace import SACTrace
import shutil
import re
import matplotlib.pyplot as plt
from seisloc.utils import get_st
from seisloc.geometry import spherical_dist
from seisloc.text_io import load_y2000
from seisloc.sta import load_sta
import numpy as np
from tqdm import tqdm
import pickle
import logging
from obspy.taup import TauPyModel

def cut_eve_wf(phs_dict="out.arc",
               sta_dict="sta.txt",
               src_root="day_data",
               tar_root="eve_wf",
               tb=-20,te=60,
               shift_hour=0):
    """
    Parameters:
          phs_dict: HYPOINVERSE output archive file
          sta_dict: station file in format "net sta lon lat ele label"
          src_root: root path for input waveform data
          tar_root: output for event waveform data
                tb: time before event time
                te: time after event time
        shift_hour: shift hours for event time zone, +8 for Beijing time zone
    """
    eve_list = phs_dict.keys()
    print(">>> cut event waveform ... ")
    for eve in tqdm(eve_list):
        str_day = eve[:8]
        #create day folder
        if os.path.exists(os.path.join(tar_root,str_day)):
            pass
        else:
            os.makedirs(os.path.join(tar_root,str_day))
        if os.path.exists(os.path.join(tar_root,str_day,eve)):
            shutil.rmtree(os.path.join(tar_root,str_day,eve))
        os.makedirs(os.path.join(tar_root,str_day,eve))
        e_time = UTCDateTime.strptime(eve,"%Y%m%d%H%M%S%f")-shift_hour*60*60
        starttime = e_time + tb
        endtime = e_time + te
        phases = phs_dict[eve]["phase"]
        processed_sta = []
        for [net,sta,pha,rel_time] in phases:
            src_folder=os.path.join(src_root,sta)
            if sta not in processed_sta:
                processed_sta.append(sta)
                st = get_st(net,sta,starttime,endtime,src_folder)
                if len(st)==0:
                    continue
                for trace in st:
                    chn = trace.stats.channel
                    starttime = trace.stats.starttime
                    b_value = starttime - e_time
                    out_file = os.path.join(tar_root,str_day,eve,f"{sta}.{chn}")
                    trace.write(out_file,format="SAC")
                    sac = SACTrace.read(out_file,headonly=True)
                    net = sac.knetwk
                    sac.lovrok = 1
                    sac.evlo = phs_dict[eve]["eve_loc"][0]
                    sac.evla = phs_dict[eve]["eve_loc"][1]
                    sac.evdp = phs_dict[eve]["eve_loc"][2]
                    sac.stlo = sta_dict[net][sta][0]
                    sac.stla = sta_dict[net][sta][1]
                    sac.stel = sta_dict[net][sta][2]
                    sac.nzyear = e_time.year
                    sac.nzjday = e_time.julday
                    sac.nzhour = e_time.hour
                    sac.nzmin = e_time.minute
                    sac.nzsec = e_time.second
                    sac.nzmsec = e_time.microsecond/1000
                    sac.b = b_value
                    sac.o = 0
                    sac.write(out_file)
                    
def write_eve_wf(phs_dict,tar_root,depth=2):
    '''
    depth = 2 means tar_root/str_day/eve/sta.*
    depth = 1 means tar_root/eve/sta.*
    '''
    eve_list = phs_dict.keys()
    print(">>> write arrival time information ...")
    for eve in tqdm(eve_list):
        str_day = eve[:8]
        phases = phs_dict[eve]["phase"]
        for [net,sta,pha,rel_time] in phases:
            if depth == 2:
                sac_files = glob.glob(os.path.join(tar_root,str_day,eve,sta+"*"))
            if depth == 1:
                sac_files = glob.glob(os.path.join(tar_root,eve,sta+"*"))
 
            if len(sac_files) == 0:
                continue
            if pha == "P":
                for sac_file in sac_files:
                    sac = SACTrace.read(sac_file,headonly=True)
                    sac.a = rel_time
                    logging.info(f">>> write in {eve} {sta} with P arrival {sac.a}")
                    sac.write(sac_file)
            if pha == "S":
                for sac_file in sac_files:
                    sac = SACTrace.read(sac_file,headonly=True)
                    sac.t0 = rel_time
                    logging.info(f">>> write in {eve} {sta} with S arrival {sac.t0}")
                    sac.write(sac_file)

def gen_eve_wf(y2000_file,sta_file,src_root="day_data",tar_root="eve_wf",tb = -20,te = 60,shift_hour=0):
    """
    cut event waveform from continous waveform based on phase file generated by 
    HYPOINVERSE
 
    Parameters:
    y2000_file: phase file, which is the output from HYPOINVERSE
    sta_file: station file in format: net sta lon lat ele
    src_root: source path for continous waveform
    tar_root: target path of eve based waveform
    tb: waveform begin time in relative to the event occurence time
    te: waveform end time in relative to the event ouccurence time
    """
    if not os.path.exists(tar_root):
        os.mkdir(tar_root)
    base_path = os.path.abspath("./")
    sta_dict = load_sta(sta_file)
    phs_dict = load_y2000(y2000_file)
 
    cut_eve_wf(phs_dict,sta_dict,src_root,tar_root,tb,te,shift_hour=shift_hour)
    write_eve_wf(phs_dict,tar_root)

def get_e_time(tr):
    """
    Applied to a sac trace. This function first read sac reference from from sac head.
    Then read o_value from the sac head. Return is the event time
    """
    sac = SACTrace.from_obspy_trace(tr)
    sac_ref_time = sac.reftime    # In UTCDateTime format
    o_value = tr.stats.sac.o 
    event_time = sac_ref_time + o_value     # event origin time
    
    return event_time

def wf_dist_plot(st,
                 xb=0,
                 xe=10,
                 dist_threshold=50,
                 scaling_factor=2,
                 color=None,
                 label_sta=True,
                 out_format="PNG",
                 out_folder="./",
                 figsize=(8,10),
                 align=None,
                 normalize=True):
    '''    
    Description
    ------------
    Plot event waveform by distance. The start time is event origin time.
    Data should be in sac format. The output is a saved file with the title
    of reference time.

    Parameters
    -----------
                st: obspy Stream object
                xb: begin of x in reference to the event time, default 0
                xe: end of x in reference to the event time, default 10
    dist_threshold: set maximum distance for plot
    scaling_factor: The waveform are normalized, increase scaling_facotr to
                    make the waveform plot more obvious
             color: The usage of color is the same as matplotlib.pyplot. 
                    Using default color if not defined.
         label_sta: Whether label station name on the plot.
        out_format: "png","jpg","pdf"..., The same as matplotlib.pyplot.savefig
        out_folder: output folder for the plot file

    Below data information needed:
    |   P arrival: tr.stats.sac.a
    |   S arrival: tr.stats.sac.t0
    |        evla: tr.stats.sac.evla
    |        evlo: tr.stats.sac.evlo
    |        stla: tr.stats.sac.stla
    |        stlo: tr.stats.sac.stlo
    '''
    
    # >>>>>>>>>>>>> Initiation >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    st.detrend("linear")
    st.detrend("constant")
    try:
        e_mag = st[0].stats.sac.mag
    except:
        e_mag = -9
        
    starttime=st[0].stats.starttime
    endtime =st[0].stats.endtime
    sampling_rate = st[0].stats.sampling_rate
    
    event_time = get_e_time(st[0])

    if len(st) == 0:
        print("Warning: Nothing to plot!")

    for tr in st:
        evla = tr.stats.sac.evla
        evlo = tr.stats.sac.evlo
        stla = tr.stats.sac.stla
        stlo = tr.stats.sac.stlo
        #It is recommend to set tr.stats.distance in meters by osbpy guideline
        tr.stats.distance = spherical_dist(evlo,evla,stlo,stla)*111*1000
    
    #Get mininumtime, maximum time, and max distance
    min_time = st[0].stats.starttime
    max_time = st[0].stats.endtime
    max_dist = st[0].stats.distance/1000 # in km
    min_dist = st[0].stats.distance/1000 # in km
    min_dist_idx = 0
    max_dist_idx = 0
    for i,tr in enumerate(st[1:]):
        if tr.stats.starttime < min_time:
            min_time = tr.stats.starttime
        if tr.stats.endtime > max_time:
            max_time = tr.stats.endtime
        if tr.stats.distance/1000 > max_dist:
            max_dist = tr.stats.distance/1000
            max_dist_idx = i
        if tr.stats.distance/1000 < min_dist:
            min_dist = tr.stats.distance/1000
            min_dist_idx = i
    if max_dist>dist_threshold:
        max_dist = dist_threshold
    
    # >>>>>>>>>>>>>>> Plot >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    fig,ax = plt.subplots(1,1,figsize = figsize)
    plt.xlim(xb,xe)
    plt.ylim(0,max_dist*1.1)
    plt.xticks(size=12)
    plt.yticks(size=12)
    plt.xlabel("Time (s)",fontdict={'size':16})
    plt.ylabel("Distance (km)",fontdict={'size':16})
       
    for tr in st:            # Plot trace by trace
        sta = tr.stats.station
        sac = SACTrace.from_obspy_trace(tr)
        tr_ref_time = sac.reftime
        tr_o_value = tr.stats.sac.o
        event_time = tr_ref_time + tr_o_value
        x_start = min_time-event_time
        dist = tr.stats.distance/1000
        
        if dist>max_dist:
            continue
        
        disp_data = tr.data
        if normalize==True:  # Normalize the waveform
            disp_data = tr.data/(max(tr.data) - min(tr.data))
        disp_data = disp_data*scaling_factor
        
        # >>>>>>>>>>>>>> Get align_shift for alignment >>>>>>>>>>>>>>>>>>
        if align == None:
            align_shift = 0
        else:
            align_shift = -tr.stats.sac.__getattribute__(align)

        
        if color == None:
            plt.plot(np.arange(0,len(tr.data))/sampling_rate+x_start+align_shift,
                    disp_data+dist,
                    linewidth = 0.5)
        else:
            plt.plot(np.arange(0,len(tr.data))/sampling_rate+x_start+align_shift,
                    disp_data+dist,
                    color = color,
                    linewidth = 0.5)
        if label_sta:
            plt.text(0.1+xb,dist+0.02,sta,fontsize=12)
        
        #>>>>>>>>>>>>>>>>>>>>> Plot Phase Markers >>>>>>>>>>>>>>>>>>>>>>>>>>>>
        # Plot P arrival if available
        try:
            a = tr.stats.sac.a - tr.stats.sac.o
            gap = 0.5*max_dist/25
            plt.plot([a+align_shift,a+align_shift],[dist-gap,dist+gap],color='b',linewidth=2)
        except:
            pass
        
        # Plot S arrival if available
        try:
            t0 = tr.stats.sac.t0 - tr.stats.sac.o
            gap = 0.5*max_dist/25
            plt.plot([t0+align_shift,t0+align_shift],[dist-gap,dist+gap],color='r',linewidth=2)
        except:
            pass
        # Plot t5 arrival if available
        try:
            t5 = tr.stats.sac.t5 - tr.stats.sac.o
            gap = 0.5*max_dist/25
            plt.plot([t5+align_shift,t5+align_shift],[dist-gap,dist+gap],color='c',linewidth=2)
        except:
            pass
        
        # >>>>> Plot title with magnitude if available >>>>>>>>>>>>>>>>>>>>>>>
        if e_mag != -9:
            plt.title(str(tr_ref_time)+f"_M{e_mag}",fontdict={'size':18})
        else:
            plt.title(str(tr_ref_time),fontdict={'size':18})
    
    # >>>>>>>> Save fig and release memory >>>>>>>>>>>>>>>>>>>>>>
    plt.savefig(os.path.join(out_folder,f"{event_time.strftime('%Y%m%d%H%M%S')}."+out_format.lower()),format=out_format)
    plt.close(fig)
