import joblib
import pandas as pd
import streamlit as st 

# Load the models
model_acc = joblib.load('model_acc.joblib')  # Model for accommodation cost prediction
model_tp = joblib.load('model_tp.joblib')    # Model for transportation cost prediction

# Load the unique values for accommodation and transportation types
unique_values_acc = joblib.load('unique_values_acc.joblib')
unique_values_tp = joblib.load('unique_values_tp.joblib')

def main():
    st.title("Angie Travel Assistant")

    with st.form("questionnaire"):
        day = st.slider("Duration of Traveling", min_value=1, max_value=60)
        acc_type = st.selectbox("Preferred Accommodation Type", unique_values_acc)
        tp_type = st.selectbox("Preferred Transportation Type", unique_values_tp)

        clicked = st.form_submit_button("Predict total expense")
        if clicked:
            # Prepare input data for both models
            data_acc = pd.DataFrame({
                'duration_(days)': day,
                'accommodation_type_Airbnb': 0,  
                'accommodation_type_Guesthouse': 0,
                'accommodation_type_Hostel': 0,
                'accommodation_type_Hotel': 0,
                'accommodation_type_Resort': 0,
                'accommodation_type_Vacation rental': 0,
                'accommodation_type_Villa': 0
            }, index=[0])
            
            data_tp = pd.DataFrame({
                'duration_(days)': day,
                'transportation_type_Bus': 0,
                'transportation_type_Car': 0,
                'transportation_type_Ferry': 0,
                'transportation_type_Flight': 0,
                'transportation_type_Subway': 0,
                'transportation_type_Train': 0
            }, index=[0])

            # Set the selected transportation type to 1
            data_tp[tp_type] = 1
            data_acc[acc_type] = 1

            # Predict expenses using both models
            result1 = model_acc.predict(data_acc)
            result2 = model_tp.predict(data_tp)

            # Calculate total expense
            total_expense = result1[0] + result2[0]

            st.success(f'The predicted total expense is {total_expense}')

if __name__ == '__main__':
    main()