import hashlib
import os

import pandas as pd
import plotly.express as px
import streamlit as st
from bat import Benchmark, Config, Reporter, Tester


def get_nice_benchmark_name(bench_name):
    prettified_names = {
        "holmes": "Holmes",
        "helm_lite_narrativeqa": "Helm Lite NarrativeQA",
        "helm_lite_naturalquestionsopen": "Helm Lite NaturalQuestionsOpen",
        "helm_lite_naturalquestionsclosed": "Helm Lite NaturalQuestionsClosed",
        "helm_lite_openbookqa": "Helm Lite OpenBookQA",
        "helm_lite_mmlu": "Helm Lite MMLU",
        "helm_lite_math_equivalentcot": "Helm Lite MathEquivalentCOT",
        "helm_lite_gsm8k": "Helm Lite GSM8K",
        "helm_lite_legalbench": "Helm Lite LegalBench",
        "helm_lite_medqa": "Helm Lite MedQA",
        "helm_lite_wmt2014": "Helm Lite WMT2014",
        "hfv2_bbh": "HFv2 BBH",
        "hfv2_bbh_raw": "HFv2 BBH Raw",
        "hfv2_gpqa": "HFv2 GPQA",
        "hfv2_ifeval": "HFv2 IFEval",
        "hfv2_math_lvl_5": "HFv2 Math Level 5",
        "hfv2_mmlu_pro": "HFv2 MMLU Pro",
        "hfv2_musr": "HFv2 MuSR",
        "oc_mmlu": "OpenCompass MMLU",
        "oc_mmlu_pro": "OpenCompass MMLU Pro",
        "oc_cmmlu": "OpenCompass CMMLU",
        "oc_bbh": "OpenCompass BBH",
        "oc_gqpa_dimand": "OpenCompass GQPA-Dimand",
        "oc_humaneval": "OpenCompass HumanEval",
        "oc_ifeval": "OpenCompass IFEval",
        "helm_mmlu": "Helm MMLU",
        "helm_boolq": "Helm BoolQ",
        "helm_narrativeqa": "Helm NarrativeQA",
        "helm_naturalquestionsclosed": "Helm NaturalQuestionsClosed",
        "helm_naturalquestionsopen": "Helm NaturalQuestionsOpen",
        "helm_quac": "Helm QuAC",
        "helm_openbookqa": "Helm OpenBookQA",
        "helm_imdb": "Helm IMDB",
        "helm_civilcomments": "Helm CivilComments",
        "helm_raft": "Helm RAFT",
        "mmlu_pro": "MMLU Pro",
        "mixeval_triviaqa": "MixEval TriviaQA",
        "mixeval_mmlu": "MixEval MMLU",
        "mixeval_drop": "MixEval DROP",
        "mixeval_hellaswag": "MixEval HellaSwag",
        "mixeval_commonsenseqa": "MixEval CommonsenseQA",
        "mixeval_triviaqa_hard": "MixEval TriviaQA Hard",
        "mixeval_mmlu_hard": "MixEval MMLU Hard",
        "mixeval_drop_hard": "MixEval DROP Hard",
        "oc_language": "OpenCompass Language",
        "oc_knowledge": "OpenCompass Knowledge",
        "oc_reasoning": "OpenCompass Reasoning",
        "oc_math": "OpenCompass Math",
        "oc_code": "OpenCompass Code",
        "oc_instruct": "OpenCompass Instruction",
        "oc_agent": "OpenCompass Agent",
        "oc_arena": "OpenCompass Arena",
        "lb_reasoning": "LiveBench Reasoning",
        "lb_coding": "LiveBench Coding",
        "lb_mathematics": "LiveBench Mathematics",
        "lb_data_analysis": "LiveBench Data Analysis",
        "lb_language": "LiveBench Language",
        "lb_if": "LiveBench Instruction Following",
        "wb_info_seek": "WildBench Information Seeking",
        "wb_creative": "WildBench Creative",
        "wb_code_debug": "WildBench Code Debugging",
        "wb_math_data": "WildBench Math & Data",
        "wb_reason_plan": "WildBench Reasoning & Planning",
        "wb_score": "WildBench Score",
        "hfv1_arc": "HFv1 ARC",
        "hfv1_gsm8k": "HFv1 GSM8K",
        "hfv1_hellaswag": "HFv1 HellaSwag",
        "hfv1_mmlu": "HFv1 MMLU",
        "hfv1_truthfulqa": "HFv1 TruthfulQA",
        "hfv1_winogrande": "HFv1 Winogrande",
        "biggen_grounding": "BigBench Grounding",
        "biggen_instruction_following": "BigBench Instruction Following",
        "biggen_planning": "BigBench Planning",
        "biggen_reasoning": "BigBench Reasoning",
        "biggen_refinement": "BigBench Refinement",
        "biggen_safety": "BigBench Safety",
        "biggen_theory_of_mind": "BigBench Theory of Mind",
        "biggen_tool_usage": "BigBench Tool Usage",
        "biggen_multilingual": "BigBench Multilingual",
        "lb_reasoning_average": "LiveBench Reasoning Average",
        "lb_coding_average": "LiveBench Coding Average",
        "lb_mathematics_average": "LiveBench Mathematics Average",
        "lb_data_analysis_average": "LiveBench Data Analysis Average",
        "lb_language_average": "LiveBench Language Average",
        "lb_if_average": "LiveBench Instruction Following Average",
        "helm_lite": "Helm Lite",
        "hf_open_llm_v2": "HF OpenLLM v2",
        "opencompass_academic": "OpenCompass Academic",
        "arena_elo": "Arena Elo",
        "helm_classic": "Helm Classic",
        "mixeval": "MixEval",
        "mixeval_hard": "MixEval Hard",
        "opencompass": "OpenCompass",
        "alphacaeval_v2lc": "AlphacaEval v2lc",
        "livebench_240725": "LiveBench 240725",
        "wb_elo_lc": "WildBench Elo LC",
        "arena_hard": "Arena Hard",
        "agentbench": "AgentBench",
        "hf_open_llm_v1": "HF OpenLLM v1",
        "biggen": "BigBench",
        "livebench_240624": "LiveBench 240624",
        "mt_bench": "MT-Bench",
    }

    if bench_name in prettified_names:
        return prettified_names[bench_name]
    else:
        return bench_name


holistic_scenarios = [
    get_nice_benchmark_name(scen)
    for scen in [
        # "holmes",
        "helm_lite",
        # "narrativeqa",
        # "naturalquestionsopen",
        # "naturalquestionsclosed",
        # "openbookqa",
        # "mmlu",
        # "math_equivalentcot",
        # "gsm8k",
        # "legalbench",
        # "medqa",
        # "wmt2014",
        # "arc_c",
        # "arc_e",
        # "boolq",
        # "csqa",
        # "hellaswag",
        # "piqa",
        # "siqa",
        # "winogrande",
        # "olmes_average",
        # "bbh",
        # "bbh_raw",
        # "gpqa",
        "hf_open_llm_v2",
        # "ifeval",
        # "math_lvl_5",
        # "mmlu_pro",
        # "musr",
        "opencompass_academic",
        # "oc_mmlu",
        # "oc_mmlu_pro",
        # "oc_cmmlu",
        # "oc_bbh",
        # "oc_gqpa_dimand",
        # "oc_math",
        # "oc_humaneval",
        # "oc_ifeval",
        # "helm_mmlu",
        "arena_elo",
        "helm_classic",
        # "quac",
        # "truthfulqa",
        # "ms_marcoregular",
        # "ms_marcotrec",
        # "cnn/dailymail",
        # "xsum",
        # "imdb",
        # "civilcomments",
        # "raft",
        "mixeval_hard",
        "mixeval",
        # "arena_elo0527",
        "opencompass",
        # "oc_language",
        # "oc_knowledge",
        # "oc_reasoning",
        # "oc_code",
        # "oc_instruct",
        # "oc_agent",
        # "oc_arena",
        "alphacaeval_v2lc",
        "livebench_240725",
        "livebench_240624",
        # "lb_reasoning",
        # "lb_coding",
        # "lb_mathematics",
        # "lb_data_analysis",
        # "lb_language",
        # "lb_if",
        "wb_elo_lc",
        # "wb_info_seek",
        # "wb_creative",
        # "wb_code_debug",
        # "wb_math_data",
        # "wb_reason_plan",
        # "wb_score",
        # "boolqmixed",
        "arena_hard",
        "agentbench",
        # "arc",
        "hf_open_llm_v1",
        "biggen",
        # "biggen_grounding",
        # "biggen_instruction_following",
        # "biggen_planning",
        # "biggen_reasoning",
        # "biggen_refinement",
        # "biggen_safety",
        # "biggen_theory_of_mind",
        # "biggen_tool_usage",
        # "biggen_multilingual",
        # "lb_global_average",
        # "lb_reasoning_average",
        # "lb_coding_average",
        # "lb_mathematics_average",
        # "lb_data_analysis_average",
        # "lb_language_average",
        # "lb_if_average",
    ]
]


st.markdown(
    """<h1 style='text-align: center; color: black;'>🏋️‍♂️ BenchBench Leaderboard 🏋️‍♂️</h1>""",
    unsafe_allow_html=True,
)

st.markdown(
    "We are excited to share the BenchBench-Leaderboard, a crucial component of our comprehensive research work -- [Benchmark Agreement Testing Done Right: A Guide for LLM Benchmark Evaluation](https://arxiv.org/abs/2407.13696). "
    "This leaderboard is a meta-benchmark that ranks benchmarks based on their agreement with the crowd harnessing many different references. "
)


all_scenarios_for_aggragate = Benchmark()
all_scenarios_for_aggragate.load_local_catalog()
all_scenarios_for_aggragate = (
    all_scenarios_for_aggragate.df["scenario"].unique().tolist()
)

st.subheader("The Leaderboard", divider=True)
# st.subheader("🏋️‍♂️ BenchBench Leaderboard 🏋", divider=True)

leftcol, rightcol = st.columns([2, 1])

with st.expander("Leaderboard configurations (defaults are great BTW)", icon="⚙️"):
    with st.form("my_form"):
        all_scenarios_for_aggragate_with_all = [
            get_nice_benchmark_name(scenario)
            for scenario in all_scenarios_for_aggragate
        ]

        aggragate_scenarios = st.multiselect(
            "Scenarios in Aggregate (defualts are the 'Holistic' benchmarks)",
            all_scenarios_for_aggragate,
            holistic_scenarios,
        )

        corr_type = st.selectbox(
            label="Select Correlation type", options=["kendall", "pearson"], index=0
        )

        aggragate_scenario_blacklist = [
            scen
            for scen in all_scenarios_for_aggragate
            if scen not in aggragate_scenarios
        ]

        model_select_strategy = st.selectbox(
            label="Select strategy",
            options=["random", "top_aggregate", "somewhere_aggregate"],
            index=0,
        )

        n_models_taken_list = st.slider(
            label="Select number of models to use",
            min_value=3,
            max_value=20,
            value=10,
        )

        n_models_taken_list = [n_models_taken_list]

        n_exps = 10

        submitted = st.form_submit_button(label="Run BAT")


uploaded_file = st.file_uploader("add your benchmark as a CSV")
st.download_button(
    label="Download example CSV",
    data=pd.read_csv("assets/mybench.csv").to_csv(index=False).encode("utf-8"),
    file_name="mybench.csv",
    mime="text/csv",
)

my_benchmark = Benchmark()
if uploaded_file is not None:
    df = pd.read_csv(uploaded_file)
    my_benchmark.assign_df(df, data_source="Uploaded Benchmark")


def run_load(
    aggragate_scenario_blacklist=[],
    n_models_taken_list=[5],
    model_select_strategy_list=["random"],
    corr_types=["kendall"],
    n_exps=10,
    my_benchmark=Benchmark(),
    use_caching=True,
):
    # Create a hash of the inputs to generate a unique cache file for each set of inputs
    input_str = (
        str(aggragate_scenario_blacklist)
        + str(n_models_taken_list)
        + str(model_select_strategy_list)
        + str(corr_types)
        + str(n_exps)
    )

    if not my_benchmark.is_empty:
        input_str += str(
            hashlib.sha256(
                my_benchmark.df.to_csv(index=False).encode("utf-8")
            ).hexdigest()
        )

    input_hash = hashlib.md5(input_str.encode()).hexdigest()
    cache_file = f"agreements_cache_{input_hash}.csv"

    # Define the cache directory
    cache_dir = "cache"
    os.makedirs(cache_dir, exist_ok=True)
    cache_path = os.path.join(cache_dir, cache_file)

    # Check if the cache file exists
    if os.path.exists(cache_path) and use_caching:
        print("Loading cached results...")
        agreements = pd.read_csv(cache_path)
        aggregate_scores = pd.read_csv(
            cache_path.replace("agreement", "aggregate_scores")
        )

        return agreements, aggregate_scores

    else:
        print("Cached results not found, calculating")

        cfg = Config(
            exp_to_run="example",
            n_models_taken_list=n_models_taken_list,
            model_select_strategy_list=model_select_strategy_list,
            corr_types=corr_types,
            n_exps=n_exps if n_models_taken_list != [0] else 1,
        )

        holistic = Benchmark()
        holistic.load_local_catalog()
        holistic.df = holistic.df.query("scenario in @holistic_scenarios")

        holistic.clear_repeated_scenarios()
        holistic.add_aggragete(
            new_col_name="aggregate",
            agg_source_name="holistic",
            scenario_blacklist=aggragate_scenario_blacklist,
            min_scenario_for_models_to_appear_in_agg=5,
        )

        aggragate_scores = holistic.df.query('scenario=="aggregate"')[
            ["model", "score"]
        ].sort_values(by="score", ascending=False)

        allbench = Benchmark()
        allbench.load_local_catalog()

        # allbench.df = allbench.df[~allbench.df["source"].str.contains("livebench")]

        allbench.extend(my_benchmark)
        # allbench.df = allbench.df.drop(columns=["tag"])
        allbench.clear_repeated_scenarios()

        # removing and adding the holistic scenarios
        allbench.df = allbench.df.query("scenario not in @holistic_scenarios")
        allbench = allbench.extend(holistic)

        tester = Tester(cfg=cfg)

        agreements = tester.all_vs_all_agreement_testing(
            allbench,
            single_source_scenario="aggregate",  # olny measuring all with the aggragate
        )

        agreements.to_csv(cache_path, index=False)
        aggragate_scores.to_csv(
            cache_path.replace("agreement", "aggregate_scores"), index=False
        )

    return agreements, aggragate_scores


agreements, aggragare_score_df = run_load(
    aggragate_scenario_blacklist=aggragate_scenario_blacklist,
    n_models_taken_list=n_models_taken_list,
    model_select_strategy_list=[model_select_strategy],
    corr_types=[corr_type],
    n_exps=n_exps,
    my_benchmark=my_benchmark,
)

reporter = Reporter()
z_scores = reporter.get_all_z_scores(agreements=agreements, aggragate_name="aggregate")
z_scores.drop(columns=["n_models_of_corr_with_agg"], inplace=True)

corr_name = f"{'Kendall Tau' if corr_type=='kendall' else 'Per.'} Corr."

z_scores["z_score"] = z_scores["z_score"].round(2)
z_scores["corr_with_agg"] = z_scores["corr_with_agg"].round(2)
z_scores["p_value_of_corr_with_agg"] = z_scores["p_value_of_corr_with_agg"].round(2)
# z_scores["n_models_of_corr_with_agg"] = z_scores["n_models_of_corr_with_agg"].round(1)

z_scores["source"] = z_scores["source"].apply(lambda x: x.split(".csv")[0])

# print(z_scores["scenario"].unique().tolist())

z_scores["scenario"] = z_scores["scenario"].apply(lambda x: get_nice_benchmark_name(x))

data = (
    z_scores.rename(
        columns={
            "scenario": "Benchmark",
            "z_score": "Z Score",
            "corr_with_agg": corr_name,
            "p_value_of_corr_with_agg": "p-value of Corr.",
            # "n_models_of_corr_with_agg": "# Models Used",
            "source": "Source",
        }
    )
    .sort_values("Z Score", ascending=False)
    .reset_index(drop=True)
)


# Apply coloring based on 'Z' valuesz
def highlight_uploaded_benchmark(row):
    if row["Source"] == "Uploaded Benchmark":
        return ["background-color: rgba(100,100,100,0.1)"] * len(row)
    else:
        return [""] * len(row)


styled_data = (
    data.style.background_gradient(
        subset=["Z Score"],
        cmap="RdYlGn",
        vmin=-data["Z Score"].abs().max(),
        vmax=data["Z Score"].abs().max(),
    )
    .apply(highlight_uploaded_benchmark, axis=1)
    .background_gradient(
        subset=["p-value of Corr."],
        cmap="Reds",
        vmin=0.1,
        vmax=1,
    )
    .format(subset=["Z Score", corr_name, "p-value of Corr."], formatter="{:.2}")
)

print(data["Benchmark"].unique().tolist())

st.dataframe(
    data=styled_data,
    hide_index=True,
    use_container_width=True,
    height=500,
)

aggragare_score_df.rename(
    columns={
        "model": "Model",
        "score": "Mean Win Rate over Selected Scenarios for Aggragate",
    },
    inplace=True,
)
with st.expander(label="Model scored by the aggragate"):
    st.dataframe(
        data=aggragare_score_df,
        hide_index=True,
        height=500,
        use_container_width=True,
    )


st.markdown(
    "BenchBench-Leaderboard complements our study, where we analyzed over 40 prominent benchmarks and introduced standardized practices to enhance the robustness and validity of benchmark evaluations through the [BenchBench Python package](#). "
    "The BenchBench-Leaderboard serves as a dynamic platform for benchmark comparison and is an essential tool for researchers and practitioners in the language model field aiming to select and utilize benchmarks effectively. "
)

st.subheader("How did we get the Z Scores?", divider=True)

st.write(r"""
            Section 3.1 in our work shows how using a single reference benchmark drastically hurts the roubustness and validity of BAT.
            To remedy this, we propose to test benchmark agreement with an aggragate benchmark and compare the agreement to other benchmarks.
            We recommend to perform this comparison using the [Z score](https://en.wikipedia.org/wiki/Standard_score) and demonstrate obtaining it to a benchmark of your selection.
            In the follwing way: $z_i=(x_i-\mu_{i...N}) / \sigma_{i...N}$ where $x_i$ is the agreement of the $i$th benchmark to the aggragate and $\mu_{i...N}$,$\sigma_{i...N}$ are the 
            mean and standard deviation of the agreements of the other benchmarks to the aggragate.
            """)


benchmarks = data["Benchmark"].unique().tolist()
plotted_scenario = st.selectbox(
    "Choose Benchmark to plot",
    benchmarks,
    index=benchmarks.index("Arena Elo"),
)


fig = px.histogram(
    data.query("Benchmark!=@plotted_scenario"), x=corr_name, nbins=len(data) - 1
)
# Add a vertical line at a specific x-coordinate
# Replace 'x_value' with the actual value where you want the line
x_value = 0.5  # Example value, adjust as necessary
fig.add_vline(
    x=data.query("Benchmark==@plotted_scenario")[corr_name].iloc[0],
    line_dash="dash",
    line_color="red",
)
# Update layout to add a title
fig.update_layout(
    title="Histogram of Correlation Values",  # Change the title text as needed
    title_x=0.3,  # Centers the title
    title_font=dict(size=20, family="CMU"),  # Customize font if needed
)

# # Plot!
st.plotly_chart(fig, use_container_width=True)

st.subheader("Why should you use the BenchBench Leaderboard?")

st.markdown(
    """

    Current practices in Benchmark Agreement Testing (BAT) often suffer from a lack of standardization and transparency, which can lead to inconsistent results and diminished trust in benchmark evaluations. Several key issues are prevalent in the field:

    """
)

st.markdown(
    """
    - **Lack of Standard Methodologies:** Unlike other scientific procedures that follow rigorous methodologies, BAT lacks uniform procedures across different studies. Researchers often employ varied criteria for selecting benchmarks and models for comparison, which leads to results that cannot be easily compared or replicated. This variation undermines the reliability of conclusions drawn from BAT and makes it difficult for other researchers to build on existing work.
    """
)

st.image(
    "images/motivation.png",
    caption="Conclusions depend on the models considered. Kendall-tau correlations between the LMSys Arena benchmark and three other benchmarks: BBH, MMLU, and Alpaca v2. Each group of bars represents the correlation for different sets of top models, specifically the top 5, top 10, and top 15 (overlapping) models (according to the Arena). The results indicate that the degree of agreement between benchmarks varies with the number of top models considered, highlighting that different selections of models can lead to varying conclusions about benchmark agreement.",
    use_column_width=True,
)

st.markdown(
    """
    - **Arbitrary Selection of Reference Benchmarks:** One of the most critical decisions in BAT is the choice of reference benchmarks. Currently, this choice is often arbitrary and lacks a clear rationale, influenced by availability or personal preference rather than strategic alignment with the benchmark’s purpose. This can skew the results significantly, as different benchmarks may not be equally representative or relevant to the models being tested.
    """
)
st.markdown(
    """
    - **Inadequate Model Representation:** BAT frequently relies on a limited subset of models, which may not comprehensively represent the diversity of architectures and training paradigms in modern language models. This selective representation can lead to biased agreement scores that favor certain types of models over others, failing to provide a holistic view of model performance across different benchmarks.
    """
)

st.image(
    "images/pointplot_granularity_matters.png",
    caption="Correlations increase with number of models. Mean correlation (y) between each benchmark (lines) and the rest, given different numbers of models. The Blue and Orange lines are the average of all benchmark pair correlations with models sampled randomly (orange) or in contiguous sets (blue). The shaded lines represents adjacent sampling for the different benchmarks.",
    use_column_width=True,
)

st.markdown(
    """
    - **Overemphasis on Correlation Metrics:** Current BAT practices tend to over-rely on correlation metrics without adequately considering their limitations and the context of their application. While these metrics can provide useful insights, they are often treated as definitive evidence of agreement without acknowledging that high correlation does not necessarily imply conceptual alignment between benchmarks.
    """
)

st.markdown(
    """
    To address these issues, there is a critical need for a more structured approach to BAT that includes clear guidelines for benchmark and model selection, a broader consideration of agreement metrics, and an acknowledgment of the evolving nature of technology in this space. By reforming BAT practices, the research community can improve the reliability and utility of benchmarks as tools for evaluating and advancing language models.
    """
)


st.image(
    "images/ablations.png",
    caption="Our recommendations substantially reduce the variance of BAT. Ablation analysis for each BAT recommendation separately and their combinations.",
    use_column_width=True,
)


st.header("The BenchBench package")

st.markdown("""
### Overview

The BAT package is designed to facilitate benchmark agreement testing for NLP models. It allows users to easily compare multiple models against various benchmarks and generate comprehensive reports on their agreement.

### Installation

To install the BAT package, you can use pip:

```
pip install bat-package
```

### Usage Example

Below is a step-by-step example of how to use the BAT package to perform agreement testing.

#### Step 1: Configuration

First, set up the configuration for the tests:

```python
import pandas as pd
from bat import Tester, Config, Benchmark, Reporter
from bat.utils import get_holistic_benchmark

cfg = Config(
    exp_to_run="example",
    n_models_taken_list=[0],
    model_select_strategy_list=["random"],
    n_exps=10
)
```

#### Step 2: Fetch Model Names

Fetch the names of the reference models to be used for scoring:

```python
tester = Tester(cfg=cfg)
models_for_benchmark_scoring = tester.fetch_reference_models_names(
    reference_benchmark=get_holistic_benchmark(), n_models=20
)
print(models_for_benchmark_scoring)
```

#### Step 3: Load and Prepare Benchmark

Load a new benchmark and add an aggregate column:

```python
newbench_name = "fakebench"
newbench = Benchmark(
    pd.read_csv(f"src/bat/assets/{newbench_name}.csv"),
    data_source=newbench_name,
)
newbench.add_aggregate(new_col_name=f"{newbench_name}_mwr")
```

#### Step 4: Agreement Testing

Perform all-vs-all agreement testing on the new benchmark:

```python
newbench_agreements = tester.all_vs_all_agreement_testing(newbench)
reporter = Reporter()
reporter.draw_agreements(newbench_agreements)
```

#### Step 5: Extend and Clean Benchmark

Extend the new benchmark with holistic data and clear repeated scenarios:

```python
allbench = newbench.extend(get_holistic_benchmark())
allbench.clear_repeated_scenarios(source_to_keep=newbench_name)
```

#### Step 6: Comprehensive Agreement Testing

Perform comprehensive agreement testing and visualize:

```python
all_agreements = tester.all_vs_all_agreement_testing(allbench)
reporter.draw_agreements(all_agreements)
```
""")