import gradio as gr
from huggingface_hub import hf_hub_download
import pickle
import numpy as np
import os

REPO_ID = "Umranz/mediscan-symptom-classifier"

def load_models():
    files = ["svm.pkl", "logistic.pkl", "random_forest.pkl", "naive_bayes.pkl", "voting_ensemble.pkl", "label_encoder.pkl", "tfidf.pkl"]
    loaded = {}
    for f in files:
        path = hf_hub_download(repo_id=REPO_ID, filename=f)
        with open(path, "rb") as file:
            loaded[f.replace(".pkl", "")] = pickle.load(file)
    return loaded

print("Loading models...")
M  = load_models()
tfidf    = M["tfidf"]
le       = M["label_encoder"]
ensemble = M["voting_ensemble"]
models   = {
    "SVM"            : M["svm"],
    "Logistic Reg"   : M["logistic"],
    "Random Forest"  : M["random_forest"],
    "Naive Bayes"    : M["naive_bayes"],
}
print("✅ Models loaded!")

SEVERITY = {
    "Fungal infection"            : ("🟡", "Mild"),
    "Allergy"                     : ("🟡", "Mild"),
    "GERD"                        : ("🟡", "Mild"),
    "Chronic cholestasis"         : ("🟠", "Moderate"),
    "Drug Reaction"               : ("🟠", "Moderate"),
    "Peptic ulcer disease"        : ("🟠", "Moderate"),
    "AIDS"                        : ("🔴", "Severe"),
    "Diabetes"                    : ("🟠", "Moderate"),
    "Gastroenteritis"             : ("🟡", "Mild"),
    "Bronchial Asthma"            : ("🟠", "Moderate"),
    "Hypertension"                : ("🔴", "Severe"),
    "Migraine"                    : ("🟡", "Mild"),
    "Cervical spondylosis"        : ("🟡", "Mild"),
    "Paralysis (brain hemorrhage)": ("🔴", "Severe"),
    "Jaundice"                    : ("🟠", "Moderate"),
    "Malaria"                     : ("🔴", "Severe"),
    "Chicken pox"                 : ("🟡", "Mild"),
    "Dengue"                      : ("🔴", "Severe"),
    "Typhoid"                     : ("🟠", "Moderate"),
    "hepatitis A"                 : ("🟠", "Moderate"),
    "Hepatitis B"                 : ("🔴", "Severe"),
    "Hepatitis C"                 : ("🔴", "Severe"),
    "Hepatitis D"                 : ("🔴", "Severe"),
    "Hepatitis E"                 : ("🟠", "Moderate"),
    "Alcoholic hepatitis"         : ("🟠", "Moderate"),
    "Tuberculosis"                : ("🔴", "Severe"),
    "Common Cold"                 : ("🟢", "Low"),
    "Pneumonia"                   : ("🔴", "Severe"),
    "Dimorphic hemmorhoids(piles)": ("🟡", "Mild"),
    "Heart attack"                : ("🔴", "Critical"),
    "Varicose veins"              : ("🟡", "Mild"),
    "Hypothyroidism"              : ("🟠", "Moderate"),
    "Hyperthyroidism"             : ("🟠", "Moderate"),
    "Hypoglycemia"                : ("🔴", "Severe"),
    "Osteoarthristis"             : ("🟡", "Mild"),
    "Arthritis"                   : ("🟡", "Mild"),
    "Vertigo"                     : ("🟡", "Mild"),
    "Acne"                        : ("🟢", "Low"),
    "Urinary tract infection"     : ("🟡", "Mild"),
    "Psoriasis"                   : ("🟡", "Mild"),
    "Impetigo"                    : ("🟡", "Mild"),
}

def predict(symptoms, threshold):
    if not symptoms.strip():
        return (
            "⚠️ Please enter your symptoms.",
            "",
            "",
            ""
        )

    vec   = tfidf.transform([symptoms])
    proba = ensemble.predict_proba(vec)[0]
    top3  = np.argsort(proba)[::-1][:3]

    top_idx    = top3[0]
    top_label  = le.classes_[top_idx]
    top_conf   = proba[top_idx] * 100
    sev_emoji, sev_label = SEVERITY.get(top_label, ("⚪", "Unknown"))

    if top_conf < threshold:
        main_result = (
            f"⚠️ **Low Confidence ({top_conf:.1f}%)** — Please provide more specific symptoms.\n\n"
            f"Best guess: **{top_label}** but confidence is below your threshold of {threshold}%."
        )
        return main_result, "", "", ""
    else:
        main_result = (
            f"## {sev_emoji} {top_label}\n"
            f"**Confidence:** {top_conf:.1f}%\n\n"
            f"**Severity:** {sev_emoji} {sev_label}\n\n"
            f"{'█' * int(top_conf // 5)}{'░' * (20 - int(top_conf // 5))} {top_conf:.1f}%"
        )

    top3_result = "## 🏆 Top 3 Predictions\n\n"
    for rank, idx in enumerate(top3):
        label  = le.classes_[idx]
        conf   = proba[idx] * 100
        s_emoji, s_label = SEVERITY.get(label, ("⚪", "Unknown"))
        bar    = "█" * int(conf // 5) + "░" * (20 - int(conf // 5))
        top3_result += (
            f"**{rank+1}. {label}**  {s_emoji} {s_label}\n"
            f"{bar} {conf:.1f}%\n\n"
        )

    agreement = "## 🤖 Model Votes\n\n"
    votes     = {}
    for name, model in models.items():
        pred        = le.classes_[model.predict(vec)[0]]
        votes[name] = pred
        match       = "✅" if pred == top_label else "🔄"
        agreement  += f"{match} **{name}** → {pred}\n\n"

    all_agree = len(set(votes.values())) == 1
    agreement += (
        "\n🟢 **All models agree!**" if all_agree
        else "\n🟡 **Models have different opinions — consider consulting a doctor.**"
    )

    disclaimer = (
        "## ⚠️ Medical Disclaimer\n\n"
        "This tool is for **educational purposes only** and does **NOT** replace "
        "professional medical advice. Always consult a qualified healthcare provider "
        "for diagnosis and treatment.\n\n"
        "**If you have a medical emergency, call your local emergency number immediately.**"
    )

    return main_result, top3_result, agreement, disclaimer

EXAMPLES = [
    ["fever, chills, headache, muscle pain, sweating",               50],
    ["itching, skin rash, nodal skin eruptions, dischromic patches", 50],
    ["chest pain, shortness of breath, fatigue, sweating",           50],
    ["sneezing, runny nose, cough, sore throat, congestion",         50],
    ["fatigue, weight loss, high fever, night sweats, cough",        50],
]

with gr.Blocks(title="MediScan AI") as demo:

    gr.Markdown("""
    # 🩺 MediScan AI — Medical Symptom Classifier
    **4 ML Models + Voting Ensemble** | DistilBERT-level accuracy with traditional ML
    > Enter your symptoms separated by commas for instant multi-model analysis.
    """)

    with gr.Row():
        with gr.Column(scale=2):
            symptoms_input = gr.Textbox(
                lines=4,
                placeholder="e.g. fever, chills, headache, muscle pain, fatigue...",
                label="🔍 Describe Your Symptoms",
                max_lines=8
            )
            threshold_slider = gr.Slider(
                minimum=10,
                maximum=90,
                value=50,
                step=5,
                label="⚙️ Confidence Threshold (%)",
                info="Predictions below this % will show a low-confidence warning"
            )
            analyze_btn = gr.Button(
                "🔍 Analyze Symptoms",
                variant="primary",
                size="lg"
            )

        with gr.Column(scale=3):
            main_output = gr.Markdown(label="Primary Diagnosis")

    with gr.Row():
        top3_output     = gr.Markdown(label="Top 3 Predictions")
        agreement_output = gr.Markdown(label="Model Agreement")

    disclaimer_output = gr.Markdown()

    gr.Examples(
        examples=EXAMPLES,
        inputs=[symptoms_input, threshold_slider],
        label="💡 Try These Examples"
    )

    with gr.Accordion("ℹ️ About MediScan AI", open=False):
        gr.Markdown("""
        ## 🧠 How It Works
        MediScan AI runs your symptoms through **4 independent ML models simultaneously:**

        | Model | Strength |
        |---|---|
        | **SVM** | Best accuracy on text classification |
        | **Logistic Regression** | Fast, reliable baseline |
        | **Random Forest** | Handles noisy input well |
        | **Naive Bayes** | Great for keyword-based symptoms |

        A **Soft Voting Ensemble** combines all 4 predictions for the final result.

        ## 📊 Dataset
        - **Source:** Gretel AI Symptom to Diagnosis dataset
        - **Diseases:** 24 unique conditions
        - **Features:** TF-IDF with bigrams (5000 features)

        ## 👨‍💻 Built By
        Umranz — [HuggingFace Profile](https://huggingface.co/Umranz)
        """)

    analyze_btn.click(
        fn=predict,
        inputs=[symptoms_input, threshold_slider],
        outputs=[main_output, top3_output, agreement_output, disclaimer_output]
    )

    symptoms_input.submit(
        fn=predict,
        inputs=[symptoms_input, threshold_slider],
        outputs=[main_output, top3_output, agreement_output, disclaimer_output]
    )

demo.launch()