import gradio as gr
import numpy as np
from joblib import load

rf = load("model.pkl")

columns = [
    "BMI",
    "Smoking",
    "AlcoholDrinking",
    "Stroke",
    "PhysicalHealth",
    "MentalHealth",
    "DiffWalking",
    "Sex",
    "AgeEstimate",
    "Race",
    "Diabetic",
    "PhysicalActivity",
    "GenHealth",
    "SleepTime",
    "Asthma",
    "KidneyDisease",
    "SkinCancer"
]

def convert_to_int(s):
    return int(s[:2])

def predict(
    BMI,
    Smoking,
    AlcoholDrinking,
    Stroke,
    PhysicalHealth,
    MentalHealth,
    DiffWalking,
    Sex,
    AgeEstimate,
    Race,
    Diabetic,
    PhysicalActivity,
    GenHealth,
    SleepTime,
    Asthma,
    KidneyDisease,
    SkinCancer
):
    data = np.array(
        [
            [
                BMI,
                Smoking,
                AlcoholDrinking,
                Stroke,
                PhysicalHealth,
                MentalHealth,
                DiffWalking,
                Sex,
                convert_to_int(AgeEstimate),
                Race+1,
                Diabetic,
                PhysicalActivity,
                GenHealth+1,
                SleepTime,
                Asthma,
                KidneyDisease,
                SkinCancer
            ]
        ]
    )
    pred = rf.predict(data)[0]
    if pred == 1:
        return "Has heart disease"
    else:
        return "Does not have any heart disease"


inputs = [
    gr.Slider(minimum=0, maximum=150, label="BMI"),
    gr.Checkbox(label="Smoking"),
    gr.Checkbox(label="Alcohol Drinking"),
    gr.Checkbox(label="Stroke"),
    gr.Number(label="Physical Health"),
    gr.Number(label="MentalHealth"),
    gr.Checkbox(label="Diff Walking"),
    gr.Dropdown(["Female", "Male"], type="index", label="Sex"),
    gr.Dropdown(['18-24',
            '25-29',
            '30-34',
            '35-39',
            '40-44',
            '45-49',
            '50-54',
            '55-59',
            '60-64',
            '65-69',
            '70-74',
            '75-79',
            '80 or older'], label="Age Category"),
    gr.Dropdown(["White", "Black", "Asian", "American Indian/Alaskan Native", "Hispanic", "Other"], type="index", label="Race"),
    gr.Dropdown(["No", "Yes", "Yes(during pregnancy)", "No, borderline diabetes"], type="index", label="Diebetic"),
    gr.Checkbox(label="Physical Activity"),
    gr.Dropdown(["Poor", "Fair", "Good", "Very good", "Excellent"], type="index", label="General Health"),
    gr.Number(label="Sleep time"),
    gr.Checkbox(label="Asthma"),
    gr.Checkbox(label="Kidney Disease"),
    gr.Checkbox(label="Skin Cancer"),
]

iface = gr.Interface(
    fn=predict,
    inputs=inputs,
    outputs="text",
    description="The purpose of this model is to predict whether or not a person has any heart diseases or not.",
)

iface.launch(share=True)