import streamlit as st
from PIL import Image
from chest_utils import predict_chest, condition_questions
from heart_utils import predict_heart

st.set_page_config(page_title="🩺 Disease Prediction App", layout="centered")

# Initialize session state
if "page" not in st.session_state:
    st.session_state.page = "heart"

st.title("🩺 Disease Prediction App")
st.markdown("Follow the steps to assess your Heart, Chest, and Combined Health Risk.")

# Custom styled result display
def styled_result(label, result, icon, disease_detected):
    if disease_detected:
        bg = "#f8d7da"  # light red
        color = "#a94442"
    else:
        bg = "#d4edda"  # light green
        color = "#155724"
    st.markdown(
        f'''
        <div style="background-color:{bg}; padding:10px; border-radius:5px; color:{color};">
            {icon} <strong>{label}: {result}</strong>
        </div>
        ''',
        unsafe_allow_html=True
    )

# ───────────── Step 1: Heart Disease Prediction ───────────── #
if st.session_state.page == "heart":
    # Common Inputs
    st.session_state.age = st.slider("Age", 1, 100, 30)
    gender = st.selectbox("Gender", ["Male", "Female"])
    st.session_state.gender = gender
    st.session_state.sex = 1 if gender == "Male" else 0
    
    st.header("❤️ Step 1: Heart Disease Prediction")

    # Heart-related Inputs
    st.session_state.cp = st.selectbox("Chest Pain Type", [0, 1, 2, 3])
    st.session_state.trestbps = st.number_input("Resting Blood Pressure", min_value=80, max_value=200, value=120)
    st.session_state.chol = st.number_input("Serum Cholesterol (mg/dl)", min_value=100, max_value=600, value=200)
    st.session_state.fbs = st.selectbox("Fasting Blood Sugar > 120 mg/dl", [1, 0])
    st.session_state.restecg = st.selectbox("Resting ECG", [0, 1, 2])
    st.session_state.thalach = st.number_input("Max Heart Rate Achieved", min_value=60, max_value=220, value=150)
    st.session_state.exang = st.selectbox("Exercise Induced Angina", [1, 0])
    st.session_state.oldpeak = st.number_input("ST Depression", min_value=0.0, max_value=6.0, value=1.0)
    st.session_state.slope = st.selectbox("Slope of ST Segment", [0, 1, 2])
    st.session_state.ca = st.selectbox("Number of Major Vessels", [0, 1, 2, 3])
    st.session_state.thal = st.selectbox("Thalassemia", [1, 2, 3])

    if st.button("Next → Chest Symptoms"):
        st.session_state.heart_result = predict_heart(
            st.session_state.age, st.session_state.sex, st.session_state.cp, st.session_state.trestbps,
            st.session_state.chol, st.session_state.fbs, st.session_state.restecg, st.session_state.thalach,
            st.session_state.exang, st.session_state.oldpeak, st.session_state.slope,
            st.session_state.ca, st.session_state.thal
        )
        
        # Show heart prediction immediately before proceeding
        styled_result("Heart Disease Prediction", st.session_state.heart_result, "❤️", st.session_state.heart_result == "Heart Disease Detected")

        st.session_state.page = "chest"
        st.rerun()

# ───────────── Step 2: Chest Disease Symptoms ───────────── #
elif st.session_state.page == "chest":
    st.header("🫁 Step 2: Chest Disease Symptoms")

    st.markdown(f"ℹ️ Using Age: **{st.session_state.age}**, Gender: **{st.session_state.gender}**")

    st.session_state.position = st.selectbox("View Position", ["PA", "AP"])

    uploaded_file = st.file_uploader("Upload Chest X-ray", type=["jpg", "jpeg", "png"])
    if uploaded_file:
        st.session_state.image = Image.open(uploaded_file)
        st.image(st.session_state.image, caption="Uploaded Chest X-ray", use_container_width=True)

    st.subheader("Answer the following clinical questions:")
    st.session_state.conditions = []
    for question in condition_questions:
        st.session_state.conditions.append(1 if st.checkbox(question) else 0)

    if st.button("Next → Combined Prediction"):
        if uploaded_file:
            st.session_state.page = "combined"
            st.rerun()
        else:
            st.warning("⚠️ Please upload a chest X-ray to continue.")

# ───────────── Step 3: Combined Prediction Results ───────────── #
elif st.session_state.page == "combined":
    st.header("🔁 Step 3: Combined Prediction Results")

    chest_result = predict_chest(
        st.session_state.age,
        st.session_state.gender,
        st.session_state.position,
        st.session_state.conditions
    )
    heart_result = st.session_state.heart_result

    st.image(st.session_state.image, caption="Uploaded Chest X-ray", use_container_width=True)

    styled_result("Chest Disease Prediction", chest_result, "🫁", chest_result == "Chest Disease Detected")
    styled_result("Heart Disease Prediction", heart_result, "❤️", heart_result == "Heart Disease Detected")

    # Combined risk display
    if "Detected" in chest_result or "Detected" in heart_result:
        st.markdown(
            '''
            <div style="background-color:#f8d7da; padding:10px; border-radius:5px; color:#a94442;">
                ⚠️ <strong>Combined Risk: High</strong> — Immediate attention recommended.
            </div>
            ''',
            unsafe_allow_html=True
        )
    else:
        st.markdown(
            '''
            <div style="background-color:#d4edda; padding:10px; border-radius:5px; color:#155724;">
                ✅ <strong>Combined Risk: Low</strong> — You seem to be in good health.4
            </div>
            ''',
            unsafe_allow_html=True
        )

    if st.button("🔁 Start Over"):
        st.session_state.clear()
        st.rerun()