import streamlit as st
import page_home
import page_attitudes
import page_demographics
import page_shopping
import page_investing
import page_environment
import page_ai
import page_personas
import page_tests
from urllib.parse import quote, unquote
from datasets import load_dataset

# Default to wide mode
st.set_page_config(layout="wide")

# Initialize a session state variable for page navigation
if 'page' not in st.session_state:
    # Check if a page is specified in the query parameters
    query_params = st.experimental_get_query_params()
    initial_page = unquote(query_params.get("page", ["Home"])[0])
    st.session_state['page'] = initial_page


@st.cache_data
def load_data():
    # Load data from Huggingface
    dataset = load_dataset(
        "krishaamer/taiwanese-college-students", data_files={'train': 'clean.csv'})
    return dataset['train'].to_pandas()


# Sidebar navigation using buttons
st.sidebar.title("Taiwanese College Students")
if st.sidebar.button("Introduction"):
    st.session_state['page'] = 'Home'
if st.sidebar.button("Demographics"):
    st.session_state['page'] = 'Demographics'
if st.sidebar.button("Attitudes"):
    st.session_state['page'] = 'Attitudes'
if st.sidebar.button("Personas ⭐"):
    st.session_state['page'] = 'Personas'
if st.sidebar.button("Shopping"):
    st.session_state['page'] = 'Shopping'
if st.sidebar.button("Investing"):
    st.session_state['page'] = 'Investing'
if st.sidebar.button("Environment"):
    st.session_state['page'] = 'Environment'
if st.sidebar.button("AI Companion"):
    st.session_state['page'] = 'AI'
if st.sidebar.button("Tests"):
    st.session_state['page'] = 'Tests'

# Update the query parameters based on the session state
st.experimental_set_query_params(page=quote(st.session_state['page']))

# Page contents based on session state
df = load_data()

if st.session_state['page'] == 'Home':
    page_home.show()
elif st.session_state['page'] == 'Attitudes':
    page_attitudes.show(df)
elif st.session_state['page'] == 'Demographics':
    page_demographics.show(df)
elif st.session_state['page'] == 'Personas':
    page_personas.show(df)
elif st.session_state['page'] == 'Shopping':
    page_shopping.show(df)
elif st.session_state['page'] == 'Investing':
    page_investing.show(df)
elif st.session_state['page'] == 'Environment':
    page_environment.show(df)
elif st.session_state['page'] == 'AI':
    page_ai.show(df)
elif st.session_state['page'] == 'Tests':
    page_tests.show(df)

st.sidebar.divider()

# Rerun Calculations
if st.sidebar.button("Rerun Calculations"):
    st.rerun()