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

df_framework = pd.read_csv('./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 convert_yes_no(data):
  if data == "Yes":
    return 1
  else:
    return 0


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

    data = [[str(gender), float(age), int(convert_yes_no(hypertension)), int(convert_yes_no(heart_disease)), str(married), str(work_type), str(residence_type), float(avg_glucose_level), float(bmi), str(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 f"O paciente tem a seguinte possibilidade de infarto: {preds[0][1]:.2f}"


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()