"""
Simon Two-stage Design 互動式教學工具
賽門式二階段試驗設計 - 基於 Tan et al. (2002) BJC 論文
完整版：含 Trial C、Trial P、4種事前分佈
修改版：使用 Google Gemini API
"""

import gradio as gr
import numpy as np
from scipy import stats
import matplotlib.pyplot as plt

# ============== 計算函數 ==============

def beta_mean(a, b):
    """Beta 分佈平均值"""
    return a / (a + b)

def beta_var(a, b):
    """Beta 分佈變異數"""
    return (a * b) / ((a + b)**2 * (a + b + 1))

def beta_ci(a, b, level=0.95):
    """Beta 分佈可信區間"""
    lower = (1 - level) / 2
    return stats.beta.ppf(lower, a, b), stats.beta.ppf(1-lower, a, b)

def calc_clinical_prior(r0, r1):
    """
    Clinical Prior - 使用論文的精確數值
    Trial C: Beta(0.7, 2.1)
    Trial P: Beta(0.6, 3.0)
    """
    if r0 == 0.10 and r1 == 0.30:
        return 0.7, 2.1  # Trial C (論文 Table 1)
    elif r0 == 0.05 and r1 == 0.20:
        return 0.6, 3.0  # Trial P (論文 Table 2)
    else:
        # 其他情況用計算
        m = (r0 + r1) / 2
        return m * 2, (1 - m) * 2

def calc_sceptical_prior(r0, r1):
    """
    Sceptical Prior - 使用論文的精確數值
    Trial C: Beta(1, 9)
    Trial P: Beta(0.4, 7.6)
    """
    if r0 == 0.10 and r1 == 0.30:
        return 1, 9  # Trial C (論文 Table 1)
    elif r0 == 0.05 and r1 == 0.20:
        return 0.4, 7.6  # Trial P (論文 Table 2)
    else:
        return r0 * 10, (1 - r0) * 10

def calc_enthusiastic_prior(r0, r1):
    """
    Enthusiastic Prior - 使用論文的精確數值
    Trial C: Beta(3, 7)
    Trial P: Beta(2.4, 9.6)
    """
    if r0 == 0.10 and r1 == 0.30:
        return 3, 7  # Trial C (論文 Table 1)
    elif r0 == 0.05 and r1 == 0.20:
        return 2.4, 9.6  # Trial P (論文 Table 2)
    else:
        return r1 * 10, (1 - r1) * 10

# ============== 主分析函數 ==============

def analyze(trial_type, 
            s1_success, s1_total, s2_success, s2_total,
            prior_type, custom_a, custom_b):
    
    # 根據試驗類型設定參數
    if trial_type == "Trial C (初治組)":
        r0, r1 = 0.10, 0.30
        trial_name = "Trial C：Chemotherapy-naïve（初治組）"
        trial_name_en = "Trial C: Chemotherapy-naive"
    else:
        r0, r1 = 0.05, 0.20
        trial_name = "Trial P：Previously treated（曾治療組）"
        trial_name_en = "Trial P: Previously treated"
    
    # 基本計算
    s1_success, s1_total = int(s1_success), int(s1_total)
    s2_success, s2_total = int(s2_success), int(s2_total)
    total_success = s1_success + s2_success
    total_n = s1_total + s2_total
    
    # 根據事前分佈類型計算參數
    if prior_type == "Clinical Prior（臨床事前分佈）":
        prior_a, prior_b = calc_clinical_prior(r0, r1)
        prior_desc = f"根據專家意見，P(θ<R₀) = P(R₀<θ<R₁) = P(θ>R₁) = 1/3"
        prior_name_en = "Clinical Prior"
    elif prior_type == "Reference Prior（參考事前分佈）":
        prior_a, prior_b = 1, 1  # Uniform
        prior_desc = "無資訊事前分佈（Uniform），代表完全不確定"
        prior_name_en = "Reference Prior"
    elif prior_type == "Sceptical Prior（懷疑性事前分佈）":
        prior_a, prior_b = calc_sceptical_prior(r0, r1)
        prior_desc = f"平均值 = R₀ = {r0*100:.0f}%，P(θ>R₁) = 5%（偏向不相信藥效）"
        prior_name_en = "Sceptical Prior"
    elif prior_type == "Enthusiastic Prior（樂觀性事前分佈）":
        prior_a, prior_b = calc_enthusiastic_prior(r0, r1)
        prior_desc = f"平均值 = R₁ = {r1*100:.0f}%，P(θ<R₀) = 5%（偏向相信藥效）"
        prior_name_en = "Enthusiastic Prior"
    else:  # 自訂
        prior_a, prior_b = custom_a, custom_b
        prior_desc = "使用者自訂參數"
        prior_name_en = "Custom Prior"
    
    # === 計算事後分佈 ===
    # Stage 1 後
    post1_a = prior_a + s1_success
    post1_b = prior_b + (s1_total - s1_success)
    
    # Stage 2 後（最終）
    post2_a = prior_a + total_success
    post2_b = prior_b + (total_n - total_success)
    
    # === 計算機率 ===
    # 事前
    prior_prob_lt_r0 = stats.beta.cdf(r0, prior_a, prior_b)
    prior_prob_r0_r1 = stats.beta.cdf(r1, prior_a, prior_b) - prior_prob_lt_r0
    prior_prob_gt_r1 = 1 - stats.beta.cdf(r1, prior_a, prior_b)
    
    # Stage 1 後
    post1_prob_lt_r0 = stats.beta.cdf(r0, post1_a, post1_b)
    post1_prob_r0_r1 = stats.beta.cdf(r1, post1_a, post1_b) - post1_prob_lt_r0
    post1_prob_gt_r1 = 1 - stats.beta.cdf(r1, post1_a, post1_b)
    
    # Stage 2 後
    post2_prob_lt_r0 = stats.beta.cdf(r0, post2_a, post2_b)
    post2_prob_r0_r1 = stats.beta.cdf(r1, post2_a, post2_b) - post2_prob_lt_r0
    post2_prob_gt_r1 = 1 - stats.beta.cdf(r1, post2_a, post2_b)
    
    # === 頻率學派判定 ===
    if trial_type == "Trial C (初治組)":
        stage1_threshold = 2  # 至少 2 人有反應
        final_threshold = 6   # 至少 6 人有反應
    else:
        stage1_threshold = 1  # 至少 1 人有反應
        final_threshold = 4   # 至少 4 人有反應
    
    stage1_pass = s1_success >= stage1_threshold
    final_pass = total_success >= final_threshold
    
    # === 輸出結果 ===
    result = f"""
## 📋 試驗資訊：{trial_name}

### 設計參數
- **R₀（無興趣反應率）**: {r0*100:.0f}%
- **R₁（理想反應率）**: {r1*100:.0f}%

---

## 📊 頻率學派分析

### 試驗結果
| 階段 | 成功 (R) | 失敗 (N-R) | 總數 (N) | 反應率 |
|------|----------|------------|----------|--------|
| Stage I | {s1_success} | {s1_total - s1_success} | {s1_total} | {s1_success/s1_total*100:.1f}% |
| Stage II | {s2_success} | {s2_total - s2_success} | {s2_total} | {s2_success/s2_total*100:.1f}% |
| **合計** | **{total_success}** | **{total_n - total_success}** | **{total_n}** | **{total_success/total_n*100:.1f}%** |

### Simon Minimax Design 判定
- Stage I：需 ≥ {stage1_threshold} 人有反應 → {'✅ 通過' if stage1_pass else '❌ 未通過'}（實際 {s1_success} 人）
- 最終：需 ≥ {final_threshold} 人有反應 → {'✅ 有效' if final_pass else '❌ 無效'}（實際 {total_success} 人）

---

## 🔮 貝氏分析

### 事前分佈：{prior_type}
{prior_desc}

**Beta({prior_a:.1f}, {prior_b:.1f})**，平均值 = {beta_mean(prior_a, prior_b)*100:.1f}%

### 機率表（對應論文 Table 1 & 2）

| 階段 | Beta 參數 | π（平均值） | θ ≤ R₀ | R₀ < θ ≤ R₁ | θ > R₁ |
|------|-----------|-------------|--------|-------------|--------|
| Prior | Beta({prior_a:.1f}, {prior_b:.1f}) | **{beta_mean(prior_a, prior_b):.3f}** | {prior_prob_lt_r0:.3f} | {prior_prob_r0_r1:.3f} | {prior_prob_gt_r1:.3f} |
| Stage 1 posterior | Beta({post1_a:.1f}, {post1_b:.1f}) | **{beta_mean(post1_a, post1_b):.3f}** | {post1_prob_lt_r0:.3f} | {post1_prob_r0_r1:.3f} | {post1_prob_gt_r1:.3f} |
| Stage 2 posterior | Beta({post2_a:.1f}, {post2_b:.1f}) | **{beta_mean(post2_a, post2_b):.3f}** | {post2_prob_lt_r0:.3f} | {post2_prob_r0_r1:.3f} | {post2_prob_gt_r1:.3f} |

### 貝氏更新過程（練習用）

| 事前分佈 | Stage I 結果 | Stage I 後分佈 | Stage II 結果 | Stage II 後分佈 |
|----------|-------------|---------------|--------------|----------------|
| Beta({prior_a:.1f},{prior_b:.1f}), π={beta_mean(prior_a, prior_b):.2f} | 成功:{s1_success}, 失敗:{s1_total - s1_success} | Beta({post1_a:.1f},{post1_b:.1f}), π={beta_mean(post1_a, post1_b):.2f} | 成功:{s2_success}, 失敗:{s2_total - s2_success} | Beta({post2_a:.1f},{post2_b:.1f}), π={beta_mean(post2_a, post2_b):.3f} |

### 結果解讀
- **P(θ > R₁) = {post2_prob_gt_r1:.1%}**：藥物反應率超過理想值的機率
- **P(θ > R₀) = {(1-post2_prob_lt_r0):.1%}**：藥物值得繼續研究的機率

"""
    
    if post2_prob_gt_r1 > 0.5:
        result += "📌 **結論**：強烈支持進入 Phase III 試驗\n"
    elif post2_prob_gt_r1 > 0.25:
        result += "📌 **結論**：中度支持，可考慮進一步研究\n"
    else:
        result += "📌 **結論**：證據不足，需謹慎評估\n"
    
    # === 繪圖（英文標籤）===
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    x = np.linspace(0, 1, 500)
    
    # 左圖：三條分佈曲線
    ax1 = axes[0]
    ax1.plot(x, stats.beta.pdf(x, prior_a, prior_b), 'b--', lw=2, 
             label=f'Prior Beta({prior_a:.1f},{prior_b:.1f})')
    ax1.plot(x, stats.beta.pdf(x, post1_a, post1_b), 'g-.', lw=2,
             label=f'Stage 1 Post Beta({post1_a:.1f},{post1_b:.1f})')
    ax1.plot(x, stats.beta.pdf(x, post2_a, post2_b), 'r-', lw=2.5,
             label=f'Stage 2 Post Beta({post2_a:.1f},{post2_b:.1f})')
    ax1.axvline(r0, color='gray', ls=':', alpha=0.7, label=f'R0={r0}')
    ax1.axvline(r1, color='orange', ls=':', alpha=0.7, label=f'R1={r1}')
    ax1.set_xlabel('Response Rate (theta)', fontsize=12)
    ax1.set_ylabel('Density', fontsize=12)
    ax1.set_title(f'{prior_name_en} - Bayesian Update\n{trial_name_en}', fontsize=12)
    ax1.legend(fontsize=9)
    ax1.grid(alpha=0.3)
    ax1.set_xlim(0, 0.8)
    
    # 右圖：4 種事前分佈比較
    ax2 = axes[1]
    
    # 計算所有 4 種事前分佈
    clin_a, clin_b = calc_clinical_prior(r0, r1)
    ref_a, ref_b = 1, 1
    scep_a, scep_b = calc_sceptical_prior(r0, r1)
    enth_a, enth_b = calc_enthusiastic_prior(r0, r1)
    
    ax2.plot(x, stats.beta.pdf(x, clin_a, clin_b), 'b-', lw=1.5, alpha=0.7,
             label=f'Clinical ({clin_a:.1f},{clin_b:.1f})')
    ax2.plot(x, stats.beta.pdf(x, ref_a, ref_b), 'gray', ls='--', lw=1.5, alpha=0.7,
             label=f'Reference ({ref_a},{ref_b})')
    ax2.plot(x, stats.beta.pdf(x, scep_a, scep_b), 'r-', lw=1.5, alpha=0.7,
             label=f'Sceptical ({scep_a:.1f},{scep_b:.1f})')
    ax2.plot(x, stats.beta.pdf(x, enth_a, enth_b), 'g-', lw=1.5, alpha=0.7,
             label=f'Enthusiastic ({enth_a:.1f},{enth_b:.1f})')
    ax2.axvline(r0, color='gray', ls=':', alpha=0.5)
    ax2.axvline(r1, color='orange', ls=':', alpha=0.5)
    ax2.set_xlabel('Response Rate (theta)', fontsize=12)
    ax2.set_ylabel('Density', fontsize=12)
    ax2.set_title(f'Comparison of 4 Prior Types\n{trial_name_en} (Figure 1 & 4 style)', fontsize=12)
    ax2.legend(fontsize=9)
    ax2.grid(alpha=0.3)
    ax2.set_xlim(0, 0.8)
    
    plt.tight_layout()
    
    # 建立摘要給 LLM
    summary = f"""
Trial: {trial_name}
R0={r0}, R1={r1}
Stage1: {s1_success}/{s1_total}, Stage2: {s2_success}/{s2_total}, Total: {total_success}/{total_n}
Prior: {prior_type}, Beta({prior_a:.1f},{prior_b:.1f})
Posterior: Beta({post2_a:.1f},{post2_b:.1f})
P(theta>R0)={1-post2_prob_lt_r0:.3f}, P(theta>R1)={post2_prob_gt_r1:.3f}
Frequentist: Stage1={'Pass' if stage1_pass else 'Fail'}, Final={'Effective' if final_pass else 'Ineffective'}
"""
    
    return result, fig, summary

# ============== LLM 解釋功能（OpenAI ChatGPT）==============

def get_llm_explanation(summary_str, question, api_key):
    """呼叫 OpenAI API (ChatGPT) 獲得 LLM 解釋"""
    
    if not api_key or api_key.strip() == "":
        return "⚠️ 請先輸入 OpenAI API Key\n\n（到 https://platform.openai.com/api-keys 取得）"
    
    if not summary_str:
        return "⚠️ 請先執行分析，再請 AI 解釋"
    
    try:
        from openai import OpenAI
        
        client = OpenAI(api_key=api_key)
        
        system_prompt = """你是一位專業的生物統計學家和臨床試驗專家，專門教導學生理解 Simon Two-stage Design 和貝氏分析。

你的教學基於 Tan et al. (2002) 發表在 British Journal of Cancer 的論文：
"A Bayesian re-assessment of two Phase II trials of gemcitabine in metastatic nasopharyngeal cancer"

請用繁體中文回答，語氣友善且易懂，像是老師在跟學生解釋。回答時請：
1. 用簡單的語言解釋統計結果的意義
2. 說明頻率學派和貝氏方法的差異
3. 解釋不同事前分佈（Clinical, Reference, Sceptical, Enthusiastic）的意義
4. 給出是否建議進入下一期試驗的建議"""

        user_prompt = f"""以下是一個 Simon Two-stage Design 臨床試驗的分析結果：

{summary_str}

學生的問題：{question if question else "請幫我解釋這個結果代表什麼意思？不同的事前分佈會如何影響結論？"}

請用簡單易懂的方式回答。"""

        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            max_tokens=1500,
            temperature=0.7
        )
        
        return response.choices[0].message.content
        
    except Exception as e:
        return f"❌ API 呼叫失敗：{str(e)}\n\n請確認：\n1. API Key 是否正確\n2. 帳戶是否有餘額"

# ============== Gradio 介面 ==============

with gr.Blocks(
    title="Simon Two-stage Design 教學工具",
    theme=gr.themes.Soft(primary_hue="blue")
) as demo:
    
    # 儲存分析結果
    analysis_summary = gr.State("")
    
    # 標題
    gr.Markdown("""
    # 🔬 Simon Two-stage Design 互動式教學工具
    ## 基於 Tan et al. (2002) British Journal of Cancer 論文
    **A Bayesian re-assessment of two Phase II trials of gemcitabine in metastatic nasopharyngeal cancer**
    """)
    
    # 故事背景
    with gr.Accordion("📖 教學故事線", open=True):
        gr.Markdown("""
        > 🎯 **核心問題**：新藥開發很燒錢，怎麼「早點知道該不該繼續」？
        
        - 💰 一個新藥開發平均要 **10-20 億美元**
        - ⏰ 耗時 **10-15 年**
        - 📉 Phase II 失敗率高達 **70%**
        
        ↓
        
        「如果藥沒效，能不能早點停？」
        「如果有希望，再多收病人確認」
        
        ↓
        
        > **1989 年 Richard Simon 提出解決方案：**
        > 👉 **Simon Two-stage Design**
        
        ---
        
        ### 🏥 今天的案例：Gemcitabine 治療鼻咽癌
        
        新加坡國家癌症中心要測試 **gemcitabine（健擇）** 對轉移性鼻咽癌的療效，設計了兩個試驗：
        
        | 試驗 | 對象 | R₀（沒興趣） | R₁（理想） |
        |------|------|-------------|-----------|
        | **Trial C** | 初治患者 | 10% | 30% |
        | **Trial P** | 曾治療患者 | 5% | 20% |
        
        - 若反應率 ≤ R₀ → 藥物不值得繼續研究
        - 若反應率 ≥ R₁ → 值得進入 Phase III
        
        👇 **請選擇試驗、輸入數據，看看結果如何！**
        """)
    
    with gr.Accordion("📐 方法說明：四種事前分佈", open=False):
        gr.Markdown("""
        貝氏分析需要設定「事前分佈」，代表看到數據**之前**的信念：
        
        | 類型 | 定義 | 立場 |
        |------|------|------|
        | **Clinical** | P(θ<R₀) = P(R₀<θ<R₁) = P(θ>R₁) = 1/3 | 均衡看法 |
        | **Reference** | Uniform (1,1) | 完全不確定 |
        | **Sceptical** | 平均=R₀，P(θ>R₁)=5% | 懷疑，不太信 |
        | **Enthusiastic** | 平均=R₁，P(θ<R₀)=5% | 樂觀，相信藥效 |
        
        **Beta-Binomial 共軛更新**：
        ```
        事前：Beta(α, β) + 數據：n人中r人成功 → 事後：Beta(α+r, β+n-r)
        ```
        
        📚 *Tan SB et al. Br J Cancer. 2002;86(6):843-850.*
        """)
    
    gr.Markdown("---")
    
    # 主要設定區
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ 試驗設定")
            
            trial_type = gr.Radio(
                choices=["Trial C (初治組)", "Trial P (曾治療組)"],
                value="Trial C (初治組)",
                label="選擇試驗類型"
            )
            
            gr.Markdown("### 📝 輸入試驗數據")
            gr.Markdown("*請根據論文填入數據，或點下方按鈕載入*")
            
            with gr.Row():
                s1_success = gr.Number(value=None, label="Stage I 成功人數", precision=0)
                s1_total = gr.Number(value=None, label="Stage I 總人數", precision=0)
            
            with gr.Row():
                s2_success = gr.Number(value=None, label="Stage II 成功人數", precision=0)
                s2_total = gr.Number(value=None, label="Stage II 總人數", precision=0)
            
            gr.Markdown("### 🔮 選擇事前分佈")
            
            prior_type = gr.Radio(
                choices=[
                    "Clinical Prior（臨床事前分佈）",
                    "Reference Prior（參考事前分佈）",
                    "Sceptical Prior（懷疑性事前分佈）",
                    "Enthusiastic Prior（樂觀性事前分佈）",
                    "Custom（自訂）"
                ],
                value="Clinical Prior（臨床事前分佈）",
                label="事前分佈類型"
            )
            
            with gr.Row(visible=False) as custom_row:
                custom_a = gr.Number(value=1, label="自訂 α", precision=1)
                custom_b = gr.Number(value=1, label="自訂 β", precision=1)
            
            # 載入預設值按鈕
            gr.Markdown("### 📚 載入論文數據（或自己填完後對答案）")
            with gr.Row():
                load_c_btn = gr.Button("📥 載入 Trial C 數據", size="sm")
                load_p_btn = gr.Button("📥 載入 Trial P 數據", size="sm")
            
            btn = gr.Button("🚀 執行分析", variant="primary", size="lg")
        
        with gr.Column(scale=2):
            gr.Markdown("### 📊 分析結果")
            output_text = gr.Markdown()
            output_plot = gr.Plot()
    
    # 載入預設值函數
    def load_trial_c():
        return "Trial C (初治組)", 3, 15, 4, 10
    
    def load_trial_p():
        return "Trial P (曾治療組)", 7, 13, 6, 14
    
    load_c_btn.click(load_trial_c, outputs=[trial_type, s1_success, s1_total, s2_success, s2_total])
    load_p_btn.click(load_trial_p, outputs=[trial_type, s1_success, s1_total, s2_success, s2_total])
    
    # 顯示/隱藏自訂參數
    def toggle_custom(prior):
        return gr.Row(visible=(prior == "Custom（自訂）"))
    
    prior_type.change(toggle_custom, prior_type, custom_row)
    
    # 執行分析
    btn.click(
        fn=analyze,
        inputs=[trial_type, s1_success, s1_total, s2_success, s2_total, 
                prior_type, custom_a, custom_b],
        outputs=[output_text, output_plot, analysis_summary]
    )
    
    # ========== LLM 解釋區 ==========
    gr.Markdown("---")
    gr.Markdown("### 🤖 AI 助教解釋")
    
    with gr.Row():
        with gr.Column(scale=2):
            user_question = gr.Textbox(
                label="你的問題（選填）",
                placeholder="例如：為什麼 Sceptical Prior 和 Enthusiastic Prior 的結論會不同？",
                lines=2
            )
        with gr.Column(scale=1):
            api_key = gr.Textbox(
                label="OpenAI API Key",
                placeholder="sk-...",
                type="password"
            )
    
    ask_ai_btn = gr.Button("🧠 請 AI 助教解釋", variant="secondary")
    ai_response = gr.Markdown()
    
    ask_ai_btn.click(
        fn=get_llm_explanation,
        inputs=[analysis_summary, user_question, api_key],
        outputs=[ai_response]
    )
    
    # 習題
    gr.Markdown("---")
    with gr.Accordion("✏️ 論文導讀問題", open=False):
        gr.Markdown("""
        ### 問題 1：頻率學派 vs 貝氏方法
        論文中提到「The conclusions drawn using the Bayesian approach were in general agreement with those obtained from the frequentist analysis」，但貝氏方法有什麼額外的優點？
        
        ### 問題 2：事前分佈的影響
        比較 Trial C 使用 Sceptical Prior 和 Enthusiastic Prior 的結果（論文 Table 1），為什麼 P(θ>R₁) 的值差異很大？
        
        ### 問題 3：Trial C vs Trial P
        Trial P 的結果（48% 反應率）比 Trial C（28%）好很多，論文作者如何解釋這個現象？
        
        ### 問題 4：實務應用
        如果你是藥廠的統計師，你會選擇哪種事前分佈來分析結果？為什麼？
        """)
    
    # 頁尾
    gr.Markdown("""
    ---
    <center>
    🔬 Simon Two-stage Design 教學工具<br>
    Based on Tan et al. (2002) British Journal of Cancer<br>
    <small>Powered by Gradio + OpenAI</small>
    </center>
    """)

if __name__ == "__main__":
    demo.launch()