import plotly.graph_objects as go
import pandas as pd
import pydeck as pdk
from obspy import UTCDateTime
import numpy as np  # 我們需要 numpy 來輔助定位註釋

def plot_waveforms_with_picks(p_picks):
    """
    為每個成功撿拾到P波的測站繪製波形圖並標示P波。
    (已修正 TypeError)
    """
    fig = go.Figure()
    
    # 為了讓Y軸範圍更好看，找到所有波形的最大振幅
    max_amplitude = 0
    if p_picks:
        max_amplitude = max(np.max(np.abs(p['trace'].data)) for p in p_picks)

    for i, pick in enumerate(p_picks):
        tr = pick['trace']
        time_axis = [t.datetime for t in tr.times("utcdatetime")]
        
        fig.add_trace(go.Scatter(
            x=time_axis,
            y=tr.data,
            mode='lines',
            name=tr.stats.station
        ))
        
        # --- 變更開始 ---
        # 步驟 1: 畫垂直線 (不帶註釋)
        fig.add_vline(
            x=pick['p_arrival_time'].datetime,
            line_width=1.5,
            line_dash="dash",
            line_color="red"
        )
        
        # 步驟 2: 單獨添加文字註釋，以獲得更好的控制並避免 bug
        fig.add_annotation(
            x=pick['p_arrival_time'].datetime,
            y=max_amplitude * 0.9,  # 將註釋放在圖表頂部附近
            yref="y",
            text=f"<b>{tr.stats.station} P</b>", # 使用簡短、粗體的文字
            showarrow=False,
            font=dict(color="red", size=11),
            bgcolor="rgba(255, 255, 255, 0.7)", # 半透明背景
            xshift=5, # 稍微向右移動以避免與線重疊
            yshift=15  # 稍微向上移動
        )
        # --- 變更結束 ---
        
    fig.update_layout(
        title="測站波形與 P 波初達標示",
        xaxis_title="時間 (UTC)",
        yaxis_title="振幅 (Counts)",
        legend_title="測站",
        height=500,
        yaxis_range=[-max_amplitude*1.1, max_amplitude*1.1] # 設定一個對稱且略大的Y軸範圍
    )
    return fig


def plot_event_map(inventory, p_picks, location_result):
    """
    在地圖上呈現測站、P波撿拾站點及震央位置。
    (此函式無須變更)
    """
    station_data = []
    picked_stations_codes = [p['station'] for p in p_picks]

    for net in inventory:
        for sta in net:
            if sta.code in picked_stations_codes:
                color = [0, 0, 255, 160] # Blue
            else:
                color = [128, 128, 128, 80] # Grey
                
            station_data.append({
                "name": f"測站: {sta.code}",
                "coordinates": [sta.longitude, sta.latitude],
                "color": color,
            })
            
    station_df = pd.DataFrame(station_data)

    if location_result:
        epicenter_df = pd.DataFrame([{
            "name": f"震央\n緯度: {location_result['latitude']:.3f}\n經度: {location_result['longitude']:.3f}\n深度: {location_result['depth_km']:.1f} km",
            "coordinates": [location_result['longitude'], location_result['latitude']],
            "color": [255, 0, 0, 255], # Red
        }])

    station_layer = pdk.Layer(
        "ScatterplotLayer",
        data=station_df,
        get_position="coordinates",
        get_fill_color="color",
        get_radius=200,
        pickable=True,
        auto_highlight=True,
    )
    
    layers = [station_layer]
    
    if location_result:
        epicenter_layer = pdk.Layer(
            "ScatterplotLayer",
            data=epicenter_df,
            get_position="coordinates",
            get_fill_color="color",
            get_radius=500,
            pickable=True,
        )
        layers.append(epicenter_layer)

    view_state = pdk.ViewState(
        latitude=23.9,
        longitude=121.0,
        zoom=6.5,
        pitch=45,
    )

    tooltip = {"text": "{name}"}

    deck = pdk.Deck(layers=layers, initial_view_state=view_state, map_style='mapbox://styles/mapbox/light-v9', tooltip=tooltip)
    return deck