import gradio as gr
import pandas as pd
import numpy as np
import pickle
from joblib import load


trees = load('randomFlorest.joblib')

def convert(options, input):
  for i in range(0, len(options)):
    if options[i] == input:
        return i
      

def predict(age, sex, chestPainType, restingBP, cholesterol, fastingBS, restingECG, maxHR, exerciseAngina, oldpeak, stSlope):
    
    convertSex = convert(["F", "M"], sex)
    convertCPT = convert(['ASY', 'ATA', 'NAP', 'TA'], chestPainType)
    convertRECG = convert(['LVH', 'Normal', 'ST'], restingECG)
    convertEA = convert(["N", "Y"],exerciseAngina)
    convertSTS = convert(['Down', 'Flat', 'Up'], stSlope)
    
    data = [[int(convertSex), int(convertCPT), int(convertRECG), int(convertEA), int(convertSTS), int(age), int(restingBP), int(cholesterol), int(fastingBS), int(maxHR), np.float32(oldpeak)]]
    columns_df = columns_df = ['Sex', 'ChestPainType', 'RestingECG', 'ExerciseAngina', 'ST_Slope', 'Age', 'RestingBP', 'Cholesterol', 'FastingBS', 'MaxHR', 'Oldpeak']
    df_row = pd.DataFrame(data,columns=columns_df)

    predictions = [tree.predict(df_row) for tree in trees] 
    avg_probs = np.stack(predictions).mean(0)
    
    return f"A probabilidade do paciente ter doença cardíca é de: {avg_probs[0]*100:.2f}%"
    

example = [
    [40, 'M', 'ATA', 140, 289, 0, 'Normal', 172, 'N', 0.0, 'Up'],
    [49, 'F', 'NAP', 160, 180, 0, 'Normal', 156, 'N', 1.0, 'Flat'],
    [50, 'M', 'ASY', 140, 341, 0, 'ST', 125, 'Y', 1.8, 'Flat'],
    [53, 'M', 'NAP', 145, 518, 0, 'Normal', 125, 'N', 0.0, 'Flat']
]

gr.Interface(
    fn=predict,
    title="Heart Failure predict Model",
    allow_flagging="never",
    share=True,
    inputs=[
        gr.inputs.Number(default=30, label="Age"),
        gr.inputs.Radio(choices=["F", "M"], default="M", label="Sex"),
        gr.Dropdown(['ASY', 'ATA', 'NAP', 'TA'], label="Chest Pain Type", default="ASY"),
        gr.inputs.Number(default=120, label="Resting Blood Pressure"),
        gr.inputs.Number(default=0, label="Cholesterol"),
        gr.Dropdown([0, 1], default=0, label="Fasting Blood Sugar"),
        gr.Dropdown(['LVH', 'Normal', 'ST'], default="Normal", label="Resting ECG"),
        gr.inputs.Number(default=150, label="maximum heart rate"),
        gr.inputs.Radio(choices=["N", "Y"], default="No", label="ExerciseAngina"),
        gr.inputs.Number(default=0.0, label="Old peak"),
        gr.Dropdown(['Down', 'Flat', 'Up'],default="Flat", label="ST Slope"),
    ],
    examples=example,
    outputs="text").launch()