import streamlit as st
st.set_page_config(layout="wide")
import pandas  as pd
import os
import json
import shutil
from huggingface_hub import Repository

REFERENCE_NAME = "references"
SUBMISSION_NAME = "vbench_leaderboard_submission"


SUBMISSION_URL = os.path.join("https://huggingface.co/datasets/Vchitect/", SUBMISSION_NAME)


TEST_SETS = [
    "subject consistency",
    "background consistency",
    "temporal flickering",
    "motion smoothness",
    "dynamic degree",
    "aesthetic quality",
    "imaging quality",
    "object class",
    "multiple objects",
    "human action",
    "color",
    "spatial relationship",
    "scene",
    "appearance style",
    "temporal style",
    "overall consistency"
    ]

# style = """
# <style>
#     th {
#         font-size: 18px;
#         font-family: Arial, sans-serif;
#         text-transform: capitalize;
#     }
# </style>
# """
style = """
<style>
    th, td {
        font-size: 14px;
        font-family: Arial, sans-serif;
        text-transform: capitalize;
        text-align: center; /* Center-align text for headers and cells */
    }
    table {
        width: 100%; /* Optional: makes the table take full width */
        margin-left: auto; 
        margin-right: auto;
        text-align: center;
    }
    .dataframe th, .dataframe td {
        text-align: center !important;
    }
</style>
"""

CSV_RESULTS_FILE = os.path.join(SUBMISSION_NAME, "results.csv")

HF_TOKEN = os.environ.get("HF_TOKEN")

try:
    submission_repo = Repository(
        local_dir="vbench_leaderboard_submission", clone_from=SUBMISSION_URL, use_auth_token=HF_TOKEN, repo_type="dataset"
    )
    submission_repo.git_pull()
except Exception as e:
    print(e)


all_submissions = [
    file_name
    for file_name in os.listdir(SUBMISSION_NAME)
    if file_name.endswith('.json')
]

all_results = pd.read_csv(CSV_RESULTS_FILE)
with open(os.path.join(SUBMISSION_NAME, "verified_model.txt")) as f:
    verified_model = [i.strip().lower() for i in f.readlines()]
all_results['verified'] = all_results['Model Name (clickable)'].apply(lambda x: '√' if x.lower() in verified_model else ' ')


# Write table form CSV
table = all_results.copy()
table = table.round(4)
# columns_to_convert = table.columns[1:-1]
# table[columns_to_convert] = table[columns_to_convert].applymap(lambda x: f'{x * 100}'[:5]+'%')



# Streamlit
st.markdown("# VBench ")


# st.markdown(
#     f"""
#     This is the leaderboard of VBench.
# > **VBench: Comprehensive Benchmark Suite for Video Generative Models**
# > [[Paper](https://vchitect.github.io/VBench-project/assets/vbench/VBench_paper.pdf)] | [[Project Page](https://vchitect.github.io/VBench-project/)] | [[GitHub Code](https://github.com/Vchitect/VBench)] | [[Video](https://www.youtube.com/watch?v=7IhCC8Qqn8Y)]
#     """
# )
st.markdown("This is the leaderboard of VBench.")
st.markdown("> **VBench: Comprehensive Benchmark Suite for Video Generative Models**  ([Paper](https://vchitect.github.io/VBench-project/assets/vbench/VBench_paper.pdf), [Project Page](https://vchitect.github.io/VBench-project/), [GitHub Code](https://github.com/Vchitect/VBench), [Video](https://www.youtube.com/watch?v=7IhCC8Qqn8Y))")

default_index = list(table.columns[1:-1]).index("overall consistency") if "overall consistency" in table.columns else 0
sort_option = st.selectbox(
    'Choose a column to sort by',
    table.columns[1:-1], 
    index=default_index,
)
table = table.sort_values(by=sort_option, ascending=False)

st.write(style + table.to_markdown(index=False), unsafe_allow_html=True)
# st.markdown(
#     """
#     For more information, refer to the paper submission on [Arxiv](https://).
#     """
# )

st.markdown(
    """
    ## Submitting to VBench
    \n
    To submit to VBench, download the prompt suite from [VBench/Prompt](https://github.com/Vchitect/VBench/tree/master/prompts). Upload your zipped submissions for scoring and placement on the leaderboard. 
    \n
    Should you experience any issues, open an issue using the link [new discussion](https://huggingface.co/spaces/Vchitect/VBench_Leaderboard/discussions) and tag `@Ziqi` or `@ynhe`.
 """
)

# Using the "with" syntax
with st.form(key="my_form"):
    uploaded_file = st.file_uploader("Choose a json file")
    submit_button = st.form_submit_button(label="Submit")

if submit_button:
    if uploaded_file is None:
        raise ValueError("Please make sure to have uploaded a json file.")

    submission = uploaded_file.name.split(".json")[0]
    with st.spinner(f"Uploading {submission}..."):
        with open(os.path.join(submission_repo.local_dir, os.path.basename(uploaded_file.name)),'wb') as f:
            f.write(uploaded_file.getvalue())
        submission_repo.push_to_hub()

    with st.spinner(f"Update Score for {submission}..."):
        results = {"name": submission}
        upload_score = json.loads(uploaded_file.getvalue())
        for info in upload_score:
            results[info['dimension']] = info['final_score']
        all_results.loc[len(all_results)] = results
        all_results.to_csv(CSV_RESULTS_FILE, index=False)
        commit_url = submission_repo.push_to_hub()

    st.success('Please refresh this space (CTRL+R) to see your result')