import gradio as gr
import joblib

m = joblib.load('tree_model.pkl')

def predict(gender, ever_married, work_type, residence_type, smoking_status, age, hypertension, heart_disease, avg_glucose_level, bmi):
    input_data = {
        "gender": gender,
        "ever_married": ever_married,
        "work_type": work_type,
        "residence_type": residence_type,
        "smoking_status": smoking_status,
        "age": age,
        "hypertension": hypertension,
        "heart_disease": heart_disease,
        "avg_glucose_level": avg_glucose_level,
        "bmi": bmi
    }

    prediction = m.predict([input_data])[0]
    return prediction

inputs = [
    gr.inputs.Dropdown(choices=["Male", "Female", "Other"], label="Gender"),
    gr.inputs.Dropdown(choices=["Yes", "No"], label="Ever Married"),
    gr.inputs.Dropdown(choices=['Govt_job', 'Never_worked', 'Private', 'Self-employed', 'children'], label="Work Type"),
    gr.inputs.Dropdown(choices=["Urban", "Rural"], label="Residence Type"),
    gr.inputs.Dropdown(choices=["Unknown", "never smoked", "formerly smoked", "smokes"], label="Smoking Status"),
    gr.inputs.Number(label="Age"),
    gr.inputs.Number(label="Hypertension"),
    gr.inputs.Number(label="Heart Disease"),
    gr.inputs.Number(label="Average Glucose Level"),
    gr.inputs.Number(label="BMI")
]

output = gr.outputs.Label(label="Predição")

intf = gr.Interface(fn=predict, inputs=inputs, outputs=output)

intf.launch()