import joblib
import pandas as pd
import gradio as gr

# Load model and preprocessor
try:
    model = joblib.load("stroke_rf_model.pkl")
    preprocessor = joblib.load("preprocessor.pkl")
    print("✅ Model and preprocessor loaded successfully.")
except Exception as e:
    print(f"❌ Error loading model/preprocessor: {str(e)}")
    model = None
    preprocessor = None

# Define prediction function
def predict_stroke(gender, age, hypertension, heart_disease, ever_married,
                   work_type, residence_type, avg_glucose_level, smoking_status, bmi):
    
    if model is None or preprocessor is None:
        return "Error: Model or preprocessor not loaded."

    data = pd.DataFrame([{
        'gender': gender,
        'age': age,
        'hypertension': hypertension,
        'heart_disease': heart_disease,
        'ever_married': ever_married,
        'work_type': work_type,
        'Residence_type': residence_type,
        'avg_glucose_level': avg_glucose_level,
        'smoking_status': smoking_status,
        'bmi': bmi
    }])

    try:
        processed_data = preprocessor.transform(data)
        prediction = model.predict(processed_data)
        return "⚠️ Stroke Risk" if prediction[0] == 1 else "✅ No Stroke Risk"
    except Exception as e:
        print(f"❌ Prediction error: {str(e)}")
        return f"Error: {str(e)}"

# Gradio Interface
iface = gr.Interface(
    fn=predict_stroke,
    inputs=[
        gr.Radio(choices=["Male", "Female"], label="Gender"),
        gr.Slider(minimum=1, maximum=100, step=1, label="Age"),
        gr.Radio(choices=[0, 1], label="Hypertension (0=No, 1=Yes)"),
        gr.Radio(choices=[0, 1], label="Heart Disease (0=No, 1=Yes)"),
        gr.Dropdown(choices=["Yes", "No"], label="Ever Married"),
        gr.Dropdown(choices=["Private", "Self-employed", "Govt_job", "children", "Never_worked"], label="Work Type"),
        gr.Radio(choices=["Urban", "Rural"], label="Residence Type"),
        gr.Number(label="Average Glucose Level"),
        gr.Dropdown(choices=["never smoked", "formerly smoked", "smokes", "Unknown"], label="Smoking Status"),
        gr.Number(label="BMI")
    ],
    outputs="text",
    title="🩺 Stroke Risk Prediction App",
    description="Predict the likelihood of stroke based on health metrics.",
    allow_flagging="never"
)

# Launch the app
if __name__ == "__main__":
    iface.launch()