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/VBench/", 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: 10px;
    }
</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['name'].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: Comprehensive Benchmark Suite for Video Generative Models (VBench).
    """
)

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/VBench/T2V-Leaderboard/discussions) and tag `@Ziqi`.
 """
)

# 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')