pages/Gallery.py

import json
import os
import requests

import altair as alt
import extra_streamlit_components as stx
import numpy as np
import pandas as pd
import streamlit as st
import streamlit.components.v1 as components

from bs4 import BeautifulSoup
from datasets import load_dataset, Dataset, load_from_disk
from huggingface_hub import login
from streamlit_agraph import agraph, Node, Edge, Config
from streamlit_extras.switch_page_button import switch_page
from streamlit_extras.no_default_selectbox import selectbox
from sklearn.svm import LinearSVC

SCORE_NAME_MAPPING = {'clip': 'clip_score', 'rank': 'msq_score', 'pop': 'model_download_count'}


class GalleryApp:
    def __init__(self, promptBook, images_ds):
        self.promptBook = promptBook
        self.images_ds = images_ds

        # init gallery state
        if 'gallery_state' not in st.session_state:
            st.session_state.gallery_state = {}

        # initialize selected_dict
        if 'selected_dict' not in st.session_state:
            st.session_state['selected_dict'] = {}

        if 'gallery_focus' not in st.session_state:
            st.session_state.gallery_focus = {'tag': None, 'prompt': None}

    def gallery_standard(self, items, col_num, info):
        rows = len(items) // col_num + 1
        containers = [st.container() for _ in range(rows)]
        for idx in range(0, len(items), col_num):
            row_idx = idx // col_num
            with containers[row_idx]:
                cols = st.columns(col_num)
                for j in range(col_num):
                    if idx + j < len(items):
                        with cols[j]:
                            # show image
                            # image = self.images_ds[items.iloc[idx + j]['row_idx'].item()]['image']
                            image = f"https://modelcofferbucket.s3-accelerate.amazonaws.com/{items.iloc[idx + j]['image_id']}.png"
                            st.image(image, use_column_width=True)

                            # handel checkbox information
                            prompt_id = items.iloc[idx + j]['prompt_id']
                            modelVersion_id = items.iloc[idx + j]['modelVersion_id']

                            check_init = True if modelVersion_id in st.session_state.selected_dict.get(prompt_id, []) else False

                            # st.write("Position: ", idx + j)

                            # show checkbox
                            st.checkbox('Select', key=f'select_{prompt_id}_{modelVersion_id}', value=check_init)

                            # show selected info
                            for key in info:
                                st.write(f"**{key}**: {items.iloc[idx + j][key]}")

    def gallery_graph(self, items):
        items = load_tsne_coordinates(items)

        # sort items to be popularity from low to high, so that most popular ones will be on the top
        items = items.sort_values(by=['model_download_count'], ascending=True).reset_index(drop=True)

        scale = 50
        items.loc[:, 'x'] = items['x'] * scale
        items.loc[:, 'y'] = items['y'] * scale

        nodes = []
        edges = []

        for idx in items.index:
            # if items.loc[idx, 'modelVersion_id'] in st.session_state.selected_dict.get(items.loc[idx, 'prompt_id'], 0):
            #     opacity = 0.2
            # else:
            #     opacity = 1.0

            nodes.append(Node(id=items.loc[idx, 'image_id'],
                              # label=str(items.loc[idx, 'model_name']),
                              title=f"model name: {items.loc[idx, 'model_name']}\nmodelVersion name: {items.loc[idx, 'modelVersion_name']}\nclip score: {items.loc[idx, 'clip_score']}\nmcos score: {items.loc[idx, 'mcos_score']}\npopularity: {items.loc[idx, 'model_download_count']}",
                              size=20,
                              shape='image',
                              image=f"https://modelcofferbucket.s3-accelerate.amazonaws.com/{items.loc[idx, 'image_id']}.png",
                              x=items.loc[idx, 'x'].item(),
                              y=items.loc[idx, 'y'].item(),
                              # fixed=True,
                              color={'background': '#E0E0E1', 'border': '#ffffff', 'highlight': {'border': '#F04542'}},
                              # opacity=opacity,
                              shadow={'enabled': True, 'color': 'rgba(0,0,0,0.4)', 'size': 10, 'x': 1, 'y': 1},
                              borderWidth=2,
                              shapeProperties={'useBorderWithImage': True},
                              )
                         )

        config = Config(width='100%',
                        height='600',
                        directed=True,
                        physics=False,
                        hierarchical=False,
                        interaction={'navigationButtons': True, 'dragNodes': False, 'multiselect': False},
                        # **kwargs
                        )

        return agraph(nodes=nodes,
                      edges=edges,
                      config=config,
                      )

    def selection_panel(self, items):
        # temperal function

        selecters = st.columns([1, 4])

        if 'score_weights' not in st.session_state:
            st.session_state.score_weights = [1.0, 0.8, 0.2, 0.8]

        # select sort type
        with selecters[0]:
            sort_type = st.selectbox('Sort by', ['Scores', 'IDs and Names'])
            if sort_type == 'Scores':
                sort_by = 'weighted_score_sum'

        # select other options
        with selecters[1]:
            if sort_type == 'IDs and Names':
                sub_selecters = st.columns([3, 1])
                # select sort by
                with sub_selecters[0]:
                    sort_by = st.selectbox('Sort by',
                                           ['model_name', 'model_id', 'modelVersion_name', 'modelVersion_id', 'norm_nsfw'],
                                           label_visibility='hidden')

                continue_idx = 1

            else:
                # add custom weights
                sub_selecters = st.columns([1, 1, 1, 1])

                with sub_selecters[0]:
                    clip_weight = st.number_input('Clip Score Weight', min_value=-100.0, max_value=100.0, value=1.0, step=0.1, help='the weight for normalized clip score')
                with sub_selecters[1]:
                    mcos_weight = st.number_input('Dissimilarity Weight', min_value=-100.0, max_value=100.0, value=0.8, step=0.1, help='the weight for m(eam) s(imilarity) q(antile) score for measuring distinctiveness')
                with sub_selecters[2]:
                    pop_weight = st.number_input('Popularity Weight', min_value=-100.0, max_value=100.0, value=0.2, step=0.1, help='the weight for normalized popularity score')

                items.loc[:, 'weighted_score_sum'] = round(items[f'norm_clip'] * clip_weight + items[f'norm_mcos'] * mcos_weight + items[
                    'norm_pop'] * pop_weight, 4)

                continue_idx = 3

                # save latest weights
                st.session_state.score_weights[0] = round(clip_weight, 2)
                st.session_state.score_weights[1] = round(mcos_weight, 2)
                st.session_state.score_weights[2] = round(pop_weight, 2)

            # select threshold
            with sub_selecters[continue_idx]:
                nsfw_threshold = st.number_input('NSFW Score Threshold', min_value=0.0, max_value=1.0, value=0.8, step=0.01, help='Only show models with nsfw score lower than this threshold, set 1.0 to show all images')
                items = items[items['norm_nsfw'] <= nsfw_threshold].reset_index(drop=True)

            # save latest threshold
            st.session_state.score_weights[3] = nsfw_threshold

        # # draw a distribution histogram
        # if sort_type == 'Scores':
        #     try:
        #         with st.expander('Show score distribution histogram and select score range'):
        #             st.write('**Score distribution histogram**')
        #             chart_space = st.container()
        #             # st.write('Select the range of scores to show')
        #             hist_data = pd.DataFrame(items[sort_by])
        #             mini = hist_data[sort_by].min().item()
        #             mini = mini//0.1 * 0.1
        #             maxi = hist_data[sort_by].max().item()
        #             maxi = maxi//0.1 * 0.1 + 0.1
        #             st.write('**Select the range of scores to show**')
        #             r = st.slider('Select the range of scores to show', min_value=mini, max_value=maxi, value=(mini, maxi), step=0.05, label_visibility='collapsed')
        #             with chart_space:
        #                 st.altair_chart(altair_histogram(hist_data, sort_by, r[0], r[1]), use_container_width=True)
        #             # event_dict = altair_component(altair_chart=altair_histogram(hist_data, sort_by))
        #             # r = event_dict.get(sort_by)
        #             if r:
        #                 items = items[(items[sort_by] >= r[0]) & (items[sort_by] <= r[1])].reset_index(drop=True)
        #                 # st.write(r)
        #     except:
        #         pass

        display_options = st.columns([1, 4])

        with display_options[0]:
            # select order
            order = st.selectbox('Order', ['Ascending', 'Descending'], index=1 if sort_type == 'Scores' else 0)
            if order == 'Ascending':
                order = True
            else:
                order = False

        with display_options[1]:

            # select info to show
            info = st.multiselect('Show Info',
                                  ['model_name', 'model_id', 'modelVersion_name', 'modelVersion_id',
                                   'weighted_score_sum', 'model_download_count', 'clip_score', 'mcos_score',
                                   'nsfw_score', 'norm_nsfw'],
                                  default=sort_by)

        # apply sorting to dataframe
        items = items.sort_values(by=[sort_by], ascending=order).reset_index(drop=True)

        # select number of columns
        col_num = st.slider('Number of columns', min_value=1, max_value=9, value=4, step=1, key='col_num')

        return items, info, col_num

    def sidebar(self, items, prompt_id, note):
        with st.sidebar:
            # prompt_tags = self.promptBook['tag'].unique()
            # # sort tags by alphabetical order
            # prompt_tags = np.sort(prompt_tags)[::1]
            #
            # tag = st.selectbox('Select a tag', prompt_tags, index=5)
            #
            # items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True)
            #
            # prompts = np.sort(items['prompt'].unique())[::1]
            #
            # selected_prompt = st.selectbox('Select prompt', prompts, index=3)

            # mode = st.radio('Select a mode', ['Gallery', 'Graph'], horizontal=True, index=1)

            # items = items[items['prompt'] == selected_prompt].reset_index(drop=True)

            # st.title('Model Visualization and Retrieval')

            # show source
            if isinstance(note, str):
                if note.isdigit():
                    st.caption(f"`Source: civitai`")
                else:
                    st.caption(f"`Source: {note}`")
            else:
                st.caption("`Source: Parti-prompts`")

            # show image metadata
            image_metadatas = ['prompt', 'negativePrompt', 'sampler', 'cfgScale', 'size', 'seed']
            for key in image_metadatas:
                label = ' '.join(key.split('_')).capitalize()
                st.write(f"**{label}**")
                if items[key][0] == ' ':
                    st.write('`None`')
                else:
                    st.caption(f"{items[key][0]}")

            # for note as civitai image id, add civitai reference
            if isinstance(note, str) and note.isdigit():
                try:
                    st.write(f'**[Civitai Reference](https://civitai.com/images/{note})**')
                    res = requests.get(f'https://civitai.com/images/{note}')
                    # st.write(res.text)
                    soup = BeautifulSoup(res.text, 'html.parser')
                    image_section = soup.find('div', {'class': 'mantine-12rlksp'})
                    image_url = image_section.find('img')['src']
                    st.image(image_url, use_column_width=True)
                except:
                    pass

        # return prompt_tags, tag, prompt_id, items

    def app(self):
        st.write('### Model Visualization and Retrieval')
        # st.write('This is a gallery of images generated by the models')

        # build the tabular view
        prompt_tags = self.promptBook['tag'].unique()
        # sort tags by alphabetical order
        prompt_tags = np.sort(prompt_tags)[::1].tolist()

        # chosen_data = [stx.TabBarItemData(id=tag, title=tag, description='') for tag in prompt_tags]
        # tag = stx.tab_bar(chosen_data, key='tag', default='food')

        # save tag to session state on change
        tag = st.radio('Select a tag', prompt_tags, index=5, horizontal=True, key='tag', label_visibility='collapsed')

        # tabs = st.tabs(prompt_tags)
        # for i in range(len(prompt_tags)):
        #     with tabs[i]:
        #         tag = prompt_tags[i]
        items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True)

        prompts = np.sort(items['prompt'].unique())[::1].tolist()

        # st.caption('Select a prompt')
        subset_selector = st.columns([3, 1])
        with subset_selector[0]:
            selected_prompt = selectbox('Select prompt', prompts, key=f'prompt_{tag}', no_selection_label='---', label_visibility='collapsed', index=0)
            # st.session_state.prompt_idx_last_time = prompts.index(selected_prompt) if selected_prompt else 0

        if selected_prompt is None:
            # st.markdown(':orange[Please select a prompt above👆]')
            st.write('**Feel free to navigate among tags and pages! Your selection will be saved within one log-in session.**')

            with subset_selector[-1]:
                st.write(':orange[👈 **Please select a prompt**]')

        else:
            items = items[items['prompt'] == selected_prompt].reset_index(drop=True)
            prompt_id = items['prompt_id'].unique()[0]
            note = items['note'].unique()[0]

            # add state to session state
            if prompt_id not in st.session_state.gallery_state:
                st.session_state.gallery_state[prompt_id] = 'graph'

            # add focus to session state
            st.session_state.gallery_focus['tag'] = tag
            st.session_state.gallery_focus['prompt'] = selected_prompt

            # add safety check for some prompts
            safety_check = True

            # load unsafe prompts
            unsafe_prompts = json.load(open('./data/unsafe_prompts.json', 'r'))
            for prompt_tag in prompt_tags:
                if prompt_tag not in unsafe_prompts:
                    unsafe_prompts[prompt_tag] = []
            # # manually add unsafe prompts
            # unsafe_prompts['world knowledge'] = [83]
            # unsafe_prompts['abstract'] = [1, 3]

            if int(prompt_id.item()) in unsafe_prompts[tag]:
                st.warning('This prompt may contain unsafe content. They might be offensive, depressing, or sexual.')
                safety_check = st.checkbox('I understand that this prompt may contain unsafe content. Show these images anyway.', key=f'safety_{prompt_id}')

            print('current state: ', st.session_state.gallery_state[prompt_id])

            if st.session_state.gallery_state[prompt_id] == 'graph':
                if safety_check:
                    self.graph_mode(prompt_id, items)
                with subset_selector[-1]:
                    has_selection = False
                    try:
                        if len(st.session_state.selected_dict.get(prompt_id, [])) > 0:
                            has_selection = True
                    except:
                        pass

                    if has_selection:
                        checkout = st.button('Check out selections', use_container_width=True, type='primary')
                        if checkout:
                            print('checkout')

                            st.session_state.gallery_state[prompt_id] = 'gallery'
                            print(st.session_state.gallery_state[prompt_id])
                            st.experimental_rerun()
                    else:
                        st.write(':orange[👇 **Select images you like below**]')

            elif st.session_state.gallery_state[prompt_id] == 'gallery':
                items = items[items['modelVersion_id'].isin(st.session_state.selected_dict[prompt_id])].reset_index(
                    drop=True)
                self.gallery_mode(prompt_id, items)

                with subset_selector[-1]:
                    state_operations = st.columns([1, 1])
                    with state_operations[0]:
                        back = st.button('Back to 🖼️', use_container_width=True)
                        if back:
                            st.session_state.gallery_state[prompt_id] = 'graph'
                            st.experimental_rerun()

                    with state_operations[1]:
                        forward = st.button('Check out', use_container_width=True, type='primary', on_click=self.submit_actions, args=('Continue', prompt_id))
                        if forward:
                            switch_page('ranking')

        try:
            self.sidebar(items, prompt_id, note)
        except:
            pass

    def graph_mode(self, prompt_id, items):
        graph_cols = st.columns([3, 1])
        # prompt = st.chat_input(f"Selected model version ids: {str(st.session_state.selected_dict.get(prompt_id, []))}",
        #                        disabled=False, key=f'{prompt_id}')
        # if prompt:
        #     switch_page("ranking")

        with graph_cols[0]:
            graph_space = st.empty()

            with graph_space.container():
                return_value = self.gallery_graph(items)

        with graph_cols[1]:
            if return_value:
                with st.form(key=f'{prompt_id}'):
                    image_url = f"https://modelcofferbucket.s3-accelerate.amazonaws.com/{return_value}.png"

                    st.image(image_url)

                    item = items[items['image_id'] == return_value].reset_index(drop=True).iloc[0]
                    modelVersion_id = item['modelVersion_id']

                    # handle selection
                    if 'selected_dict' in st.session_state:
                        if item['prompt_id'] not in st.session_state.selected_dict:
                            st.session_state.selected_dict[item['prompt_id']] = []

                        if modelVersion_id in st.session_state.selected_dict[item['prompt_id']]:
                            checked = True
                        else:
                            checked = False

                    if checked:
                        # deselect = st.button('Deselect', key=f'select_{item["prompt_id"]}_{item["modelVersion_id"]}', use_container_width=True)
                        deselect = st.form_submit_button('Deselect', use_container_width=True)
                        if deselect:
                            st.session_state.selected_dict[item['prompt_id']].remove(item['modelVersion_id'])
                            self.remove_ranking_states(item['prompt_id'])
                            st.experimental_rerun()

                    else:
                        # select = st.button('Select', key=f'select_{item["prompt_id"]}_{item["modelVersion_id"]}', use_container_width=True, type='primary')
                        select = st.form_submit_button('Select', use_container_width=True, type='primary')
                        if select:
                            st.session_state.selected_dict[item['prompt_id']].append(item['modelVersion_id'])
                            self.remove_ranking_states(item['prompt_id'])
                            st.experimental_rerun()

                    # st.write(item)
                    infos = ['model_name', 'modelVersion_name', 'model_download_count', 'clip_score', 'mcos_score',
                             'nsfw_score']

                    infos_df = item[infos]
                    # rename columns
                    infos_df = infos_df.rename(index={'model_name': 'Model', 'modelVersion_name': 'Version', 'model_download_count': 'Downloads', 'clip_score': 'Clip Score', 'mcos_score': 'mcos Score', 'nsfw_score': 'NSFW Score'})
                    st.table(infos_df)

                    # for info in infos:
                    #     st.write(f"**{info}**:")
                    #     st.write(item[info])

            else:
                st.info('Please click on an image to show')

    def gallery_mode(self, prompt_id, items):
        items, info, col_num = self.selection_panel(items)

        # if 'selected_dict' in st.session_state:
        #     # st.write('checked: ', str(st.session_state.selected_dict.get(prompt_id, [])))
        #     dynamic_weight_options = ['Grid Search', 'SVM', 'Greedy']
        #     dynamic_weight_panel = st.columns(len(dynamic_weight_options))
        #
        #     if len(st.session_state.selected_dict.get(prompt_id, [])) > 0:
        #         btn_disable = False
        #     else:
        #         btn_disable = True
        #
        #     for i in range(len(dynamic_weight_options)):
        #         method = dynamic_weight_options[i]
        #         with dynamic_weight_panel[i]:
        #             btn = st.button(method, use_container_width=True, disabled=btn_disable, on_click=self.dynamic_weight, args=(prompt_id, items, method))

        # prompt = st.chat_input(f"Selected model version ids: {str(st.session_state.selected_dict.get(prompt_id, []))}", disabled=False, key=f'{prompt_id}')
        # if prompt:
        #     switch_page("ranking")

        # with st.form(key=f'{prompt_id}'):
            # buttons = st.columns([1, 1, 1])
        # buttons_space = st.columns([1, 1, 1])
        gallery_space = st.empty()

        # with buttons_space[0]:
        #     continue_btn = st.button('Proceed selections to ranking', use_container_width=True, type='primary')
        #     if continue_btn:
        #         # self.submit_actions('Continue', prompt_id)
        #         switch_page("ranking")
        #
        # with buttons_space[1]:
        #     deselect_btn = st.button('Deselect All', use_container_width=True)
        #     if deselect_btn:
        #         self.submit_actions('Deselect', prompt_id)
        #
        # with buttons_space[2]:
        #     refresh_btn = st.button('Refresh', on_click=gallery_space.empty, use_container_width=True)

        with gallery_space.container():
            with st.spinner('Loading images...'):
                self.gallery_standard(items, col_num, info)

                # st.info("Don't forget to scroll back to top and click the 'Confirm Selection' button to save your selection!!!")

    def submit_actions(self, status, prompt_id):
        # remove counter from session state
        # st.session_state.pop('counter', None)
        self.remove_ranking_states('prompt_id')
        if status == 'Select':
            modelVersions = self.promptBook[self.promptBook['prompt_id'] == prompt_id]['modelVersion_id'].unique()
            st.session_state.selected_dict[prompt_id] = modelVersions.tolist()
            print(st.session_state.selected_dict, 'select')
            st.experimental_rerun()
        elif status == 'Deselect':
            st.session_state.selected_dict[prompt_id] = []
            print(st.session_state.selected_dict, 'deselect')
            st.experimental_rerun()
            # self.promptBook.loc[self.promptBook['prompt_id'] == prompt_id, 'checked'] = False
        elif status == 'Continue':
            st.session_state.selected_dict[prompt_id] = []
            for key in st.session_state:
                keys = key.split('_')
                if keys[0] == 'select' and keys[1] == str(prompt_id):
                    if st.session_state[key]:
                        st.session_state.selected_dict[prompt_id].append(int(keys[2]))
            # switch_page("ranking")
            print(st.session_state.selected_dict, 'continue')
            # st.experimental_rerun()

    def dynamic_weight(self, prompt_id, items, method='Grid Search'):
        selected = items[
            items['modelVersion_id'].isin(st.session_state.selected_dict[prompt_id])].reset_index(drop=True)
        optimal_weight = [0, 0, 0]

        if method == 'Grid Search':
            # grid search method
            top_ranking = len(items) * len(selected)

            for clip_weight in np.arange(-1, 1, 0.1):
                for mcos_weight in np.arange(-1, 1, 0.1):
                    for pop_weight in np.arange(-1, 1, 0.1):

                        weight_all = clip_weight*items[f'norm_clip'] + mcos_weight*items[f'norm_mcos'] + pop_weight*items['norm_pop']
                        weight_all_sorted = weight_all.sort_values(ascending=False).reset_index(drop=True)
                        # print('weight_all_sorted:', weight_all_sorted)
                        weight_selected = clip_weight*selected[f'norm_clip'] + mcos_weight*selected[f'norm_mcos'] + pop_weight*selected['norm_pop']

                        # get the index of values of weight_selected in weight_all_sorted
                        rankings = []
                        for weight in weight_selected:
                            rankings.append(weight_all_sorted.index[weight_all_sorted == weight].tolist()[0])
                        if sum(rankings) <= top_ranking:
                            top_ranking = sum(rankings)
                            print('current top ranking:', top_ranking, rankings)
                            optimal_weight = [clip_weight, mcos_weight, pop_weight]
            print('optimal weight:', optimal_weight)

        elif method == 'SVM':
            # svm method
            print('start svm method')
            # get residual dataframe that contains models not selected
            residual = items[~items['modelVersion_id'].isin(selected['modelVersion_id'])].reset_index(drop=True)
            residual = residual[['norm_clip_crop', 'norm_mcos_crop', 'norm_pop']]
            residual = residual.to_numpy()
            selected = selected[['norm_clip_crop', 'norm_mcos_crop', 'norm_pop']]
            selected = selected.to_numpy()

            y = np.concatenate((np.full((len(selected), 1), -1), np.full((len(residual), 1), 1)), axis=0).ravel()
            X = np.concatenate((selected, residual), axis=0)

            # fit svm model, and get parameters for the hyperplane
            clf = LinearSVC(random_state=0, C=1.0, fit_intercept=False, dual='auto')
            clf.fit(X, y)
            optimal_weight = clf.coef_[0].tolist()
            print('optimal weight:', optimal_weight)
            pass

        elif method == 'Greedy':
            for idx in selected.index:
                # find which score is the highest, clip, mcos, or pop
                clip_score = selected.loc[idx, 'norm_clip_crop']
                mcos_score = selected.loc[idx, 'norm_mcos_crop']
                pop_score = selected.loc[idx, 'norm_pop']
                if clip_score >= mcos_score and clip_score >= pop_score:
                    optimal_weight[0] += 1
                elif mcos_score >= clip_score and mcos_score >= pop_score:
                    optimal_weight[1] += 1
                elif pop_score >= clip_score and pop_score >= mcos_score:
                    optimal_weight[2] += 1

            # normalize optimal_weight
            optimal_weight = [round(weight/len(selected), 2) for weight in optimal_weight]
            print('optimal weight:', optimal_weight)
            print('optimal weight:', optimal_weight)

        st.session_state.score_weights[0: 3] = optimal_weight


    def remove_ranking_states(self, prompt_id):
        # for drag sort
        try:
            st.session_state.counter[prompt_id] = 0
            st.session_state.ranking[prompt_id] = {}
            print('remove ranking states')
        except:
            print('no sort ranking states to remove')

        # for battles
        try:
            st.session_state.pointer[prompt_id] = {'left': 0, 'right': 1}
            print('remove battles states')
        except:
            print('no battles states to remove')

        # for page progress
        try:
            st.session_state.progress[prompt_id] = 'ranking'
            print('reset page progress states')
        except:
            print('no page progress states to be reset')


# hist_data = pd.DataFrame(np.random.normal(42, 10, (200, 1)), columns=["x"])
@st.cache_resource
def altair_histogram(hist_data, sort_by, mini, maxi):
    brushed = alt.selection_interval(encodings=['x'], name="brushed")

    chart = (
        alt.Chart(hist_data)
        .mark_bar(opacity=0.7, cornerRadius=2)
        .encode(alt.X(f"{sort_by}:Q", bin=alt.Bin(maxbins=25)), y="count()")
        # .add_selection(brushed)
        # .properties(width=800, height=300)
    )

    # Create a transparent rectangle for highlighting the range
    highlight = (
        alt.Chart(pd.DataFrame({'x1': [mini], 'x2': [maxi]}))
        .mark_rect(opacity=0.3)
        .encode(x='x1', x2='x2')
        # .properties(width=800, height=300)
    )

    # Layer the chart and the highlight rectangle
    layered_chart = alt.layer(chart, highlight)

    return layered_chart


@st.cache_data
def load_hf_dataset(show_NSFW=False):
    # login to huggingface
    login(token=os.environ.get("HF_TOKEN"))

    # load from huggingface
    roster = pd.DataFrame(load_dataset('MAPS-research/GEMRec-Roster', split='train'))
    promptBook = pd.DataFrame(load_dataset('MAPS-research/GEMRec-Metadata', split='train'))
    # images_ds = load_from_disk(os.path.join(os.getcwd(), 'data', 'promptbook'))
    images_ds = None  # set to None for now since we use s3 bucket to store images

    # # process dataset
    # roster = roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name',
    #                                                    'model_download_count']].drop_duplicates().reset_index(drop=True)

    # add 'custom_score_weights' column to promptBook if not exist
    if 'weighted_score_sum' not in promptBook.columns:
        promptBook.loc[:, 'weighted_score_sum'] = 0

    # merge roster and promptbook
    promptBook = promptBook.merge(roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name', 'model_download_count']],
                                                                    on=['model_id', 'modelVersion_id'], how='left')

    # add column to record current row index
    promptBook.loc[:, 'row_idx'] = promptBook.index

    # apply a nsfw filter
    if not show_NSFW:
        promptBook = promptBook[promptBook['norm_nsfw'] <= 0.8].reset_index(drop=True)
        print('nsfw filter applied', len(promptBook))

    # add a column that adds up 'norm_clip', 'norm_mcos', and 'norm_pop'
    score_weights = [1.0, 0.8, 0.2]
    promptBook.loc[:, 'total_score'] = round(promptBook['norm_clip'] * score_weights[0] + promptBook['norm_mcos'] * score_weights[1] + promptBook['norm_pop'] * score_weights[2], 4)

    return roster, promptBook, images_ds

@st.cache_data
def load_tsne_coordinates(items):
    # load tsne coordinates
    tsne_df = pd.read_parquet('./data/feats_tsne.parquet')

    # print(tsne_df['modelVersion_id'].dtype)

    # print('before merge:', items)
    items = items.merge(tsne_df, on=['modelVersion_id', 'prompt_id'], how='left')
    # print('after merge:', items)
    return items


if __name__ == "__main__":
    st.set_page_config(page_title="Model Coffer Gallery", 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:
        # st.write('You have already logged in as ' + st.session_state.user_id[0])
        roster, promptBook, images_ds = load_hf_dataset(st.session_state.show_NSFW)
        # print(promptBook.columns)

        # # initialize selected_dict
        # if 'selected_dict' not in st.session_state:
        #     st.session_state['selected_dict'] = {}

        app = GalleryApp(promptBook=promptBook, images_ds=images_ds)
        app.app()

    with open('./css/style.css') as f:
        st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)