import gradio as gr
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime
import os
from transformers import pipeline

# CSV 檔案位置
DATA_FILE = "/tmp/stress_data.csv"

# 初始化 transformers 分析 pipeline
# 這裡使用假設的中文情緒分析模型
# 如果你有心理壓力專用模型，替換 model="你的模型名稱"
classifier = pipeline(
    "text-classification",
    model="uer/roberta-base-finetuned-chinanews-chinese",  # 中文分類範例
    device=-1  # CPU；若有 GPU，可設成 0
)

# 分析函數
def analyze(text):
    # 用模型預測
    pred = classifier(text)[0]  # 回傳 dict: {'label': 'POSITIVE', 'score': 0.95}
    
    label = pred['label']
    score = pred['score']
    
    # 假設我們把正面 -> 壓力低，負面 -> 壓力高
    if label in ["POSITIVE", "positive"]:
        stress_score = int((1 - score) * 100)  # 正面越高，壓力越低
    else:
        stress_score = int(score * 100)        # 負面越高，壓力越高
    
    result_text = f"文字分析結果：情緒標籤 {label}，壓力指數 {stress_score}"
    
    # 儲存資料
    save_data(text, stress_score)
    
    # 生成趨勢圖
    fig = plot_trend()
    
    return result_text, fig

# 以下函數維持不變
def save_data(text, score):
    now = datetime.now()
    df = pd.DataFrame([[now, text, score]], columns=["timestamp", "text", "stress"])
    if os.path.exists(DATA_FILE):
        df.to_csv(DATA_FILE, mode="a", header=False, index=False)
    else:
        df.to_csv(DATA_FILE, index=False)

def plot_trend():
    fig, ax = plt.subplots(figsize=(6,3))
    if not os.path.exists(DATA_FILE):
        ax.text(0.5, 0.5, "暫無資料", ha='center', va='center')
        ax.set_axis_off()
        return fig

    df = pd.read_csv(DATA_FILE)
    df['timestamp'] = pd.to_datetime(df['timestamp'])
    ax.plot(df['timestamp'], df['stress'], marker='o', linestyle='-')
    ax.set_title("心理壓力趨勢")
    ax.set_xlabel("時間")
    ax.set_ylabel("壓力指數")
    ax.grid(True)
    plt.xticks(rotation=30)
    return fig

def reset_data():
    if os.path.exists(DATA_FILE):
        os.remove(DATA_FILE)
    fig, ax = plt.subplots()
    ax.text(0.5, 0.5, "資料已重置", ha='center', va='center')
    ax.set_axis_off()
    return "資料已重置", fig

# Gradio 介面維持不變
with gr.Blocks() as demo:
    gr.Markdown("# 中文心理壓力分析器")
    gr.Markdown("""這是一款結合 BERT 深度學習模型與關鍵詞分析的心理壓力評估工具。
輸入中文文字，它就能分析文字隱含的心理壓力，並生成壓力分數與趨勢圖。""")
    
    with gr.Row():
        text_input = gr.Textbox(
            label="和我聊聊吧", 
            placeholder="說說今天的心情或煩惱...", 
            lines=5
        )
    with gr.Row():
        submit_btn = gr.Button("分析")
        reset_btn = gr.Button("重置資料")
    
    result_text = gr.Textbox(label="分析結果")
    trend_plot = gr.Plot(label="心理壓力趨勢")
    
    submit_btn.click(analyze, inputs=text_input, outputs=[result_text, trend_plot])
    reset_btn.click(reset_data, inputs=None, outputs=[result_text, trend_plot])

demo.launch()