import gradio as gr
import joblib
import pandas as pd

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 = pd.DataFrame({
        "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]
    })

    input_data['gender'] = pd.Categorical(input_data.gender)
    input_data['ever_married'] = pd.Categorical(input_data.ever_married)
    input_data['Residence_type'] = pd.Categorical(input_data.Residence_type)
    input_data['work_type'] = pd.Categorical(input_data.work_type)
    input_data['smoking_status'] = pd.Categorical(input_data.smoking_status)

    cats = ["gender", "ever_married", "work_type", "Residence_type", "smoking_status"]

    input_data[cats] = input_data[cats].apply(lambda x: x.cat.codes)    

    prediction = m.predict(input_data)
    
    if prediction == [0]:
        return "didn't have a stroke"
    return "have a stroke"
    


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

output = gr.outputs.Textbox()

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

intf.launch()