import gradio as gr
import joblib
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download

# =========================
# 1. Symptom Model Section
# =========================

model_path = hf_hub_download(
    repo_id="AWeirdDev/human-disease-prediction",
    filename="sklearn_model.joblib"
)

symptom_model = joblib.load(model_path)

TOTAL_FEATURES = 132

visible_symptoms = [
    "itching",
    "skin_rash",
    "continuous_sneezing",
    "shivering",
    "joint_pain",
    "stomach_pain",
    "acidity",
    "ulcers_on_tongue"
]

def predict_symptoms(selected):
    try:
        x = np.zeros(TOTAL_FEATURES)
        for i, symptom in enumerate(visible_symptoms):
            if symptom in selected:
                x[i] = 1

        x = x.reshape(1, -1)
        pred = symptom_model.predict(x)[0]
        return f"Predicted Disease: {pred}"
    except Exception as e:
        return f"Error: {str(e)}"


# =========================
# 2. Image Model Section
# =========================

skin_classes = [
    "Eczema",
    "Melanoma",
    "Atopic Dermatitis",
    "Basal Cell Carcinoma",
    "Melanocytic Nevi",
    "Benign Keratosis",
    "Psoriasis",
    "Tinea Ringworm"
]

def predict_skin_disease(image):
    try:
        if image is None:
            return "Please upload an image."

        # DEMO prediction only
        # Replace this block with your real CNN/ResNet model later
        image = image.resize((224, 224))
        image_array = np.array(image)

        # simple demo logic
        avg_pixel = image_array.mean()

        if avg_pixel < 85:
            pred = skin_classes[1]
        elif avg_pixel < 120:
            pred = skin_classes[0]
        elif avg_pixel < 160:
            pred = skin_classes[6]
        else:
            pred = skin_classes[4]

        return f"Predicted Skin Disease: {pred}"
    except Exception as e:
        return f"Error: {str(e)}"


# =========================
# 3. Gradio Interface
# =========================

with gr.Blocks() as demo:
    gr.Markdown("# 🩺 Hybrid Disease Prediction System")
    gr.Markdown("Predict diseases using symptoms or skin images.")

    with gr.Tab("Symptom Prediction"):
        gr.Markdown("### Select symptoms for disease prediction")
        symptom_input = gr.CheckboxGroup(
            choices=visible_symptoms,
            label="Select Symptoms"
        )
        symptom_output = gr.Textbox(label="Result")
        symptom_btn = gr.Button("Predict from Symptoms")
        symptom_btn.click(
            fn=predict_symptoms,
            inputs=symptom_input,
            outputs=symptom_output
        )

    with gr.Tab("Skin Disease from Image"):
        gr.Markdown("### Upload a skin image")
        image_input = gr.Image(type="pil", label="Upload Skin Image")
        image_output = gr.Textbox(label="Result")
        image_btn = gr.Button("Predict from Image")
        image_btn.click(
            fn=predict_skin_disease,
            inputs=image_input,
            outputs=image_output
        )

demo.launch()