import streamlit as st
import plotly.graph_objects as go
import numpy as np
import pandas as pd
from openai import OpenAI

# 初始化 tab2 專屬的 session state
def init_tab2_session_state():
    if 'tab2_chat_history' not in st.session_state:
        st.session_state.tab2_chat_history = []
    # 有加強劑的參數
    if 'tab2_booster_ct' not in st.session_state:
        st.session_state.tab2_booster_ct = 18.0
    if 'tab2_booster_days' not in st.session_state:
        st.session_state.tab2_booster_days = 7
    # 沒有加強劑的參數
    if 'tab2_no_booster_ct' not in st.session_state:
        st.session_state.tab2_no_booster_ct = 18.0
    if 'tab2_no_booster_days' not in st.session_state:
        st.session_state.tab2_no_booster_days = 7

# 基礎數據：從 Table 2 論文原始數據
BASE_DATA = {
    # Ct <= 18, 無加強劑
    'ct18_no_booster': {
        1: 0.10, 2: 0.25, 3: 0.39, 4: 0.51, 5: 0.61,
        6: 0.70, 7: 0.76, 8: 0.81, 9: 0.84, 10: 0.87,
        11: 0.89, 12: 0.91, 13: 0.93, 14: 0.94, 15: 0.95,
        16: 0.96, 17: 0.96, 18: 0.97, 19: 0.97, 20: 0.98, 21: 0.98
    },
    # Ct <= 18, 有加強劑
    'ct18_booster': {
        1: 0.44, 2: 0.65, 3: 0.77, 4: 0.84, 5: 0.89,
        6: 0.92, 7: 0.94, 8: 0.95, 9: 0.97, 10: 0.97,
        11: 0.97, 12: 0.98, 13: 0.98, 14: 0.99, 15: 0.99,
        16: 0.99, 17: 0.99, 18: 0.99, 19: 0.99, 20: 1.00, 21: 1.00
    },
    # Ct 18-25, 無加強劑
    'ct22_no_booster': {
        1: 0.49, 2: 0.57, 3: 0.65, 4: 0.72, 5: 0.78,
        6: 0.83, 7: 0.86, 8: 0.89, 9: 0.91, 10: 0.92,
        11: 0.94, 12: 0.95, 13: 0.96, 14: 0.97, 15: 0.97,
        16: 0.98, 17: 0.98, 18: 0.98, 19: 0.98, 20: 0.99, 21: 0.99
    },
    # Ct 18-25, 有加強劑
    'ct22_booster': {
        1: 0.68, 2: 0.80, 3: 0.87, 4: 0.91, 5: 0.94,
        6: 0.95, 7: 0.96, 8: 0.97, 9: 0.98, 10: 0.98,
        11: 0.98, 12: 0.99, 13: 0.99, 14: 0.99, 15: 0.99,
        16: 1.00, 17: 1.00, 18: 1.00, 19: 1.00, 20: 1.00, 21: 1.00
    }
}

def get_effectiveness_from_data(data_key, days):
    """從基礎數據中獲取效益值（帶線性插值）"""
    data = BASE_DATA[data_key]
    
    if days < 1:
        return 0.0
    
    if days > 21:
        return data[21]
    
    if days in data:
        return data[days]
    
    # 線性插值
    day_lower = int(days)
    day_upper = day_lower + 1
    
    if day_upper > 21:
        return data[21]
    
    eff_lower = data[day_lower]
    eff_upper = data[day_upper]
    
    ratio = days - day_lower
    effectiveness = eff_lower * (1 - ratio) + eff_upper * ratio
    
    return effectiveness

def get_quarantine_effectiveness(ct_value, days, has_booster):
    """獲取隔離效益"""
    if ct_value > 25:
        return 0.0
    
    booster_suffix = 'booster' if has_booster else 'no_booster'
    
    if ct_value <= 18:
        data_key = f'ct18_{booster_suffix}'
        effectiveness = get_effectiveness_from_data(data_key, days)
        
        if ct_value < 10:
            boost_factor = 1.0 + (10 - ct_value) * 0.005
            effectiveness = min(effectiveness * boost_factor, 1.0)
        
        return effectiveness
    
    else:
        ct18_key = f'ct18_{booster_suffix}'
        ct22_key = f'ct22_{booster_suffix}'
        
        eff_ct18 = get_effectiveness_from_data(ct18_key, days)
        eff_ct22 = get_effectiveness_from_data(ct22_key, days)
        
        ratio = (ct_value - 18) / 7
        effectiveness = eff_ct18 * (1 - ratio) + eff_ct22 * ratio
        
        return min(effectiveness, 1.0)

def get_ct_category(ct_value):
    """獲取 Ct 值分類"""
    if ct_value <= 18:
        return "高病毒量 (Ct ≤ 18)"
    elif ct_value <= 25:
        return "中病毒量 (18 < Ct ≤ 25)"
    else:
        return "低病毒量 (Ct > 25，假設無傳染性)"

def create_3d_plot(ct_mesh, day_mesh, effectiveness_mesh, selected_ct, quarantine_days, 
                   user_eff, has_booster, title, colorscale, marker_color):
    """創建單一 3D 曲面圖"""
    fig = go.Figure()

    fig.add_trace(go.Surface(
        x=ct_mesh,
        y=day_mesh,
        z=effectiveness_mesh,
        colorscale=colorscale,
        showscale=True,
        colorbar=dict(
            title="防疫效益",
            tickvals=[0, 0.25, 0.5, 0.75, 1.0],
            ticktext=['0%', '25%', '50%', '75%', '100%']
        ),
        opacity=0.9,
        name=title,
        contours=dict(
            x=dict(show=True, color='white', width=1),
            y=dict(show=True, color='white', width=1),
            z=dict(show=True, color='white', width=1)
        )
    ))

    fig.add_trace(go.Scatter3d(
        x=[selected_ct],
        y=[quarantine_days],
        z=[user_eff],
        mode='markers+text',
        marker=dict(size=12, color=marker_color, symbol='circle',
                    line=dict(color='white', width=3)),
        text=[f'{user_eff*100:.0f}%'],
        textposition='top center',
        textfont=dict(size=14, color='white', family='Arial Black'),
        name='當前情境',
        showlegend=True
    ))

    fig.update_layout(
        title={
            'text': title,
            'x': 0.5,
            'xanchor': 'center',
            'font': {'size': 16}
        },
        scene=dict(
            xaxis=dict(
                title='X-病毒量(Ct值)',
                range=[10, 25],
                tickvals=[10, 15, 18, 20, 25],
                showgrid=True,
                gridwidth=2,
                gridcolor='rgb(200, 200, 200)',
                showbackground=True,
                backgroundcolor='rgba(240, 240, 240, 0.9)'
            ),
            yaxis=dict(
                title='Y-隔離檢疫天數',
                range=[1, 21],
                tickvals=[1, 5, 7, 10, 14, 21],
                showgrid=True,
                gridwidth=2,
                gridcolor='rgb(200, 200, 200)',
                showbackground=True,
                backgroundcolor='rgba(240, 240, 240, 0.9)'
            ),
            zaxis=dict(
                title='Z-防疫效益',
                range=[0, 1],
                tickvals=[0, 0.25, 0.5, 0.75, 1.0],
                ticktext=['0%', '25%', '50%', '75%', '100%'],
                showgrid=True,
                gridwidth=2,
                gridcolor='rgb(200, 200, 200)',
                showbackground=True,
                backgroundcolor='rgba(240, 240, 240, 0.9)'
            ),
            camera=dict(
                eye=dict(x=1.5, y=-1.5, z=1.3),
                center=dict(x=0, y=0, z=0)
            ),
            aspectmode='manual',
            aspectratio=dict(x=1, y=1.2, z=0.8)
        ),
        height=600,
        showlegend=True,
        legend=dict(
            x=0.02,
            y=0.98,
            bgcolor='rgba(255, 255, 255, 0.9)',
            bordercolor='black',
            borderwidth=1
        ),
        margin=dict(l=0, r=0, t=40, b=0)
    )
    
    return fig

def render():
    """渲染 Tab2 的完整內容"""
    init_tab2_session_state()
    
    # 側邊欄 - 控制面板
    with st.sidebar:
        st.header("🏥 隔離檢疫決策工具")
        
        st.markdown("---")
        
        # 有加強劑的參數設定
        st.subheader("💉 有加強劑情境")
        
        booster_ct = st.slider(
            "🦠 接觸者Ct值",
            min_value=10.0,
            max_value=25.0,
            value=st.session_state.tab2_booster_ct,
            step=0.5,
            help="Ct值越低代表病毒量越高 (Ct>25視為無傳染性)",
            key="tab2_booster_ct_slider"
        )
        st.session_state.tab2_booster_ct = booster_ct
        
        booster_days = st.slider(
            "📅 隔離天數",
            min_value=1,
            max_value=21,
            value=st.session_state.tab2_booster_days,
            help="需要隔離檢疫的天數",
            key="tab2_booster_days_slider"
        )
        st.session_state.tab2_booster_days = booster_days
        
        booster_eff = get_quarantine_effectiveness(booster_ct, booster_days, True)
        st.metric(
            label="📊 防疫效益",
            value=f"{booster_eff * 100:.0f}%",
            delta="有加強劑"
        )
        
        st.markdown("---")
        
        # 沒有加強劑的參數設定
        st.subheader("⚠️ 無加強劑情境")
        
        no_booster_ct = st.slider(
            "🦠 接觸者Ct值",
            min_value=10.0,
            max_value=25.0,
            value=st.session_state.tab2_no_booster_ct,
            step=0.5,
            help="Ct值越低代表病毒量越高 (Ct>25視為無傳染性)",
            key="tab2_no_booster_ct_slider"
        )
        st.session_state.tab2_no_booster_ct = no_booster_ct
        
        no_booster_days = st.slider(
            "📅 隔離天數",
            min_value=1,
            max_value=21,
            value=st.session_state.tab2_no_booster_days,
            help="需要隔離檢疫的天數",
            key="tab2_no_booster_days_slider"
        )
        st.session_state.tab2_no_booster_days = no_booster_days
        
        no_booster_eff = get_quarantine_effectiveness(no_booster_ct, no_booster_days, False)
        st.metric(
            label="📊 防疫效益",
            value=f"{no_booster_eff * 100:.0f}%",
            delta="無加強劑"
        )
        
        st.markdown("---")
        
        # 情境說明
        with st.expander("🎯 使用情境", expanded=False):
            st.markdown("""
            當找到疑似接觸者時,防疫人員需要決定:
            **「要隔離/檢疫多少天?」**
            
            考量因素:
            - 接觸者的病毒量 (Ct值)
            - 是否已接種加強劑
            """)
        
        with st.expander("✅ Omicron 特性"):
            st.markdown("""
            相較於 Alpha 變異株:
            - **傳播更快** 但症狀較輕
            - **疫苗保護** 顯著縮短隔離時間
            - **Ct > 25** 視為無傳染性
            """)
        
        with st.expander("💉 疫苗影響"):
            st.markdown("""
            **加強劑的效益:**
            - 大幅縮短所需隔離天數
            - 相同天數下效益更高
            - 例: 達到90%效益
              - 有加強劑: 5天
              - 無加強劑: 11天
            """)

    # 主要內容區
    st.markdown("### 📊 隔離檢疫效益 3D 視覺化")
    
    # 生成 3D 數據
    ct_range = np.arange(10, 25.5, 0.5)
    day_range = np.arange(1, 22, 1)
    ct_mesh, day_mesh = np.meshgrid(ct_range, day_range)

    # 計算兩組效益值
    effectiveness_with_booster = np.zeros_like(ct_mesh)
    effectiveness_without_booster = np.zeros_like(ct_mesh)
    
    for i in range(len(day_range)):
        for j in range(len(ct_range)):
            effectiveness_with_booster[i, j] = get_quarantine_effectiveness(ct_mesh[i, j], day_mesh[i, j], True)
            effectiveness_without_booster[i, j] = get_quarantine_effectiveness(ct_mesh[i, j], day_mesh[i, j], False)

    # 創建兩個並排的圖表
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("#### 💉 有加強劑情境")
        fig1 = create_3d_plot(
            ct_mesh, day_mesh, effectiveness_with_booster,
            booster_ct, booster_days, 
            booster_eff,
            True,
            "隔離檢疫效益 - 有加強劑 (Omicron 變異株)",
            [
                [0.0, 'rgb(239, 68, 68)'],
                [0.33, 'rgb(245, 158, 11)'],
                [0.67, 'rgb(16, 185, 129)'],
                [1.0, 'rgb(59, 130, 246)']
            ],
            'blue'
        )
        st.plotly_chart(fig1, use_container_width=True)
        
        ct_category_booster = get_ct_category(booster_ct)
        st.info(f"**情境:** Ct值 = **{booster_ct}** ({ct_category_booster}), 隔離 **{booster_days}** 天 | 💉 已接種加強劑")
        
        if booster_eff >= 0.9:
            st.success("💡 **意義:** 隔離時間充足,可有效防止疫情傳播")
        elif booster_eff >= 0.7:
            st.warning("💡 **意義:** 隔離效果良好,但建議視情況延長")
        elif booster_eff >= 0.5:
            st.warning("💡 **意義:** 隔離效果一般,建議延長隔離時間")
        else:
            st.error("💡 **意義:** 隔離效果不足,需要大幅延長隔離時間")
    
    with col2:
        st.markdown("#### ⚠️ 無加強劑情境")
        fig2 = create_3d_plot(
            ct_mesh, day_mesh, effectiveness_without_booster,
            no_booster_ct, no_booster_days,
            no_booster_eff,
            False,
            "隔離檢疫效益 - 無加強劑 (Omicron 變異株)",
            [
                [0.0, 'rgb(239, 68, 68)'],
                [0.33, 'rgb(245, 158, 11)'],
                [0.67, 'rgb(16, 185, 129)'],
                [1.0, 'rgb(59, 130, 246)']
            ],
            'red'
        )
        st.plotly_chart(fig2, use_container_width=True)
        
        ct_category_no_booster = get_ct_category(no_booster_ct)
        st.info(f"**情境:** Ct值 = **{no_booster_ct}** ({ct_category_no_booster}), 隔離 **{no_booster_days}** 天 | ⚠️ 未接種加強劑")
        
        if no_booster_eff >= 0.9:
            st.success("💡 **意義:** 隔離時間充足,可有效防止疫情傳播")
        elif no_booster_eff >= 0.7:
            st.warning("💡 **意義:** 隔離效果良好,但建議視情況延長")
        elif no_booster_eff >= 0.5:
            st.warning("💡 **意義:** 隔離效果一般,建議延長隔離時間")
        else:
            st.error("💡 **意義:** 隔離效果不足,需要大幅延長隔離時間")
        
        # 疫苗建議
        improvement = (booster_eff - no_booster_eff) * 100
        if improvement > 0:
            st.info(f"💉 **提示:** 若接種加強劑,在相同參數下可提升 {improvement:.0f}% 效益")

    # 底部說明區域
    st.markdown("---")

    col_a, col_b = st.columns(2)

    with col_a:
        with st.expander("💡 操作說明", expanded=False):
            st.markdown("""
            **3D 圖表說明:**
            - 🔵 **左側圖表**: 有加強劑的隔離效益 (藍綠色)
            - 🔴 **右側圖表**: 無加強劑的隔離效益 (橙紅色)
            - 💎 **菱形標記**: 您當前選擇的情境
            - 兩圖對比可清楚看出**疫苗的效益差異**
            
            **互動操作:**
            - 🖱️ 拖曳旋轉視角
            - 🔍 滾輪縮放
            - 🎚️ 使用左側滑桿調整參數
            """)

    with col_b:
        with st.expander("📊 查看詳細數據", expanded=False):
            # 創建對比表格
            test_days_list = [3, 5, 7, 10, 14]
            
            data = {
                '隔離天數': test_days_list,
                '有加強劑 (Ct=' + str(booster_ct) + ')': [f"{get_quarantine_effectiveness(booster_ct, d, True)*100:.0f}%" for d in test_days_list],
                '無加強劑 (Ct=' + str(no_booster_ct) + ')': [f"{get_quarantine_effectiveness(no_booster_ct, d, False)*100:.0f}%" for d in test_days_list],
                '效益差異': [f"+{(get_quarantine_effectiveness(booster_ct, d, True) - get_quarantine_effectiveness(no_booster_ct, d, False))*100:.0f}%" for d in test_days_list]
            }
            
            df = pd.DataFrame(data)
            st.dataframe(df, use_container_width=True)