import gradio as gr
import joblib
import numpy as np

# Load the trained model
model = joblib.load("trained_model.pkl")

# Define the prediction function
def predict_heart_disease(
    PhysicalHealthDays: float,
    MentalHealthDays: float,
    SleepHours: float,
    BMI: float,
    PhysicalActivities: str,
    AlcoholDrinkers: str,
    HIVTesting: str,
    RemovedTeeth: str,
    HighRiskLastYear: str,
    CovidPos: str,
):
    try:
        # Encode categorical inputs as integers
        physical_activities = 1 if PhysicalActivities == "Yes" else 0
        alcohol_drinkers = 1 if AlcoholDrinkers == "Yes" else 0
        hiv_testing = 1 if HIVTesting == "Yes" else 0
        removed_teeth = 1 if RemovedTeeth == "Yes" else 0
        high_risk_last_year = 1 if HighRiskLastYear == "Yes" else 0
        covid_pos = 1 if CovidPos == "Yes" else 0

        # Prepare features as a numpy array
        features = np.array([[
            PhysicalHealthDays,
            MentalHealthDays,
            SleepHours,
            BMI,
            physical_activities,
            alcohol_drinkers,
            hiv_testing,
            removed_teeth,
            high_risk_last_year,
            covid_pos,
        ]])

        # Make prediction
        prediction = model.predict(features)
        result = "Yes" if prediction[0] == 1 else "No"
        return f"Heart Disease Prediction: {result}"
    except Exception as e:
        return f"Error in prediction: {str(e)}"

# Define the Gradio interface
inputs = [
    gr.inputs.Number(label="Physical Health Days"),
    gr.inputs.Number(label="Mental Health Days"),
    gr.inputs.Number(label="Sleep Hours"),
    gr.inputs.Number(label="BMI"),
    gr.inputs.Radio(["Yes", "No"], label="Physical Activities"),
    gr.inputs.Radio(["Yes", "No"], label="Alcohol Drinkers"),
    gr.inputs.Radio(["Yes", "No"], label="HIV Testing"),
    gr.inputs.Radio(["Yes", "No"], label="Removed Teeth"),
    gr.inputs.Radio(["Yes", "No"], label="High Risk Last Year"),
    gr.inputs.Radio(["Yes", "No"], label="Covid Positive"),
]

outputs = gr.outputs.Textbox(label="Heart Disease Prediction")

# Create the Gradio app
app = gr.Interface(
    fn=predict_heart_disease,
    inputs=inputs,
    outputs=outputs,
    title="Heart Disease Prediction App",
    description="Please enter health metrics to predict the risk of heart disease."
)

# Launch the app
if __name__ == "__main__":
    app.launch()