import gradio as gr
import pandas as pd
import joblib

# Load the saved model
model = joblib.load("stroke_prediction_model.pkl")

# Define feature names (from cleaned dataset)
feature_names = [
    'age', 'hypertension', 'heart_disease', 'avg_glucose_level', 'bmi',
    'gender_Male', 'ever_married_Yes', 'work_type_Never_worked',
    'work_type_Private', 'work_type_Self-employed', 'work_type_children',
    'Residence_type_Urban', 'smoking_status_formerly smoked',
    'smoking_status_never smoked', 'smoking_status_smokes'
]

# Define prediction function
def predict_stroke(age, hypertension, heart_disease, avg_glucose_level, bmi,
                   gender, ever_married, work_type, residence_type, smoking_status):
    
    # Encode categorical features
    gender_Male = 1 if gender == "Male" else 0
    ever_married_Yes = 1 if ever_married == "Yes" else 0
    
    work_type_Never_worked = work_type_Private = work_type_Self_employed = work_type_children = 0
    if work_type == "Never worked":
        work_type_Never_worked = 1
    elif work_type == "Private":
        work_type_Private = 1
    elif work_type == "Self_employed":
        work_type_Self_employed = 1
    elif work_type == "children":
        work_type_children = 1

    Residence_type_Urban = 1 if residence_type == "Urban" else 0

    smoke_former = smoke_never = smoke_yes = 0
    if smoking_status == "formerly smoked":
        smoke_former = 1
    elif smoking_status == "never smoked":
        smoke_never = 1
    elif smoking_status == "smokes":
        smoke_yes = 1

    # Build input DataFrame
    input_data = pd.DataFrame({
        'age': [age],
        'hypertension': [hypertension],
        'heart_disease': [heart_disease],
        'avg_glucose_level': [avg_glucose_level],
        'bmi': [bmi],
        'gender_Male': [gender_Male],
        'ever_married_Yes': [ever_married_Yes],
        'work_type_Never_worked': [work_type_Never_worked],
        'work_type_Private': [work_type_Private],
        'work_type_Self_employed': [work_type_Self_employed],
        'work_type_children': [work_type_children],
        'Residence_type_Urban': [Residence_type_Urban],
        'smoking_status_formerly smoked': [smoke_former],
        'smoking_status_never smoked': [smoke_never],
        'smoking_status_smokes': [smoke_yes]
    })

    # Make prediction
    prediction = model.predict(input_data)[0]
    probability = model.predict_proba(input_data)[0][1]

    return "⚠️ High Risk of Stroke" if prediction == 1 else "✅ Low Risk of Stroke", f"{probability:.2%}"

# Create Gradio Interface
interface = gr.Interface(
    fn=predict_stroke,
    inputs=[
        gr.Slider(0, 100, value=45, label="Age"),
        gr.Radio(["No", "Yes"], label="Hypertension"),
        gr.Radio(["No", "Yes"], label="Heart Disease"),
        gr.Number(label="Average Glucose Level"),
        gr.Number(label="BMI"),
        gr.Radio(["Female", "Male"], label="Gender"),
        gr.Radio(["No", "Yes"], label="Ever Married"),
        gr.Radio(["Never worked", "Private", "Self_employed", "children"], label="Work Type"),
        gr.Radio(["Rural", "Urban"], label="Residence Type"),
        gr.Radio(["never smoked", "formerly smoked", "smokes"], label="Smoking Status")
    ],
    outputs=[
        gr.Textbox(label="Prediction"),
        gr.Textbox(label="Stroke Probability")
    ],
    title="🩺 Stroke Risk Prediction",
    description="Predict the risk of stroke based on patient health data.",
    allow_flagging='never'
)

# Launch locally for testing
if __name__ == "__main__":
    interface.launch()