import pickle
import streamlit as st
from streamlit_option_menu import option_menu

# ================== PAGE CONFIG ==================
st.set_page_config(
    page_title="Smart Health AI",
    page_icon="🧠",
    layout="wide",
)

# ================== LOAD MODELS ==================
def load_model(file):
    try:
        return pickle.load(open(file, "rb"))
    except:
        st.error(f"❌ Could not load model: {file}")
        return None

diabetes_model   = load_model("src/models/diabetes_model.sav")
heart_model      = load_model("src/models/heart_model.sav")
parkinson_model  = load_model("src/models/parkinsons_model.sav")

# ================== SIDEBAR MENU ==================
with st.sidebar:
    st.title("🩺 Smart Health AI")
    selected = option_menu(
        "Main Menu",
        ["Diabetes Prediction", "Heart Disease Prediction", "Parkinson's Prediction"],
        icons=["activity", "heart", "person"],
        default_index=0
    )

# ================== INPUT FUNCTIONS ==================

# --------- Diabetes Inputs ---------
def get_diabetes_input():
    col1, col2, col3 = st.columns(3)
    Pregnancies = col1.number_input('Pregnancies', 0)
    Glucose = col2.number_input('Glucose Level', 0.0)
    BloodPressure = col3.number_input('Blood Pressure', 0.0)
    SkinThickness = col1.number_input('Skin Thickness', 0.0)
    Insulin = col2.number_input('Insulin Level', 0.0)
    BMI = col3.number_input('BMI', 0.0)
    DPF = col1.number_input('Diabetes Pedigree Function', 0.0)
    Age = col2.number_input('Age', 0)
    return [Pregnancies, Glucose, BloodPressure, SkinThickness, Insulin, BMI, DPF, Age]


# --------- Heart Inputs ---------
def get_heart_input():
    cols = st.columns(4)
    age      = cols[0].number_input("Age", 0)
    sex      = cols[1].selectbox("Sex", [0,1])
    cp       = cols[2].selectbox("Chest Pain Type (0–3)", [0,1,2,3])
    trestbps = cols[3].number_input("Resting Blood Pressure", 0)

    chol     = cols[0].number_input("Cholesterol", 0)
    fbs      = cols[1].selectbox("Fasting Sugar > 120 mg/dl", [0,1])
    restecg  = cols[2].selectbox("Rest ECG", [0,1,2])
    thalach  = cols[3].number_input("Max Heart Rate", 0)

    exang    = cols[0].selectbox("Exercise Angina", [0,1])
    oldpeak  = cols[1].number_input("ST Depression", 0.0)
    slope    = cols[2].selectbox("Slope", [0,1,2])
    ca       = cols[3].selectbox("Number of Major Vessels (0–4)", [0,1,2,3,4])
    thal     = cols[0].selectbox("Thal", [0,1,2,3])

    return [age, sex, cp, trestbps, chol, fbs, restecg,
            thalach, exang, oldpeak, slope, ca, thal]


# --------- Parkinson Inputs ---------
def get_parkinson_input():
    params = [
        "MDVP:Fo(Hz)", "MDVP:Fhi(Hz)", "MDVP:Flo(Hz)", "MDVP:Jitter(%)",
        "MDVP:Jitter(Abs)", "MDVP:RAP", "MDVP:PPQ", "Jitter:DDP",
        "MDVP:Shimmer", "MDVP:Shimmer(dB)", "Shimmer:APQ3",
        "Shimmer:APQ5", "MDVP:APQ", "Shimmer:DDA", "NHR", "HNR",
        "RPDE", "DFA", "spread1", "spread2", "D2", "PPE"
    ]

    cols = st.columns(3)
    values = []
    for i, p in enumerate(params):
        values.append(cols[i % 3].number_input(p, value=0.0))
    return values


# ================== PREDICT FUNCTION ==================
def predict(model, data):
    return model.predict([data])[0]


# ================== MAIN APP ==================
if selected == "Diabetes Prediction":
    st.header("🩸 Diabetes Prediction")
    data = get_diabetes_input()
    if st.button("Predict"):
        result = predict(diabetes_model, data)
        if result == 1:
            st.error("⚠️ The person is diabetic.")
        else:
            st.success("✅ The person is not diabetic.")

elif selected == "Heart Disease Prediction":
    st.header("❤️ Heart Disease Prediction")
    data = get_heart_input()
    if st.button("Predict Heart Disease"):
        result = predict(heart_model, data)
        if result == 1:
            st.error("⚠️ The patient is at risk of heart disease.")
        else:
            st.success("✅ The heart is healthy.")

elif selected == "Parkinson's Prediction":
    st.header("🧠 Parkinson's Detection")
    data = get_parkinson_input()
    if st.button("Predict Parkinson's"):
        result = predict(parkinson_model, data)
        if result == 1:
            st.error("⚠️ Parkinson's Detected.")
        else:
            st.success("✅ No Parkinson's signs detected.")

st.markdown("---")
st.markdown("Developed by Eng Youssef Samy | AI Smart Health")