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

plt.style.use('seaborn-v0_8-muted')
sns.set_theme(style="whitegrid")

import warnings
warnings.filterwarnings('ignore')

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_v1.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}")

def preprocess_input(tuoi, gioi_tinh, ben_pt, kinh_nghiem,
                     pell_gregory, goc, chan_rang, ord_lq,
                     ha_mieng, ma, so_thuoc):

    # Tách phân loại Pell & Gregory thành Cành đứng và R7
    cd = pell_gregory[:-1]  # VD: 'IIB' -> 'II'
    r7 = pell_gregory[-1]   # VD: 'IIB' -> 'B'

    # Tách lấy ID của Hình thái chân răng
    chan_val = int(str(chan_rang)[0])

    raw_data = {
        'Tuổi': tuoi, 'Giới tính': gioi_tinh, 'Bên PT': ben_pt,
        'Kinh nghiệm PTV': kinh_nghiem, 'Tương quan R7': r7,
        'Tương quan cành đứng': cd, 'Góc nghiêng': goc,
        'Hình thái chân răng': chan_val, 'Liên quan ORD': ord_lq,
        'Độ há miệng': ha_mieng, 'Độ linh động má': ma,
        'Số viên thuốc GD': so_thuoc,
        'Thời gian PT (phút)': 0 # Placeholder sẽ cập nhật sau
    }
    df = pd.DataFrame([raw_data])

    # Tính toán các đặc trưng (Feature Engineering)
    df['Age'] = 2025 - df['Tuổi']
    df['R7_Grouped'] = 'Class A' if df['Tương quan R7'].iloc[0] == 'A' else 'Class B/C'
    canh_dung_val = df['Tương quan cành đứng'].iloc[0]
    df['Ramus_Grouped'] = 'Class I/II' if canh_dung_val in ['I', 'II'] else 'Class III'

    if chan_val in [2, 3]: df['Root_Grouped'] = 'Easy (2/3)'
    elif chan_val in [1, 4]: df['Root_Grouped'] = 'Hard (1/4)'
    else: df['Root_Grouped'] = 'Other'

    df['Age_Group'] = pd.cut(df['Age'], bins=[0, 22, 30, 100], labels=['Young (<22)', 'Adult (22-30)', 'Senior (>30)'])

    return df

def predict_with_uncertainty(model, input_df):
    rf_model = model.named_steps['model']
    preprocessor = model.named_steps['preprocessor']
    X_transformed = preprocessor.transform(input_df)
    predictions = [tree.predict(X_transformed)[0] for tree in rf_model.estimators_]
    mean_pred = np.mean(predictions)
    std_pred = np.std(predictions)
    return mean_pred, std_pred

def predict_with_error(model, input_df):
    mean_pred, std_pred = predict_with_uncertainty(model, input_df)
    error_margin = 1.0 * std_pred # Giữ nguyên hệ số như code cũ của user
    return mean_pred, error_margin

def predict_general_v2(tuoi, gioi, ben, exp, pell_gregory, goc, chan, ord_lq, ha, ma, thuoc):
    df = preprocess_input(tuoi, gioi, ben, exp, pell_gregory, goc, chan, ord_lq, ha, ma, thuoc)
    time_mean, time_err = predict_with_error(models['Op_Time'], df)
    df['Thời gian PT (phút)'] = time_mean

    p1_m, p1_e = predict_with_error(models['Pain_D1'], df)
    p3_m, p3_e = predict_with_error(models['Pain_D3'], df)
    p7_m, p7_e = predict_with_error(models['Pain_D7'], df)
    end_m, end_e = predict_with_error(models['End_Day'], df)

    res_text = (
        f"⏱️ THỜI GIAN NHỔ RĂNG DỰ KIẾN: {time_mean:.1f} ± {time_err:.1f} phút\n"
        f"(Dải biến thiên: {max(0, time_mean-time_err):.1f} - {time_mean+time_err:.1f} phút)\n\n"
        f"📈 DỰ BÁO MỨC ĐỘ ĐAU:\n"
        f"• Ngày 1: {p1_m:.1f} (±{p1_e:.1f})\n"
        f"• Ngày 3: {p3_m:.1f} (±{p3_e:.1f})\n"
        f"• Ngày 7: {p7_m:.1f} (±{p7_e:.1f})\n"
        f"• Hết đau hoàn toàn: Ngày thứ {end_m:.1f} (±{end_e:.1f})"
    )

    fig, ax = plt.subplots(figsize=(8, 4.5))
    days = np.array([1, 3, 7])
    means = np.array([p1_m, p3_m, p7_m])
    errors = np.array([p1_e, p3_e, p7_e])

    ax.fill_between(days, means - errors, means + errors, color='#4A90E2', alpha=0.2, label='Dải sai số dự báo')
    ax.plot(days, means, 'o-', color='#4A90E2', linewidth=3, markersize=8, label='Mức đau trung bình')

    for i, m in enumerate(means):
        ax.text(days[i], m + 0.3, f"{m:.1f}", ha='center', fontweight='bold', color='#2C3E50')

    ax.set_ylim(0, 5.5)
    ax.set_xticks([1, 3, 7])
    ax.set_xticklabels(['Ngày 1', 'Ngày 3', 'Ngày 7'])
    ax.set_title("Biểu đồ diễn tiến đau và sai số dự báo", fontsize=14, pad=15)
    ax.set_ylabel("Mức độ đau (VAS)")
    ax.legend()
    plt.tight_layout()
    plt.close()

    return res_text, fig

def predict_experience(tuoi, gioi, ben, pell_gregory, goc, chan, ord_lq, ha, ma, thuoc):
    exp_mapping = ["<5 năm", "5-10 năm", ">10 năm"]
    times = []
    for e in exp_mapping:
        df = preprocess_input(tuoi, gioi, ben, e, pell_gregory, goc, chan, ord_lq, ha, ma, thuoc)
        times.append(models['Op_Time'].predict(df)[0])

    fig, ax = plt.subplots(figsize=(7, 3.5))
    colors = ['#FFADAD', '#FFD6A5', '#CAFFBF']
    bars = ax.barh(["Dưới 5 năm", "5-10 năm", "Trên 10 năm"], times, color=colors, height=0.6)
    ax.bar_label(bars, fmt='%.1f phút', padding=5, fontweight='bold')
    ax.set_title("So sánh thời gian nhổ theo kinh nghiệm bác sĩ", fontsize=12)
    ax.set_xlim(0, max(times) * 1.3)
    plt.tight_layout()
    plt.close()

    advice = f"💡 Bác sĩ trên 10 năm kinh nghiệm làm nhanh hơn bác sĩ mới khoảng {times[0]-times[2]:.1f} phút."
    return fig, advice

def predict_actual_with_range(tuoi, gioi, ben, exp, pell_gregory, goc, chan, ord_lq, ha, ma, thuoc_da_uong, time_actual):
    df = preprocess_input(tuoi, gioi, ben, exp, pell_gregory, goc, chan, ord_lq, ha, ma, thuoc_da_uong)
    df['Thời gian PT (phút)'] = time_actual

    targets = {
        'Ngày 1': models['Pain_D1'],
        'Ngày 3': models['Pain_D3'],
        'Ngày 7': models['Pain_D7'],
        'Hết đau': models['End_Day']
    }

    results_mean = []
    results_std = []
    labels = ['Ngày 1', 'Ngày 3', 'Ngày 7']

    for label in labels:
        m, s = predict_with_uncertainty(targets[label], df)
        results_mean.append(m)
        results_std.append(s)

    end_mean, end_std = predict_with_uncertainty(targets['Hết đau'], df)

    fig, ax = plt.subplots(figsize=(8, 4))
    x_vals = [1, 3, 7]
    means = np.array(results_mean)
    stds = np.array(results_std)

    ax.plot(x_vals, means, 'o-', color='#d9534f', linewidth=2, label='Dự báo trung bình')
    ax.fill_between(x_vals, means - 1.96*stds, means + 1.96*stds, color='#d9534f', alpha=0.2, label='Khoảng sai số dự báo')

    for i, txt in enumerate(means):
        ax.annotate(f'{txt:.1f}', (x_vals[i], means[i]), xytext=(0, 10), textcoords='offset points', ha='center')

    ax.set_ylim(0, 5.5)
    ax.set_xticks([1, 3, 7])
    ax.set_xticklabels(['Ngày 1', 'Ngày 3', 'Ngày 7'])
    ax.set_title(f"Diễn tiến đau thực tế (Thời gian: {time_actual}ph, Thuốc: {thuoc_da_uong} viên)")
    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: {time_actual} phút\n"
                f"• Thuốc giảm đau (Paracetamol 500mg) đã uống: {thuoc_da_uong} viên\n"
                f"----------------------------------------\n"
                f"🕒 DỰ KIẾN HỒI PHỤC:\n"
                f"• Ngày hết đau trung bình: Ngày thứ {end_mean:.1f} (± {end_std*1.96:.1f} ngày)\n"
                f"• Mức độ đau ngày đầu: {results_mean[0]:.1f} / 5")

    return res_text, fig
    
with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
    gr.Markdown("# 🦷 HỆ THỐNG DỰ BÁO 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)

            so_thuoc = gr.Number(label="Số viên thuốc GD", value=5, visible=False)

        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 răng", 
                "2 - Hình thành >1/3 và <2/3 chân răng", 
                "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 phía chóp/có >2 chân"
            ], 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 phía chóp/có >2 chân")
            
            ord_lq = gr.Dropdown(["Không", "Có"], label="Liên quan ORD", value="Không")

    base_inputs = [tuoi, gioi_tinh, ben_pt, kinh_nghiem, pell_gregory, goc_nghieng, chan_rang, ord_lq, ha_mieng, linh_dong_ma]
    inputs_all = base_inputs + [so_thuoc]

    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", variant="primary")
            with gr.Row():
                out_txt1 = gr.Textbox(label="Kết quả dự báo chi tiết", lines=8)
                out_plot1 = gr.Plot(label="Diễn tiến đau")
            btn1.click(predict_general_v2, inputs=inputs_all, outputs=[out_txt1, out_plot1])

        # Tab So sánh ẨN trên UI
        with gr.TabItem("👨‍⚕️ So sánh theo kinh nghiệm bác sĩ", visible=False):
            btn2 = gr.Button("📊 Dự đoán thời gian phẫu thuật theo kinh nghiệm bác sĩ")
            with gr.Row():
                out_plot2 = gr.Plot(label="Biểu đồ so sánh")
                out_txt2 = gr.Textbox(label="Nhận xét", lines=2)
            inputs_for_exp = [tuoi, gioi_tinh, ben_pt, pell_gregory, goc_nghieng, chan_rang, ord_lq, ha_mieng, linh_dong_ma, so_thuoc]
            btn2.click(predict_experience, inputs=inputs_for_exp, outputs=[out_plot2, out_txt2])

        with gr.TabItem("📝 Dự kiến sau phẫu thuật"):
            gr.Markdown("### Nhập thông số thực tế sau ca mổ để tư vấn cho bệnh nhân")
            with gr.Row():
                time_actual = gr.Number(label="Thời gian mổ thực tế (phút)", value=20)
                thuoc_post = gr.Number(label="Số viên thuốc giảm đau (Paracetamol 500mg) đã uống", value=10)

            btn3 = gr.Button("🔄 Dự đoán thời gian hết đau", variant="secondary")

            with gr.Row():
                out_txt3 = gr.Textbox(label="Chi tiết", lines=8)
                out_plot3 = gr.Plot(label="Biểu đồ diễn tiến đau")

            btn3.click(predict_actual_with_range,
                      inputs=base_inputs + [thuoc_post, time_actual],
                      outputs=[out_txt3, out_plot3])

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