import gradio as gr
import joblib
import numpy as np
import pandas as pd
import warnings
warnings.filterwarnings("ignore")
import os

# ── Load model ───────────────────────────────────────────────────────────────
MODEL_PATH = "Random_Forest.pth"
model = joblib.load(MODEL_PATH)

# ── Feature order must match training data ────────────────────────────────────
FEATURE_ORDER = [
    "Sex", "Age", "Haemoglobin", "Hematocrit", "Lymphocytes",
    "MCH", "MCHC", "MCV", "Monocytes", "MPV", "PCT", "PDW",
    "Platelet Count", "Polymorphs", "RBC Count", "RDW SD",
    "WBC Count", "Eosinophils", "Creatinine", "Urea", "GFR "
]

# ── Prediction function ───────────────────────────────────────────────────────
def predict_ckd(
    sex, age, haemoglobin, hematocrit, lymphocytes,
    mch, mchc, mcv, monocytes, mpv, pct, pdw,
    platelet_count, polymorphs, rbc_count, rdw_sd,
    wbc_count, eosinophils, creatinine, urea, gfr
):
    try:
        values = [
            sex, age, haemoglobin, hematocrit, lymphocytes,
            mch, mchc, mcv, monocytes, mpv, pct, pdw,
            platelet_count, polymorphs, rbc_count, rdw_sd,
            wbc_count, eosinophils, creatinine, urea, gfr
        ]
        features = pd.DataFrame([values], columns=FEATURE_ORDER)

        pred = model.predict(features)[0]
        proba = model.predict_proba(features)[0]

        ckd_prob    = round(float(proba[1]) * 100, 2)
        no_ckd_prob = round(float(proba[0]) * 100, 2)

        if pred == 1:
            label   = "⚠️ CKD Detected"
            details = (
                f"**Prediction: Chronic Kidney Disease (CKD)**\n\n"
                f"- CKD Probability: **{ckd_prob}%**\n"
                f"- No CKD Probability: {no_ckd_prob}%\n\n"
                f"> ⚠️ This result suggests the presence of Chronic Kidney Disease. "
                f"Please consult a qualified nephrologist for clinical evaluation."
            )
        else:
            label   = "✅ No CKD Detected"
            details = (
                f"**Prediction: No Chronic Kidney Disease**\n\n"
                f"- No CKD Probability: **{no_ckd_prob}%**\n"
                f"- CKD Probability: {ckd_prob}%\n\n"
                f"> ✅ Based on the provided blood parameters, no CKD pattern was detected. "
                f"Regular health check-ups are still recommended."
            )

        return label, details

    except Exception as e:
        import traceback
        err = traceback.format_exc()
        return "❌ Prediction Error", f"```\n{err}\n```"

# ── UI ────────────────────────────────────────────────────────────────────────
DESCRIPTION = """
## Chronic Kidney Disease (CKD) Prediction System
**Research Project — PAF-IAST, HEC-Funded**

This tool uses a **Random Forest classifier** trained on routine blood test parameters 
to assist in early screening of Chronic Kidney Disease.  

> **Disclaimer:** This application is intended for **research purposes only** and does 
> not constitute medical advice. Clinical decisions must always be made by a licensed physician.

**Project Owner:** Dr. Muhammad Rizwan
"""

ARTICLE = """
### About the Model
The prediction model was developed as part of an HEC-funded research project at PAF-IAST. 
The Random Forest classifier was selected after a comprehensive comparison of 8 machine learning 
algorithms, including XGBoost, CatBoost, Gradient Boosting, and Logistic Regression. 
Explainability was ensured using SHAP (SHapley Additive exPlanations) and LIME analyses. 
Top predictive features include **Creatinine**, **GFR**, and **Urea** levels.
"""

with gr.Blocks(
    title="CKD Prediction | PAF-IAST",
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="indigo"),
    css="""
    .result-box { border-radius: 10px; padding: 10px; }
    footer { display: none !important; }
    """
) as demo:

    gr.Markdown(DESCRIPTION)

    with gr.Row():
        # ── Column 1: Demographic + CBC ──────────────────────────────────────
        with gr.Column():
            gr.Markdown("### 👤 Demographic")
            sex = gr.Radio(
                choices=[("Male", 1), ("Female", 0)],
                label="Sex", value=1
            )
            age = gr.Slider(minimum=1, maximum=100, step=1, value=49, label="Age (years)")

            gr.Markdown("### 🩸 Complete Blood Count")
            haemoglobin   = gr.Number(value=11.0, label="Haemoglobin (g/dL)")
            hematocrit    = gr.Number(value=32.0, label="Hematocrit (%)")
            rbc_count     = gr.Number(value=3.7,  label="RBC Count (millions/µL)")
            wbc_count     = gr.Number(value=7.2,  label="WBC Count (×10³/µL)")
            platelet_count = gr.Number(value=200, label="Platelet Count (×10³/µL)")

        # ── Column 2: Red Cell Indices + Differential ─────────────────────────
        with gr.Column():
            gr.Markdown("### 📊 Red Cell Indices")
            mch    = gr.Number(value=29.2, label="MCH (pg)")
            mchc   = gr.Number(value=35.0, label="MCHC (g/dL)")
            mcv    = gr.Number(value=85.0, label="MCV (fL)")
            rdw_sd = gr.Number(value=15.7, label="RDW-SD (fL)")
            mpv    = gr.Number(value=8.7,  label="MPV (fL)")
            pct    = gr.Number(value=0.17, label="PCT (%)")
            pdw    = gr.Number(value=37.0, label="PDW (%)")

            gr.Markdown("### 🔬 Differential Count (%)")
            lymphocytes  = gr.Number(value=22, label="Lymphocytes (%)")
            monocytes    = gr.Number(value=4,  label="Monocytes (%)")
            polymorphs   = gr.Number(value=71, label="Polymorphs (%)")
            eosinophils  = gr.Number(value=2,  label="Eosinophils (%)")

        # ── Column 3: Kidney Function ─────────────────────────────────────────
        with gr.Column():
            gr.Markdown("### 🫘 Kidney Function Tests")
            creatinine = gr.Number(value=7.0,   label="Creatinine (mg/dL)")
            urea       = gr.Number(value=112.0, label="Blood Urea (mg/dL)")
            gfr        = gr.Number(value=7.2,   label="GFR (mL/min/1.73m²)")

            gr.Markdown("### 🔍 Prediction")
            predict_btn = gr.Button("🔮 Predict", variant="primary", size="lg")

            result_label   = gr.Textbox(label="Result", interactive=False, lines=1)
            result_details = gr.Markdown(label="Details")

    gr.Markdown(ARTICLE)

    # ── Example inputs ────────────────────────────────────────────────────────
    gr.Examples(
        examples=[
            # CKD positive (high creatinine, low GFR)
            [1, 49, 9.3, 25.0, 12, 28.0, 36.0, 78.0, 5, 10.9, 0.08, 41.8,
             76, 71, 3.5, 16.2, 5.5, 3, 13.2, 112, 4.06],
            # No CKD (normal values)
            [0, 36, 13.4, 40.0, 25, 31.0, 34.0, 90.0, 4, 8.5, 0.17, 36.0,
             274, 68, 4.4, 12.9, 7.5, 2, 0.80, 37, 105.01],
        ],
        inputs=[
            sex, age, haemoglobin, hematocrit, lymphocytes,
            mch, mchc, mcv, monocytes, mpv, pct, pdw,
            platelet_count, polymorphs, rbc_count, rdw_sd,
            wbc_count, eosinophils, creatinine, urea, gfr
        ],
        label="Sample Inputs (click to auto-fill)"
    )

    predict_btn.click(
        fn=predict_ckd,
        inputs=[
            sex, age, haemoglobin, hematocrit, lymphocytes,
            mch, mchc, mcv, monocytes, mpv, pct, pdw,
            platelet_count, polymorphs, rbc_count, rdw_sd,
            wbc_count, eosinophils, creatinine, urea, gfr
        ],
        outputs=[result_label, result_details]
    )

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