import gradio as gr
import pickle
import numpy as np
from sklearn.preprocessing import StandardScaler

# Load the trained models
models = {
    "Logistic": pickle.load(open("logistic_model_best.pkl", "rb")),
    "SVC": pickle.load(open("svc_model_best.pkl", "rb")),
    "Decision Trees": pickle.load(open("decisiontree_model_best.pkl", "rb")),
    "Random Forest": pickle.load(open("rfc_model_best.pkl", "rb")),
    "KNN": pickle.load(open("knn_model_best.pkl", "rb")),
    "LDA": pickle.load(open("lda_model_best.pkl", "rb"))
}

# Define the input and output dropdowns
age_dropdown = gr.inputs.Number(label='Age')
sex_dropdown = gr.inputs.Dropdown(["Male", "Female"], label="Sex")
cp_dropdown = gr.inputs.Dropdown(["1", "2", "3", "4"], label="Chest Pain Type")
restbp_dropdown = gr.inputs.Number(label='Resting Blood pressure')
chol_dropdown = gr.inputs.Number(label='Serum cholestoral')
fbs_dropdown = gr.inputs.Dropdown(["0", "1"], label="Fasting Blood Sugar")
ecg_dropdown = gr.inputs.Dropdown(["0", "1", "2"], label="Electrocardiographic Results")
max_rt_dropdown = gr.inputs.Number(label="Maximum Heart Rate")
exang_dropdown = gr.inputs.Dropdown(["0", "1"], label="Exercise-Induced Angina")
oldpeak_dropdown = gr.inputs.Number(label="oldpeak =ST Depression Induced by Exercise")
slope_dropdown = gr.inputs.Dropdown(["0", "1", "2"], label="Slope of the ST Segment")
ca_dropdown = gr.inputs.Dropdown(["0", "1", "2", "3"], label="Number of Major Vessels")
thal_dropdown = gr.inputs.Dropdown(["0", "1", "2"], label="Thalassemia")
model_dropdown = gr.inputs.Dropdown(
    ["Logistic", "SVC", "Decision Trees", "Random Forest", "KNN", "LDA"],
    label="Model"
)
xoutputs = [
    gr.outputs.Textbox(label="Probability of Heart Disease"),
]
# Define the function to make predictions
import numpy as np

def predict(age, sex, cp, restbp, chol, fbs, ecg, max_rt, exang, oldpeak, slope, ca, thal, model):
    # Convert input values to their respective data types
    age = int(age)
    sex = 1 if sex == "Male" else 0
    cp = int(cp)
    restbp = int(restbp)
    chol = int(chol)
    fbs = int(fbs)
    ecg = int(ecg)
    max_rt = int(max_rt)
    exang = int(exang)
    oldpeak = float(oldpeak)
    slope = int(slope)
    ca = int(ca)
    thal = int(thal)

    # Create the input array
    input_array = np.array([age, sex, cp, restbp, chol, fbs, ecg, max_rt, exang, oldpeak, slope, ca, thal]).reshape(1,-1)

    # Get the selected model
    selected_model = models[model]

    scaler = StandardScaler()
    normalized_input_array = scaler.fit_transform(input_array)

    predi = selected_model.predict(normalized_input_array)
    # Convert the prediction to a string
    if predi[0] == 1:
        return "Absent"
    elif predi[0] == 2:
        return "Present"
    else:
        return "Unknown"

# Create the interface
interface = gr.Interface(
    fn=predict,
    inputs=[
        age_dropdown, sex_dropdown, cp_dropdown, restbp_dropdown, chol_dropdown,
        fbs_dropdown, ecg_dropdown,max_rt_dropdown, exang_dropdown, oldpeak_dropdown,
slope_dropdown, ca_dropdown, thal_dropdown, model_dropdown
],
  outputs = xoutputs,
examples=[[70, 1.0,4.0,130.0,322.0,0.0,2.0,109.0,0.0,2.4,2.0,3.0,2.0,"SVC"]],

title="Heart Disease Prediction",
description="ML HomeWork 04"
)

interface.launch()