import streamlit as st
import joblib

# Set page configuration
st.set_page_config(page_title="Heart Disease Prediction", layout="wide")

# Load the trained model
model = joblib.load("./model-3.joblib")

# Define function to predict heart disease
def predict_heart_disease(sex, exang, cp_1, cp_2, cp_4, slope_1, slope_2, thal_3, thal_7):
    print([[sex, exang, cp_1, cp_2, cp_4, slope_1, slope_2, thal_3, thal_7]])
    prediction = model.predict([[sex, exang, cp_1, cp_2, cp_4, slope_1, slope_2, thal_3, thal_7]])
    return prediction

def run():
    st.title("Heart Disease Prediction")
    st.write("Please provide the following information to predict heart disease:")

    # Design user interface
    col1, col2 = st.columns([2, 1])

    with col1:
        sex = st.selectbox("Sex", ["Female", "Male"])
        exang = st.selectbox("Exercise Induced Angina", ["No", "Yes"])
        cp = st.selectbox("Chest Pain Type", ["Typical Angina", "Atypical Angina", "Non-Anginal Pain", "Asymptomatic"])
        slope = st.selectbox("Slope of Peak Exercise ST Segment", ["Upsloping", "Flat", "Downsloping"])
        thal = st.selectbox("Thal", ["Normal", "Fixed Defect", "Reversible Defect"])
    
    with col2:
        st.image("https://images.unsplash.com/photo-1618939304347-e91b1f33d2ab?q=80&w=1974&auto=format&fit=crop&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D", width=275)

    # Map selected options to numerical values
    sex_mapping = {"Female": 0, "Male": 1}
    exang_mapping = {"No": 0, "Yes": 1}
    cp_1_mapping = {"Typical Angina": 1, "Atypical Angina": 0, "Non-Anginal Pain": 0, "Asymptomatic": 0}
    cp_2_mapping = {"Typical Angina": 0, "Atypical Angina": 1, "Non-Anginal Pain": 0, "Asymptomatic": 0}
    cp_4_mapping = {"Typical Angina": 0, "Atypical Angina": 0, "Non-Anginal Pain": 0, "Asymptomatic": 1}
    slope_1_mapping = {"Upsloping": 1, "Flat": 0, "Downsloping": 0}
    slope_2_mapping = {"Upsloping": 0, "Flat": 1, "Downsloping": 0}
    thal_3_mapping = {"Normal": 1, "Fixed Defect": 0, "Reversible Defect": 0}
    thal_7_mapping = {"Normal": 0, "Fixed Defect": 0, "Reversible Defect": 1}

    # Predict button
    if st.button("Predict", key="predict_button"):
        result = predict_heart_disease(sex_mapping[sex], exang_mapping[exang], cp_1_mapping[cp], cp_2_mapping[cp], cp_4_mapping[cp], slope_1_mapping[slope], slope_2_mapping[slope], thal_3_mapping[thal], thal_7_mapping[thal])
        if result == 1:
            st.error("The model predicts that the patient has heart disease.")
        else:
            st.success("The model predicts that the patient does not have heart disease.")

    # Add a footer
    st.markdown("---")
    st.write("Made with ❤️ by Viga, Hanum, & Robit")

if __name__ == "__main__":
    run()