import gradio as gr
import pickle
import pandas as pd

# Load model (make sure Heart_Project.pkl is in repo root)
with open("Heart_Project.pkl", "rb") as f:
    model = pickle.load(f)

# Prediction function
def predict(Age, Sex, ChestPainType, RestingBP, Cholesterol,
            FastingBS, RestingECG, MaxHR, ExerciseAngina, Oldpeak, ST_Slope):
    SEX_MAP    = {'M': 1, 'F': 0}
    CP_MAP     = {'TA': 0, 'ATA': 1, 'NAP': 2, 'ASY': 3}
    ECG_MAP    = {'Normal': 0, 'ST': 1, 'LVH': 2}
    ANGINA_MAP = {'Y': 1, 'N': 0}
    SLOPE_MAP  = {'Up': 2, 'Flat': 1, 'Down': 0}

    # Build feature row (same order as used during training)
    features = [[
        Age,
        SEX_MAP[Sex],
        CP_MAP[ChestPainType],
        RestingBP,
        Cholesterol,
        FastingBS,
        ECG_MAP[RestingECG],
        MaxHR,
        ANGINA_MAP[ExerciseAngina],
        Oldpeak,
        SLOPE_MAP[ST_Slope]
    ]]

    df = pd.DataFrame(features, columns=[
        'Age','Sex','ChestPainType','RestingBP','Cholesterol',
        'FastingBS','RestingECG','MaxHR','ExerciseAngina',
        'Oldpeak','ST_Slope'
    ])

    # Use predict_proba if available, otherwise fallback to predict
    if hasattr(model, "predict_proba"):
        prob = model.predict_proba(df)[0, 1]
    else:
        # fallback: model might not have predict_proba
        prob = None

    pred = model.predict(df)[0]
    if prob is None:
        return f"Prediction: {'Heart Disease' if pred==1 else 'No Heart Disease'}"
    return f"Prediction: {'Heart Disease' if pred==1 else 'No Heart Disease'} | Probability: {prob:.2f}"

# Build Gradio UI
demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Number(label="Age"),
        gr.Radio(["M","F"], label="Sex"),
        gr.Radio(["TA","ATA","NAP","ASY"], label="Chest Pain Type"),
        gr.Number(label="RestingBP"),
        gr.Number(label="Cholesterol"),
        gr.Radio([0,1], label="FastingBS"),
        gr.Radio(["Normal","ST","LVH"], label="RestingECG"),
        gr.Number(label="MaxHR"),
        gr.Radio(["Y","N"], label="ExerciseAngina"),
        gr.Number(label="Oldpeak"),
        gr.Radio(["Up","Flat","Down"], label="ST Slope")
    ],
    outputs="text",
    title="❤️ Heart Disease Predictor",
    description="Enter patient details to predict the risk of heart disease."
)

if __name__ == "__main__":
    demo.launch()