# Import necessary libraries
import streamlit as st
import joblib

# 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

# Create the Streamlit web application
def main():
    # Set title and description
    st.title("Heart Disease Prediction")
    st.write("This app predicts the presence of heart disease based on selected attributes.")

    # Design user interface
    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"])

    # 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"):
        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.write("The model predicts that the patient has heart disease.")
        else:
            st.write("The model predicts that the patient does not have heart disease.")

# Run the Streamlit app
if __name__ == "__main__":
    main()