import gradio as gr
import pandas as pd
import pickle  

# Load model RF using pickle
with open('random_forest_model.joblib', 'rb') as f:  
    best_rf_model = pickle.load(f)

def predict_heart_disease(Age, Sex, ChestPainType, RestingBP, Cholesterol, FastingBS, RestingECG, MaxHR, ExerciseAngina, Oldpeak, ST_Slope):
    # Map categorical inputs to numeric values
    sex_mapping = {"Female": 0, "Male": 1}
    chest_pain_mapping = {"TA": 0, "ATA": 1, "NAP": 2, "ASY": 3}
    fasting_bs_mapping = {"< 120 mg/dl": 0, "> 120 mg/dl": 1}
    resting_ecg_mapping = {"Normal": 0, "ST": 1, "LVH": 2}
    exercise_angina_mapping = {"No": 0, "Yes": 1}
    st_slope_mapping = {"Up": 0, "Flat": 1, "Down": 2}

    # Adapt input data to model format
    input_data = {
        'Age': int(Age),
        'Sex': sex_mapping[Sex],
        'ChestPainType': chest_pain_mapping[ChestPainType],
        'RestingBP': int(RestingBP),
        'Cholesterol': int(Cholesterol),
        'FastingBS': fasting_bs_mapping[FastingBS],
        'RestingECG': resting_ecg_mapping[RestingECG],
        'MaxHR': int(MaxHR),
        'ExerciseAngina': exercise_angina_mapping[ExerciseAngina],
        'Oldpeak': float(Oldpeak),
        'ST_Slope': st_slope_mapping[ST_Slope]
    }

    # Input data to input dataframe
    input_df = pd.DataFrame([input_data])

    # Run prediction
    prediction = best_rf_model.predict(input_df)

    # Return output with doctor's advice
    prediction_message = (
        "Based on your inputs, there is no evidence of heart disease. "
        "It's important to maintain a healthy lifestyle and consider regular check-ups."
    ) if prediction[0] == 0 else (
        "Based on your inputs, there may be possible evidence of heart disease. "
        "Please consult a healthcare professional for a thorough evaluation and further guidance."
    )

    # Return both the prediction message and image
    return prediction_message, image_url

# Image URL
image_url = "heart.png"

# Set input and outputs of interface
iface = gr.Interface(
    fn=predict_heart_disease,
    inputs=[
        gr.Number(label="Age", value=40),
        gr.Radio(label="Sex", choices=["Female", "Male"], type="value", value="Male"),
        gr.Radio(label="Chest Pain Type", choices=["TA", "ATA", "NAP", "ASY"], type="value", value="ATA"),
        gr.Number(label="Resting BP (Blood Pressure)", value=140),
        gr.Number(label="Cholesterol", value=289),
        gr.Radio(label="Fasting BS", choices=["< 120 mg/dl", "> 120 mg/dl"], type="value", value="< 120 mg/dl"),
        gr.Radio(label="Resting ECG", choices=["Normal", "ST", "LVH"], type="value", value="ST"),
        gr.Number(label="Max HR (Maximum Heart Rate)", value=172),
        gr.Radio(label="Exercise Angina", choices=["No", "Yes"], type="value", value="No"),
        gr.Number(label="Oldpeak: Numeric value measured in depression", value=0),
        gr.Radio(label="ST Slope", choices=["Up", "Flat", "Down"], type="value", value="Down")
    ],
    outputs=[gr.Text(label="Prediction"), gr.Image(value=image_url, label="Image", show_label=True)],
    title="Heart Disease Prediction",
    description="Please enter the following details:",
    theme=gr.themes.Soft()
)

# Launch Gradio
iface.launch(debug=True)