import streamlit as st
import numpy as np
import pandas as pd
from streamlit_extras.switch_page_button import switch_page
from pages.Gallery import load_hf_dataset

if __name__ == "__main__":
    st.set_page_config(page_title="Personal Image Ranking", page_icon="🎖️️", layout="wide")

    if 'user_id' not in st.session_state:
        st.warning('Please log in first.')
        home_btn = st.button('Go to Home Page')
        if home_btn:
            switch_page("home")

    else:
        selected_modelVersions = []
        for key, value in st.session_state.selected_dict.items():
            for v in value:
                if v not in selected_modelVersions:
                    selected_modelVersions.append(v)

        if len(selected_modelVersions) == 0:
            st.info('You have not checked any image yet. Please go back to the gallery page and check some images.')
            gallery_btn = st.button('Go to Gallery')
            if gallery_btn:
                switch_page('gallery')
        else:
            st.write('You have checked ' + str(len(selected_modelVersions)) + ' images.')
            roster, promptBook, images_ds = load_hf_dataset()
            st.write("## roster")
            st.write(roster[roster['modelVersion_id'].isin(selected_modelVersions)])
            # st.write(roster)
            # st.write("## promptBook")
            # st.write(promptBook)