import gradio as gr
import joblib
import numpy as np
import pandas as pd

# Load the trained model
model_path = "trained_model.pkl"

def load_model(path):
    try:
        return joblib.load(path)
    except Exception as e:
        raise ValueError(f"Error loading model: {e}")

# 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:
        model = load_model(model_path)
        # Encode categorical inputs as integers
        physical_activities = 1 if PhysicalActivities.lower() == "yes" else 0
        alcohol_drinkers = 1 if AlcoholDrinkers.lower() == "yes" else 0
        hiv_testing = 1 if HIVTesting.lower() == "yes" else 0
        removed_teeth = 1 if RemovedTeeth.lower() == "yes" else 0
        high_risk_last_year = 1 if HighRiskLastYear.lower() == "yes" else 0
        covid_pos = 1 if CovidPos.lower() == "yes" else 0

        # Combine inputs into a numpy array for prediction
        features = np.array([
            PhysicalHealthDays, MentalHealthDays, SleepHours, BMI,
            physical_activities, alcohol_drinkers, hiv_testing,
            removed_teeth, high_risk_last_year, covid_pos
        ]).reshape(1, -1)

        # Predict with the model
        prediction = model.predict(features)
        return "Heart Disease Risk" if prediction[0] == 1 else "No Risk"
    except Exception as e:
        return f"Error during prediction: {e}"

# Define the Gradio interface
with gr.Blocks() as app:
    gr.Markdown("# Heart Disease Prediction App")
    gr.Markdown("### Provide input values and receive a prediction.")
    
    with gr.Row():
        PhysicalHealthDays = gr.Slider(0, 30, label="Physical Health Days")
        MentalHealthDays = gr.Slider(0, 30, label="Mental Health Days")
        SleepHours = gr.Slider(0, 24, label="Average Sleep Hours")
        BMI = gr.Slider(10, 50, label="Body Mass Index (BMI)")
    
    with gr.Row():
        PhysicalActivities = gr.Radio(["Yes", "No"], label="Engaged in Physical Activities?")
        AlcoholDrinkers = gr.Radio(["Yes", "No"], label="Consumes Alcohol?")
        HIVTesting = gr.Radio(["Yes", "No"], label="Tested for HIV?")
        RemovedTeeth = gr.Radio(["Yes", "No"], label="Has Removed Teeth?")
        HighRiskLastYear = gr.Radio(["Yes", "No"], label="High Risk Last Year?")
        CovidPos = gr.Radio(["Yes", "No"], label="Tested Positive for COVID-19?")

    predict_button = gr.Button("Predict")
    output = gr.Textbox(label="Prediction Result")

    # Connect prediction logic
    predict_button.click(
        fn=predict_heart_disease,
        inputs=[
            PhysicalHealthDays, MentalHealthDays, SleepHours, BMI,
            PhysicalActivities, AlcoholDrinkers, HIVTesting,
            RemovedTeeth, HighRiskLastYear, CovidPos
        ],
        outputs=output,
    )

# Launch the app
app.launch()