import streamlit as st
import pandas as pd
import json
import numpy as np

# 读取 JSON 数据
with open('GenAI-Bench_tags.json', 'r') as file:
    data = json.load(file)
with open('genai_image.json', 'r') as file:
    humanRatings = json.load(file)
# Streamlit 页面标题

st.set_page_config(page_title="GenAI-Bench Dataset Viewer", page_icon=None, layout="wide", initial_sidebar_state="auto",
                   menu_items=None)
st.title('GenAI-Bench Dataset Viewer')
# 多选框
basic_options = []
advanced_options = []

for i in data:
    data[i]['basic'].sort()
    data[i]['advanced'].sort()

    for basic in data[i]['basic']:
        if basic not in basic_options:
            basic_options.append(basic)
    for adv in data[i]['advanced']:
        if adv not in advanced_options:
            advanced_options.append(adv)
    data[i]['id'] = i
    # modles = ["DALLE_3","DeepFloyd_I_XL_v1","Midjourney_6","SDXL_2_1","SDXL_Base","SDXL_Turbo"]
    data[i]['DALLE_3'] = f"app/static/DALLE_3/{i}.jpeg"
    data[i]['DeepFloyd_I_XL_v1'] = f"app/static/DeepFloyd_I_XL_v1/{i}.jpeg"
    data[i]['Midjourney_6'] = f"app/static/Midjourney_6/{i}.jpeg"
    data[i]['SDXL_2_1'] = f"app/static/SDXL_2_1/{i}.jpeg"
    data[i]['SDXL_Base'] = f"app/static/SDXL_Base/{i}.jpeg"
    data[i]['SDXL_Turbo'] = f"app/static/SDXL_Turbo/{i}.jpeg"

    data[i]['DALLE_3_Human'] = round(np.mean(humanRatings[i]["models"]["DALLE_3"]), 1)
    data[i]['DeepFloyd_I_XL_v1_Human'] = round(np.mean(humanRatings[i]["models"]["DeepFloyd_I_XL_v1"]), 1)
    data[i]['Midjourney_6_Human'] = round(np.mean(humanRatings[i]["models"]["Midjourney_6"]), 1)
    data[i]['SDXL_2_1_Human'] = round(np.mean(humanRatings[i]["models"]["SDXL_2_1"]), 1)
    data[i]['SDXL_Base_Human'] = round(np.mean(humanRatings[i]["models"]["SDXL_Base"]), 1)
    data[i]['SDXL_Turbo_Human'] = round(np.mean(humanRatings[i]["models"]["SDXL_Turbo"]), 1)

selected_basic = st.multiselect('Select Basic Skills:', basic_options)
selected_advanced = st.multiselect('Select Advanced Skills:', advanced_options)

# 筛选数据
filtered_data = [
    data[item] for item in data
    if all(elem in data[item]['basic'] for elem in selected_basic) and
       all(advanced in data[item]['advanced'] for advanced in selected_advanced)
]

# 显示筛选后的数据
if filtered_data:
    df = pd.DataFrame(filtered_data)
    df = df.reindex(columns=["id", "prompt", "basic", "advanced", "DALLE_3", "DALLE_3_Human", "DeepFloyd_I_XL_v1",
                             "DeepFloyd_I_XL_v1_Human", "Midjourney_6", "Midjourney_6_Human", "SDXL_2_1",
                             "SDXL_2_1_Human", "SDXL_Base", "SDXL_Base_Human", "SDXL_Turbo", "SDXL_Turbo_Human"])
    st.dataframe(data=df, width = 4096, height = 800,
                 column_config={
                     "name": "Data Explorer",
                     "id": st.column_config.NumberColumn("ID", format="%d", width="small"),
                     'basic': st.column_config.ListColumn(label="Basic Skills", width="large", help=None),
                     'advanced': st.column_config.ListColumn(label="Advanced Skills", width="large", help=None),
                     'prompt': st.column_config.TextColumn(label="Prompt", width="large", help=None, disabled=None,
                                                           required=None,
                                                           default=None, max_chars=None, validate=None),
                     "DALLE_3": st.column_config.ImageColumn(label="DALLE_3", width="small", help=None),
                     "DALLE_3_Human": st.column_config.NumberColumn("Rating Human", format="%.1f", width="small", help="Rating Human for DALLE_3"),
                     "DeepFloyd_I_XL_v1": st.column_config.ImageColumn(label="DeepFloyd", width="small",
                                                                       help="DeepFloyd_I_XL_v1"),
                     "DeepFloyd_I_XL_v1_Human": st.column_config.NumberColumn("Rating Human", format="%.1f",
                                                                              width="small", help="Rating Human for DeepFloyd"),
                     "Midjourney_6": st.column_config.ImageColumn(label="Midjourney", width="small", help="Midjourney_6"),
                     "Midjourney_6_Human": st.column_config.NumberColumn("Rating Human", format="%.1f",
                                                                         width="small", help="Rating Human for Midjourney_6"),
                     "SDXL_2_1": st.column_config.ImageColumn(label="SDXL_2_1", width="small", help=None),
                     "SDXL_2_1_Human": st.column_config.NumberColumn("Rating Human", format="%.1f", width="small", help="Rating Human for SDXL_2_1"),
                     "SDXL_Base": st.column_config.ImageColumn(label="SDXL_Base", width="small", help=None),
                     "SDXL_Base_Human": st.column_config.NumberColumn("Rating Human", format="%.1f", width="small", help="Rating Human for SDXL_Base"),
                     "SDXL_Turbo": st.column_config.ImageColumn(label="SDXL_Turbo", width="small", help=None),
                     "SDXL_Turbo_Human": st.column_config.NumberColumn("Rating Human", format="%.1f", width="small", help="Rating Human for SDXL_Turbo"),
                 },
                 hide_index=True, selection_mode="single-row")
else:
    st.write("No data matches the selected filters.")