app.py

import gradio as gr
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import joblib
import os
import warnings
from sklearn.pipeline import Pipeline

# Cấu hình thẩm mỹ
plt.style.use('seaborn-v0_8-muted')
sns.set_theme(style="whitegrid")
warnings.filterwarnings('ignore')

# --- 1. CSS & CONFIG ---
custom_css = """
#pell-gregory-grid .wrap {
    display: grid !important;
    grid-template-columns: repeat(3, 1fr) !important;
    gap: 15px !important;
}
"""

# Load mô hình
model_path = "dental_models_v2.1.pkl"

print(f"📂 Đang kiểm tra file tại: {os.getcwd()}/{model_path}")

try:
    if os.path.exists(model_path):
        models = joblib.load(model_path)
        print("✅ Load model thành công.")
    else:
        print("⚠️ Lỗi: Không tìm thấy file model.")
except Exception as e:
    print(f"❌ Lỗi khi load model: {e}")

# --- 2. HÀM TIỀN XỬ LÝ DỮ LIỆU ---

def preprocess_input_v2(tuoi, gioi_tinh, ben_pt, kinh_nghiem, pell_gregory, goc, chan_rang, ord_lq, ha_mieng, ma, so_thuoc=0, thoi_gian_pt=0):
    cd = pell_gregory[:-1]
    r7 = pell_gregory[-1]
    flex_group = 'Kém' if ma < 45 else 'Bình thường'
    opening_group = 'Há lớn' if ha_mieng >= 45 else ('Há miệng nhỏ' if ha_mieng >= 35 else 'Há miệng hạn chế')

    def get_angle_group(x):
        if x <= -10: return '<=-10'
        elif x <= 0: return '-10-0'
        elif x <= 45: return '0-45'
        elif x <= 90: return '45-90'
        else: return '>90'
    angle_group = get_angle_group(goc)

    r7_grouped = 'Class A' if r7 == 'A' else 'Class B/C'
    ramus_grouped = 'Class I/II' if cd in ['I', 'II'] else 'Class III'
    root_val = str(chan_rang)[0]
    root_grouped = 'Easy (2/3)' if root_val in ['2', '3'] else 'Hard (1/4)'
    age = 2025 - tuoi
    age_group = 'Young' if age <= 22 else ('Adult' if age <= 30 else 'Senior')

    raw_data = {
        'Giới tính': gioi_tinh, 'Bên PT': ben_pt, 'Liên quan ORD': ord_lq, 'Kinh nghiệm PTV': kinh_nghiem,
        'R7_Grouped': r7_grouped, 'Ramus_Grouped': ramus_grouped, 'Root_Grouped': root_grouped,
        'Age_Group': age_group, 'Flex_Group': flex_group, 'Opening_Group': opening_group,
        'Angle_Grouped': angle_group, 'Thời gian PT (phút)': thoi_gian_pt, 'Số viên thuốc GD': so_thuoc
    }
    return pd.DataFrame([raw_data])

# --- 3. CORE LOGIC ---

def transform_data_through_pipeline(pipeline, df):
    X = df.copy()
    for name, step in pipeline.steps[:-1]:
        X = step.transform(X)
    return X

def get_model_uncertainty(estimator, X_input):
    """
    Hàm trợ giúp để lấy độ lệch chuẩn (std) tùy theo loại model
    """
    # Trường hợp 1: Các model dạng Ensemble (Random Forest, ExtraTrees)
    if hasattr(estimator, 'estimators_'):
        all_tree_preds = [tree.predict(X_input)[0] for tree in estimator.estimators_]
        return np.std(all_tree_preds)
    
    # Trường hợp 2: Các model không có cây (Ridge, Lasso, SVR, KNN, Linear)
    return 0

def get_prediction_with_error(model_key, input_df):
    """Dự báo kèm sai số cho model đơn lẻ (Thời gian PT)"""
    pipe = models[model_key]
    X_final = transform_data_through_pipeline(pipe, input_df)
    
    # Lấy model lõi từ pipeline
    estimator = pipe.named_steps['model']
    
    pred = estimator.predict(X_final)[0]
    std_val = get_model_uncertainty(estimator, X_final)
    
    return pred, std_val

def get_increasing_uncertainty_v3(input_df, actuals=None):
    """Xử lý chuỗi hồi phục với sai số lũy tiến cho mọi loại model"""
    full_pipe = models['Recovery_Chain']
    X_base = transform_data_through_pipeline(full_pipe, input_df)
    
    chain_engine = full_pipe.named_steps['model']
    preds, errors = [], []
    Z = 1.0 
    X_current = X_base.copy()
    
    for i, estimator in enumerate(chain_engine.estimators_):
        # Sử dụng hàm helper mới để tránh lỗi AttributeError
        m = estimator.predict(X_current)[0]
        s = get_model_uncertainty(estimator, X_current)
        
        if actuals and i in actuals:
            m = actuals[i]
            s = 0.05
            
        if i < 3: m = np.clip(m, 0, 5)

        preds.append(m)
        if i == 0:
            errors.append(s * Z)
        else:
            # sqrt(sai số hiện tại^2 + sai số truyền lại^2)
            cum_err = np.sqrt(s**2 + (errors[i-1]/Z)**2) * Z
            errors.append(cum_err)
            
        X_current = np.hstack([X_current, np.array([[m]])])
        
    return preds, errors

# --- 4. TAB FUNCTIONS ---

def predict_tab1(tuoi, gioi, ben, exp, pg, goc, chan, ord_lq, ha, ma):
    df_raw = preprocess_input_v2(tuoi, gioi, ben, exp, pg, goc, chan, ord_lq, ha, ma)

    t_m, t_s = get_prediction_with_error('Thời gian PT (phút)', df_raw)
    t_err = t_s * 1.0

    med_pipe = models['Số viên thuốc GD']
    df_with_time = df_raw.assign(**{'Thời gian PT (phút)': t_m})
    med_m = med_pipe.predict(df_with_time)[0]

    df_final = df_with_time.assign(**{'Số viên thuốc GD': med_m})
    preds, errs = get_increasing_uncertainty_v3(df_final)

    # Cập nhật logic hiển thị thời gian: Ẩn dải biến thiên nếu t_err = 0
    if t_err > 0:
        res_text = (
            f"⏱️ THỜI GIAN NHỔ RĂNG DỰ KIẾN: {t_m:.1f} ± {t_err:.1f} phút\n"
            f"(Dải biến thiên: {max(0, t_m-t_err):.1f} - {t_m+t_err:.1f} phút)\n\n"
        )
    else:
        res_text = f"⏱️ THỜI GIAN NHỔ RĂNG DỰ KIẾN: {t_m:.1f} phút\n\n"

    res_text += f"💊 SỐ VIÊN THUỐC GIẢM ĐAU DỰ KIẾN: {max(0, round(med_m))} viên\n\n"
    res_text += "📈 DỰ ĐOÁN MỨC ĐỘ ĐAU:\n"
    
    # Cập nhật logic hiển thị mức độ đau từng ngày
    res_text += f"• Ngày 1: {preds[0]:.1f}" + (f" (±{errs[0]:.1f})\n" if errs[0] > 0 else "\n")
    res_text += f"• Ngày 3: {preds[1]:.1f}" + (f" (±{errs[1]:.1f})\n" if errs[1] > 0 else "\n")
    res_text += f"• Ngày 7: {preds[2]:.1f}" + (f" (±{errs[2]:.1f})\n" if errs[2] > 0 else "\n")
    res_text += f"• Hết đau hoàn toàn: Ngày thứ {preds[3]:.1f}" + (f" (±{errs[3]:.1f})" if errs[3] > 0 else "")

    fig, ax = plt.subplots(figsize=(8, 4.5))
    days = [1, 3, 7]
    ax.plot(days, preds[:3], 'o-', color='#4A90E2', linewidth=3, markersize=8, label='Mức độ đau trung bình')
    if errs[0] > 0 or errs[1] > 0 or errs[2] > 0:
      ax.fill_between(days, [max(0, preds[i]-errs[i]) for i in range(3)], [min(5, preds[i]+errs[i]) for i in range(3)],
                      color='#4A90E2', alpha=0.15, label='Dải sai số dự kiến')
    for i, p in enumerate(preds[:3]):
        ax.text(days[i], p + 0.2, f"{p:.1f}", ha='center', fontweight='bold', color='#2c3e50')

    ax.set_ylim(-0.2, 5.5); ax.set_xticks(days); ax.set_ylabel("Mức độ đau (VAS)"); ax.legend()
    plt.tight_layout(); plt.close()
    return res_text, fig

def predict_tab3(tuoi, gioi, ben, exp, pg, goc, chan, ord_lq, ha, ma, thuoc_uong, t_thuc, thoi_diem, p1, p3, p7):
    p1v, p3v, p7v = float(p1), float(p3), float(p7)

    # Ràng buộc logic
    if thoi_diem == "Tái khám Ngày 3":
        if p3v > p1v: return "⚠️ LỖI LOGIC: Mức độ đau Ngày 3 không thể cao hơn Ngày 1.", None
    elif thoi_diem == "Tái khám Ngày 7":
        if p7v > p3v or p7v > p1v or p3v > p1v: return "⚠️ LỖI LOGIC: Mức độ đau Ngày 7 không thể cao hơn các ngày trước đó.", None

    actuals = {}
    is_fully_recovered = False
    recovery_day_actual = -1

    # Logic xác định điểm dừng nếu đã hết đau (Ngưỡng <= 1)
    if thoi_diem != "Sau khi phẫu thuật xong":
        actuals[0] = p1v
        if p1v <= 1:
            is_fully_recovered = True
            recovery_day_actual = 1

        if ("Ngày 3" in thoi_diem or "Ngày 7" in thoi_diem) and not is_fully_recovered:
            actuals[1] = p3v
            if p3v <= 1:
                is_fully_recovered = True
                recovery_day_actual = 3

        if "Ngày 7" in thoi_diem and not is_fully_recovered:
            actuals[2] = p7v
            if p7v <= 1:
                is_fully_recovered = True
                recovery_day_actual = 7

    df = preprocess_input_v2(tuoi, gioi, ben, exp, pg, goc, chan, ord_lq, ha, ma, thuoc_uong, t_thuc)
    preds, errs = get_increasing_uncertainty_v3(df, actuals=actuals)

    fig, ax = plt.subplots(figsize=(8, 4.5))
    days = [1, 3, 7]
    ax.fill_between(days, [max(0, preds[i]-errs[i]) for i in range(3)], [min(5, preds[i]+errs[i]) for i in range(3)], color='#d9534f', alpha=0.1)

    past_idx = list(actuals.keys())
    if past_idx:
        ax.plot([days[i] for i in past_idx], [preds[i] for i in past_idx], 'o-', color='black', linewidth=4, label='Mức độ đau thực tế')

    future_idx = [i for i in range(3) if i not in past_idx]
    if future_idx and not is_fully_recovered:
        if past_idx:
            ax.plot([days[past_idx[-1]], days[future_idx[0]]], [preds[past_idx[-1]], preds[future_idx[0]]], '-', color='#d9534f', linewidth=3)
        ax.plot([days[i] for i in future_idx], [preds[i] for i in future_idx], 's-', color='#d9534f', linewidth=3, label='Mức độ đau dự kiến')

    for i, p in enumerate(preds[:3]):
        if i in past_idx or (not is_fully_recovered):
            ax.text(days[i], p + 0.2, f"{p:.1f}", ha='center', fontweight='bold')

    ax.set_title(f"Mức độ đau ({thoi_diem}, Thuốc: {thuoc_uong} viên)", fontsize=12)
    ax.set_ylim(-0.2, 5.5); ax.set_xticks(days); ax.set_ylabel("Mức độ đau (VAS)"); ax.legend()
    plt.tight_layout(); plt.close()

    res_text = (f"✅ CẬP NHẬT SAU PHẪU THUẬT:\n"
                f"• Thời gian thực hiện: {t_thuc} phút\n"
                f"• Thuốc giảm đau (Paracetamol 500mg) đã uống: {thuoc_uong} viên\n"
                f"----------------------------------------\n"
                f"🕒 DỰ KIẾN HỒI PHỤC:\n")

    for i in range(3):
        if i in past_idx:
            res_text += f"• Mức độ đau ngày {days[i]}: {preds[i]:.1f} / 5 (Thực tế)\n"
            if recovery_day_actual == days[i]: break
        elif not is_fully_recovered:
            res_text += f"• Mức độ đau ngày {days[i]}: {preds[i]:.1f} / 5 (Dự kiến)\n"

    if is_fully_recovered:
        res_text += "👉 Trạng thái: Bệnh nhân ĐÃ HẾT ĐAU."
    else:
        # Ẩn sai số nếu errs[3] = 0
        res_text += f"• Ngày hết đau trung bình: Ngày thứ {preds[3]:.1f}" + (f" (± {errs[3]:.1f} ngày)" if errs[3] > 0 else "")

    return res_text, fig

def update_ui(choice):
    if choice == "Sau khi phẫu thuật xong": return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
    elif choice == "Tái khám Ngày 1": return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
    elif choice == "Tái khám Ngày 3": return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)
    else: return gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)

# --- 5. UI LAYOUT ---

with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
    gr.Markdown("# 🦷 HỆ THỐNG DỰ ĐOÁN THỜI GIAN PHẪU THUẬT RĂNG KHÔN VÀ PHÂN TÍCH MỨC ĐỘ ĐAU")

    with gr.Row():
        with gr.Column():
            gr.Markdown("### 📋 Thông tin người bệnh")
            tuoi = gr.Number(label="Tuổi", value=2005)
            gioi_tinh = gr.Dropdown(["Nam", "Nữ"], label="Giới tính", value="Nữ")
            ben_pt = gr.Dropdown(["Trái", "Phải"], label="Bên PT", value="Trái")
            ha_mieng = gr.Number(label="Độ há miệng", value=52)
            linh_dong_ma = gr.Number(label="Độ linh động má", value=48)
        with gr.Column():
            gr.Markdown("### 👨‍⚕️ Kinh nghiệm PTV")
            kinh_nghiem = gr.Dropdown(["<5 năm", "5-10 năm", ">10 năm"], label="Kinh nghiệm PTV", value="<5 năm")
            gr.Markdown("### 🔍 Đặc điểm răng (X-quang)")
            pell_gregory = gr.Radio(["IA", "IB", "IC", "IIA", "IIB", "IIC", "IIIA", "IIIB", "IIIC"], label="Phân loại Pell & Gregory", value="IIB", elem_id="pell-gregory-grid")
            goc_nghieng = gr.Number(label="Góc nghiêng", value=7)
            chan_rang = gr.Dropdown(["1 - Mầm răng/Hình thành <1/3 chân", "2 - Hình thành 1/3 đến <2/3 chân", "3 - Hình thành >2/3 chân và chân chụm", "4 - Hình thành >2/3 chân và chân phân kỳ/cong"], label="Hình thái chân răng", value="4 - Hình thành >2/3 chân và chân phân kỳ/cong")
            ord_lq = gr.Dropdown(["Không", "Có"], label="Liên quan ORD", value="Không")

    with gr.Tabs():
        with gr.TabItem("🕒 Kết quả ước đoán"):
            btn1 = gr.Button("🚀 Dự đoán thời gian phẫu thuật & Biến chứng đau", variant="primary")
            with gr.Row():
                out_txt1 = gr.Textbox(label="Kết quả chi tiết", scale=1)
                out_plot1 = gr.Plot(label="Biểu đồ đau dự kiến", scale=1)
            btn1.click(predict_tab1, inputs=[tuoi, gioi_tinh, ben_pt, kinh_nghiem, pell_gregory, goc_nghieng, chan_rang, ord_lq, ha_mieng, linh_dong_ma], outputs=[out_txt1, out_plot1])

        with gr.TabItem("📝 Dự kiến sau phẫu thuật"):
            with gr.Row():
                thoi_diem = gr.Dropdown(["Sau khi phẫu thuật xong", "Tái khám Ngày 1", "Tái khám Ngày 3", "Tái khám Ngày 7"], label="Sau phẫu thuật/Ngày tái khám", value="Sau khi phẫu thuật xong")
                t_act = gr.Number(label="Thời gian mổ thực tế (phút)", value=20)
                thuoc_p = gr.Number(label="Số viên thuốc giảm đau (Paracetamol 500mg) đã uống", value=5)
            with gr.Row():
                p_choices = [str(i) for i in range(6)]
                p1_in = gr.Dropdown(p_choices, label="Mức độ đau Ngày 1 (Thực tế)", visible=False, value="0")
                p3_in = gr.Dropdown(p_choices, label="Mức độ đau Ngày 3 (Thực tế)", visible=False, value="0")
                p7_in = gr.Dropdown(p_choices, label="Mức độ đau Ngày 7 (Thực tế)", visible=False, value="0")

            thoi_diem.change(update_ui, inputs=[thoi_diem], outputs=[p1_in, p3_in, p7_in])
            btn3 = gr.Button("🔄 Diễn tiến đau", variant="secondary")
            with gr.Row():
                out_txt3 = gr.Textbox(label="Chi tiết", scale=1)
                out_plot3 = gr.Plot(label="Biểu đồ diễn tiến đau", scale=1)
            btn3.click(predict_tab3, inputs=[tuoi, gioi_tinh, ben_pt, kinh_nghiem, pell_gregory, goc_nghieng, chan_rang, ord_lq, ha_mieng, linh_dong_ma, thuoc_p, t_act, thoi_diem, p1_in, p3_in, p7_in], outputs=[out_txt3, out_plot3])

demo.launch(theme=gr.themes.Soft(font='sans-serif'), share=True, debug=True)