import joblib
import pandas as pd
import streamlit as st

Pros= {'Engineer': 1,
               'Healthcare': 2,
               'Executive': 3,
               'Doctor': 4,
               'Artist': 5,
               'Lawyer': 6,
               'Entertainment': 7,
               'Homemaker': 8,
               'Marketing': 9}
model = joblib.load('model.joblib')
unique_values = joblib.load('unique_values.joblib')

def main():
    st.title("Customer Segmentation Prediction")
    with st.form("questionnaire"):
        Gender = st.selectbox("Gender", unique_values["Gender"])
        Ever_Married = st.selectbox("Ever Married", unique_values["Ever_Married"])
        Age = st.slider("Age", min_value=18, max_value=89)
        Graduated = st.selectbox("Graduated", unique_values["Graduated"])
        Profession = st.selectbox("Profession", unique_values["Profession"])
        Work_Experience = st.slider("Work Experience", min_value=0, max_value=14)
        Spending_Score = st.selectbox("Spending Score", unique_values["Spending_Score"])
        Family_Size = st.slider("Family Size", min_value=1, max_value=9)
        Var_1 = st.selectbox("Var_1", unique_values["Var_1"])
        ID = st.slider("ID", min_value=458982, max_value=467974)

        clicked = st.form_submit_button("Predict Segmentation")
        if clicked:
            result = model.predict(pd.DataFrame({"Gender": [Gender],
                                                 "Ever_Married": [Ever_Married],
                                                 "Age": [Age],
                                                 "ID": [ID],
                                                 "Graduated": [Graduated],
                                                 "Profession": [Pros[Profession]],
                                                 "Work_Experience": [Work_Experience],
                                                 "Spending_Score": [Spending_Score],
                                                 "Family_Size": [Family_Size],
                                                 "Var_1": [Var_1]
                                                 }))
            if result[0] == 0:
                result = "A"
            elif result[0] == 1:
                result = "B"
            elif result[0] == 2:
                result = "C"
            else:
                result = "D"
            st.success('Predicted Segmentation: {}'.format(result))

if __name__ == '__main__':
    main()