import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score
import gradio as gr

# Load dataset
df = pd.read_csv("hf://datasets/buio/heart-disease/heart.csv")

# Convert categorical columns to numeric using one-hot encoding
df = pd.get_dummies(df, drop_first=True)

# Define features and target
X = df.drop('target', axis=1)
y = df['target']

# Split data into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Standardize the data
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

# Create and train the logistic regression model
model = LogisticRegression()
model.fit(X_train, y_train)

# Evaluate the model
y_pred = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
print(f'Accuracy: {accuracy:.2f}')

# Function to predict heart disease
def predict_heart_disease(age, sex, cp, trestbps, chol, fbs, restecg, thalach, exang, oldpeak, slope, ca, thal):
    # Create input array
    input_data = np.array([age, sex, cp, trestbps, chol, fbs, restecg, thalach, exang, oldpeak, slope, ca, thal]).reshape(1, -1)
    input_data = scaler.transform(input_data)

    # Make prediction
    prediction = model.predict(input_data)
    return "The person has heart disease." if prediction[0] == 1 else "The person does not have heart disease."

# Gradio integration
def gradio_predict_heart_disease(age, sex, cp, trestbps, chol, fbs, restecg, thalach, exang, oldpeak, slope, ca, thal):
    return predict_heart_disease(age, sex, cp, trestbps, chol, fbs, restecg, thalach, exang, oldpeak, slope, ca, thal)

iface = gr.Interface(
    fn=gradio_predict_heart_disease,
    inputs=[
        gr.components.Number(label="Age"),
        gr.components.Radio(label="Sex", choices=[0, 1]),
        gr.components.Dropdown(label="Chest Pain Type (cp)", choices=[0, 1, 2, 3]),
        gr.components.Number(label="Resting Blood Pressure (trestbps)"),
        gr.components.Number(label="Serum Cholestoral in mg/dl (chol)"),
        gr.components.Radio(label="Fasting Blood Sugar > 120 mg/dl (fbs)", choices=[0, 1]),
        gr.components.Radio(label="Resting Electrocardiographic Results (restecg)", choices=[0, 1]),
        gr.components.Number(label="Maximum Heart Rate Achieved (thalach)"),
        gr.components.Radio(label="Exercise Induced Angina (exang)", choices=[0, 1]),
        gr.components.Number(label="ST depression induced by exercise relative to rest (oldpeak)"),
        gr.components.Dropdown(label="Slope of the peak exercise ST segment (slope)", choices=[0, 1, 2]),
        gr.components.Dropdown(label="Number of major vessels (0-3) colored by fluoroscopy (ca)", choices=[0, 1, 2, 3]),
        gr.components.Dropdown(label="Thalassemia (thal)", choices=[1, 2, 3])
    ],
    outputs="text"
)

iface.launch()