import gradio as gr
from fastai.tabular.all import *
import pandas as pd
import torch

df_framework = pd.read_csv('/kaggle/input/stroke-prediction-dataset/healthcare-dataset-stroke-data.csv')
splits = RandomSplitter(seed=42)(df_framework)

dls = TabularPandas(
    df_framework, splits=splits,
    procs = [Categorify, FillMissing, Normalize],
    cat_names=["gender","ever_married","work_type","Residence_type", "smoking_status"],
    cont_names=['age', 'hypertension', 'heart_disease', 'avg_glucose_level', 'bmi'],
    y_names="stroke", y_block = CategoryBlock(),
).dataloaders(path=".")

learn = tabular_learner(dls, metrics=accuracy, layers=[10,10])
learn.fit(16, lr=0.025)

def predict(age, hypertension, heart_disease, avg_glucose_level, bmi, gender, married, work_type, residence_type, smoking_status):

    data = [[gender, age, hypertension, heart_disease, married, work_type, residence_type, avg_glucose_level, bmi, smoking_status]]
    columns_df = ['gender', 'age', 'hypertension', 'heart_disease', 'ever_married', 'work_type', 'Residence_type', 'avg_glucose_level', 'bmi','smoking_status']
    df_row = pd.DataFrame(data,columns=columns_df)

    dl = learn.dls.test_dl(df_row)
    preds,_ = learn.get_preds(dl=dl)
    
    return "O paciente tem a seguinte possibilidade de infarto: " + str(preds[0])


gr.Interface(
    fn=predict,
    title="Stroke predict Model",
    allow_flagging="never",
    share=True,
    inputs=[
        gr.inputs.Number(default=30, label="Age"),
        gr.Dropdown(["Yes", "No"], label="Hypertension"),
        gr.Dropdown(["Yes", "No"], label="Heart disease"),
        gr.inputs.Number(default=100, label="Average glucose level"),
        gr.inputs.Number(default=28.8, label="Body Mass Index (BMI)"),
        gr.inputs.Radio(choices=["Female", "Male", "Other"], default="Female", label="Gender"),
        gr.Dropdown(["No", "Yes"], label="Married"),
        gr.Dropdown(["Governamental", "Never worked", "Private", "Self-employed", "Children"], label="Work type"),
        gr.Dropdown(["Rural", "Urban"],  label="Residence type"),
        gr.Dropdown(["Unknown", "Formely smoked", "Never smoked", "Smokes"], label="Smoking status"),
    ],
    outputs="text").launch()