import gradio as gr
import joblib
import pandas as pd
import json

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

# Load feature names used during training
with open("feature_names.json", "r") as f:
    feature_names = json.load(f)

# Define the prediction function
def predict_heart_failure(age_category, alcohol_drinkers, chest_scan, covid_positive, physical_health_days, mental_health_days, sleep_hours, bmi):
    # Initialize all features to 0
    input_data = {feature: 0 for feature in feature_names}

    # Handle Age Category
    input_data[f"AgeCategory_{age_category}"] = 1  # Set the selected age category to 1

    # Handle categorical inputs
    input_data[f"AlcoholDrinkers_{alcohol_drinkers}"] = 1  # Alcohol Drinkers
    input_data[f"ChestScan_{chest_scan}"] = 1  # Chest Scan
    input_data[f"CovidPos_{covid_positive}"] = 1  # Covid Positive

    # Handle numeric inputs
    input_data["PhysicalHealthDays"] = physical_health_days
    input_data["MentalHealthDays"] = mental_health_days
    input_data["SleepHours"] = sleep_hours
    input_data["BMI"] = bmi

    # Create a DataFrame with the required feature names
    input_df = pd.DataFrame([input_data])

    # Ensure all required features are present
    for feature in feature_names:
        if feature not in input_df.columns:
            input_df[feature] = 0  # Fill missing features with default value (e.g., 0)

    print("Inputs received in the app:")
    print(input_df)
    
    # Ensure no extra features are included
    input_df = input_df[feature_names]

    # Make a prediction
    prediction = model.predict(input_df)[0]
    print(f"Prediction: {prediction}")
    return "High likelihood of heart failure" if prediction == 1 else "Low likelihood of heart failure"



# Define Gradio inputs
gradio_inputs = [
    gr.Radio(
        ["18-24", "25-34", "35-44", "45-54", "55-64", "65-74", "75-79", "80 or older"],
        label="Age Category"
    ),
    gr.Radio(["Yes", "No"], label="Alcohol Drinkers"),
    gr.Radio(["Yes", "No"], label="Chest Scan"),
    gr.Radio(["Yes", "No"], label="COVID Positive"),
    gr.Slider(0, 30, step=1, label="Physical Health Days"),
    gr.Slider(0, 30, step=1, label="Mental Health Days"),
    gr.Slider(0, 12, step=0.5, label="Sleep Hours"),
    gr.Slider(10, 50, step=0.1, label="BMI")
]

# Set up the Gradio interface
interface = gr.Interface(
    fn=predict_heart_failure,
    inputs=gradio_inputs,
    outputs="text",
    title="Heart Failure Risk Prediction",
    description="This app predicts the likelihood of heart failure based on health and lifestyle inputs."
)

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