import streamlit as st
from pandas import read_csv
import os
import jiwer
from huggingface_hub import Repository
import zipfile

REFERENCE_NAME = "references"
SUBMISSION_NAME = "submissions"

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

TEST_SETS = [
    "ls-clean",
    "ls-other",
    "cv",
    "vox",
    "ted",
    "giga",
    "spgi",
    "earnings",
    "ami",
]
EXPECTED_TEST_FILES = [f + ".txt" for f in TEST_SETS]

OPTIONAL_TEST_SETS = ["swbd", "ch", "chime-4"]
OPTIONAL_TEST_FILES = [f + ".txt" for f in OPTIONAL_TEST_SETS]

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

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


def compute_wer(pred_file, ref_file):
    with open(pred_file, "r", encoding="utf-8") as pred, open(
            ref_file, "r", encoding="utf-8"
    ) as ref:
        pred_lines = [line.strip() for line in pred.readlines()]
        ref_lines = [line.strip() for line in ref.readlines()]

    wer = jiwer.wer(ref_lines, pred_lines)
    return wer


reference_repo = Repository(
    local_dir="references", clone_from=REFERENCE_URL, use_auth_token=HF_TOKEN
)
submission_repo = Repository(
    local_dir="submissions", clone_from=SUBMISSION_URL, use_auth_token=HF_TOKEN
)
submission_repo.git_pull()

all_submissions = [
    folder
    for folder in os.listdir(SUBMISSION_NAME)
    if os.path.isdir(os.path.join(SUBMISSION_NAME, folder)) and folder != ".git"
]

all_results = read_csv(CSV_RESULTS_FILE)

# Write table form CSV
table = all_results.copy()

esb_column = table.pop("esb-score")
name_column = table.pop("name")
table.insert(0, "esb-score", esb_column)
table = table.select_dtypes(exclude=['object', 'string'])
table.insert(0, "name", name_column)
table = table.sort_values(by=['esb-score'], ascending=True, ignore_index=True)
table = table.round(2)
table.index = table.index + 1

# Streamlit
st.markdown("# ESB: A Benchmark For Multi-Domain End-to-End Speech Recognition")

st.markdown(
    f"""
    This is the leaderboard of the End-to end Speech Benchmark (ESB).
    Submitted systems are ranked by the **ESB Score** which is the average of
    all non-optional datasets: {", ".join(TEST_SETS)}. The optional datasets of swbd, ch and chime-4 do not contribute to the ESB score."""
)

# st.table(table)
st.dataframe(table.style.format(subset=["esb-score", *TEST_SETS, *OPTIONAL_TEST_SETS], formatter="{:.1f}"))

st.markdown(
    """
    ESB was proposed in *ESB: A Benchmark For Multi-Domain End-to-End Speech Recognition* by ...
    \n
    The abstract of the paper is as follows:
    \n
    *Speech recognition applications cover a range of different audio and text distributions, with different speaking styles, background noise, transcription punctuation and character casing. However, many speech recognition systems require dataset-specific tuning (audio filtering, punctuation removal and normalisation of casing), therefore assuming a-priori knowledge of both the audio and text distributions. This tuning requirement can lead to systems failing to generalise to other datasets and domains. To promote the development of multi-domain speech systems, we introduce the End-to-end Speech Benchmark (ESB) for evaluating the performance of a single automatic speech recognition (ASR) system across a broad set of speech datasets. Benchmarked systems must use the same data pre- and post-processing algorithm across datasets - assuming the audio and text data distributions are a-priori unknown. We compare a series of state-of-the-art (SoTA) end-to-end (E2E) systems on this benchmark, demonstrating how a single speechsystem can be applied and evaluated on a wide range of data distributions. We find E2E systems to be effective across datasets: in a fair comparison, E2E systems achieve within 2.6% of SoTA systems tuned to a specific dataset. Our analysis reveals that transcription artefacts, such as punctuation and casing, pose difficulties for ASR systems and should be included in evaluation. We believe E2E benchmarking over a range of datasets promotes the research of multi-domain speech recognition systems.*
    \n
    For more information, please see the official submission on [OpenReview.net](https://openreview.net/forum?id=9OL2fIfDLK).
    """
)

st.markdown(
    """
    ## Submitting to ESB
    \n
    To submit to ESB, download the audio data for the nine mandatory ESB test sets from [esb/datasets](https://huggingface.co/datasets/esb/datasets). The test sets contain audio data only. Evaluate your system on the nine test sets by generating predictions for the unlabelled audio samples. For each test set, save the predictions to a .txt file in the order that the audio samples are provided, with one prediction per line. Name the .txt file according to the ESB test set names shown in the table (e.g. the predictions for LibriSpeech test-clean should be named ls-clean.txt).
    \n
    Once you have evaluated your system on all nine test sets, move the predictions into one folder and zip it. The name you assign to the zipped folder will be the name that is shown on the table (e.g. whisper-aed.zip will be displayed as whisper-aed). Upload your zipped submissions for scoring and placement on the leaderboard.
 """
)

# Using the "with" syntax
with st.form(key="my_form"):
    uploaded_file = st.file_uploader("Choose a zip 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 zip file.")

    submission = uploaded_file.name.split(".zip")[0]
    with st.spinner(f"Uploading {submission}..."):
        with zipfile.ZipFile(uploaded_file, 'r') as zip_ref:
            zip_ref.extractall(submission_repo.local_dir)
            submission_repo.push_to_hub()

    with st.spinner(f"Computing ESB Score for {submission}..."):
        results = {"name": submission}
        all_submitted_files = os.listdir(os.path.join(SUBMISSION_NAME, submission))

        submitted_files = [f for f in all_submitted_files if f in EXPECTED_TEST_FILES]
        submitted_optional_files = [f for f in all_submitted_files if f in OPTIONAL_TEST_FILES]

        if sorted(EXPECTED_TEST_FILES) != sorted(submitted_files):
            raise ValueError(
                f"{', '.join(submitted_files)} were submitted, but expected {', '.join(EXPECTED_TEST_FILES)}"
            )

        for file in submitted_files:
            ref_file = os.path.join(REFERENCE_NAME, file)
            pred_file = os.path.join(SUBMISSION_NAME, submission, file)

            wer = compute_wer(pred_file, ref_file)
            results[file.split(".")[0]] = round(100 * wer, 2)

        for file in submitted_optional_files:
            ref_file = os.path.join(REFERENCE_NAME, file)
            pred_file = os.path.join(SUBMISSION_NAME, submission, file)

            wer = compute_wer(pred_file, ref_file)
            results[file.split(".")[0]] = round(100 * wer, 2)

        # ESB score is computed over the mandatory test sets only
        wer_values = [float(results[t]) for t in TEST_SETS]
        # first average over LS test sets
        wer_values = [sum(wer_values[:2]) / 2, *wer_values[2:]]
        # then macro-average over ESB test sets
        all_wer = sum(wer_values) / len(wer_values)

        results["esb-score"] = round(all_wer, 2)
        all_results = all_results.append(results, ignore_index=True)

        # save and upload new evaluated 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')