import streamlit as st
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import statsmodels.api as sm
import random
import pandas as pd
import numpy as np
import plotly.express as px
import plotly.graph_objects as go
import matplotlib.pyplot as plt
import seaborn as sns
from itertools import combinations

# Set the layout to wide
st.set_page_config(layout="wide")


# Custom CSS to center title and header
center_css = """
<style>
h1, h2, h3, h6{
    text-align: center;
}
</style>
"""

st.markdown(center_css, unsafe_allow_html=True)


def create_agreement_heatmap(df):
    # Create a list of unique annotators and sort them by annotator index
    unique_annotators = sorted(df["annotator_index"].unique())

    # Initialize the agreement matrix and count matrix
    agreement_matrix = pd.DataFrame(
        np.nan, index=unique_annotators, columns=unique_annotators
    )
    count_matrix = pd.DataFrame(
        np.zeros((len(unique_annotators), len(unique_annotators))),
        index=unique_annotators,
        columns=unique_annotators,
    )

    # Group by (instruction, output_1, output_2)
    grouped = df.groupby(["instruction", "output_1", "output_2"])

    for name, group in grouped:
        # Extract annotators and their preferences
        annotators = group["annotator_index"].values
        preferences = group["preference"].values

        # Iterate over all pairs of annotators in the group
        for (annotator1, pref1), (annotator2, pref2) in combinations(
            zip(annotators, preferences), 2
        ):
            if pref1 == pref2:  # If they agree
                if pd.isna(agreement_matrix.loc[annotator1, annotator2]):
                    agreement_matrix.loc[annotator1, annotator2] = 0
                if pd.isna(agreement_matrix.loc[annotator2, annotator1]):
                    agreement_matrix.loc[annotator2, annotator1] = 0
                agreement_matrix.loc[annotator1, annotator2] += 1
                agreement_matrix.loc[annotator2, annotator1] += 1
            count_matrix.loc[annotator1, annotator2] += 1
            count_matrix.loc[annotator2, annotator1] += 1

    # Normalize the agreement matrix by the count matrix
    for i in unique_annotators:
        for j in unique_annotators:
            if count_matrix.loc[i, j] > 0:
                agreement_matrix.loc[i, j] /= count_matrix.loc[i, j]

    # Plot the heatmap
    plt.figure(figsize=(10, 10))  # Make the heatmap square
    sns.heatmap(
        agreement_matrix,
        annot=True,
        fmt=".2f",
        cmap="PiYG",
        cbar=True,
        mask=np.isnan(agreement_matrix),
        vmin=0.0,
        vmax=1.0,
        square=True,
    )
    plt.title("Interannotator Agreement Heatmap")
    plt.xlabel("Annotator")
    plt.ylabel("Annotator")
    plt.tight_layout()
    return agreement_matrix


def prep_rankings_table(df, y_column):
    # Create a copy of the dataframe.
    df_copy = df.copy()

    # Select the columns we care about, sort by the y column, and reset the index.
    df_copy = (
        df_copy[
            [
                "model_name",
                y_column,
                "num_words_mean",
            ]
        ]
        .sort_values(y_column, ascending=False)
        .reset_index()
    )

    # Create a rank column.
    df_copy["rank"] = df_copy.index + 1

    # Round the y column.
    df_copy[y_column] = df_copy[y_column].round(2)

    # Fix the order.
    df_copy = df_copy[["rank", "model_name", y_column, "num_words_mean"]]
    return df_copy


def get_preference(preference_score):
    rounded_preference_score = int(preference_score.round(0).iloc[0])
    return get_preference_from_rounded_score(rounded_preference_score)
    # if rounded_preference_score == 2:
    #     return "[2>1]"
    # elif rounded_preference_score == 1:
    #     return "[1>2]"


def get_preference_from_rounded_score(score):
    if score == 2:
        return "[2>1]"
    elif score == 1:
        return "[1>2]"
    return "[1=2]"
    # raise ValueError(f"Invalid score: {score}")


def is_unanimous(series):
    if len(set(series.tolist())) == 1:
        return True
    return False


def app():
    fixed_model = "gpt4_1106_preview"

    # Ensure to initialize session state variables if they do not exist
    if "selected_instruction" not in st.session_state:
        st.session_state.selected_instruction = None

    if "selected_model" not in st.session_state:
        st.session_state.selected_model = "gpt4"

    if "selected_output_human_annotations" not in st.session_state:
        st.session_state.selected_output_human_annotations = None

    if "selected_judge" not in st.session_state:
        st.session_state.selected_judge = None

    if "selected_dataset" not in st.session_state:
        st.session_state.selected_dataset = "NEW"

    if "instruction_options" not in st.session_state:
        st.session_state.instruction_options = []

    if "instruction_options_human_annotations" not in st.session_state:
        st.session_state.instruction_options_human_annotations = []

    if "selected_instruction_human_annotations" not in st.session_state:
        st.session_state.selected_instruction_human_annotations = None

    # Function to update the instruction options based on selected dataset
    def update_instruction_options():
        selected_dataset = st.session_state.dataset_selector
        if selected_dataset == "all" or selected_dataset == "NEW":
            instruction_options = df_response_judging["instruction"].unique().tolist()
        elif (
            selected_dataset == "None"
            or selected_dataset is None
            or str(selected_dataset) == ""
        ):
            instruction_options = (
                df_response_judging[pd.isna(df_response_judging["dataset"])][
                    "instruction"
                ]
                .unique()
                .tolist()
            )
        else:
            instruction_options = (
                df_response_judging[df_response_judging["dataset"] == selected_dataset][
                    "instruction"
                ]
                .unique()
                .tolist()
            )

        st.session_state.instruction_options = instruction_options

    def update_instruction_options_human_annotations():
        selected_dataset = st.session_state.dataset_selector_human_annotations
        if selected_dataset == "all" or selected_dataset == "NEW":
            instruction_options = df_human_annotations["instruction"].unique().tolist()
        elif (
            selected_dataset == "None"
            or selected_dataset is None
            or str(selected_dataset) == ""
        ):
            instruction_options = (
                df_human_annotations[pd.isna(df_human_annotations["dataset"])][
                    "instruction"
                ]
                .unique()
                .tolist()
            )
        else:
            instruction_options = (
                df_human_annotations[
                    df_human_annotations["dataset"] == selected_dataset
                ]["instruction"]
                .unique()
                .tolist()
            )
        st.session_state.instruction_options_human_annotations = instruction_options

    def update_instruction():
        st.session_state.selected_instruction = st.session_state.instruction_selector

    def update_model():
        st.session_state.selected_model = st.session_state.model_selector

    def update_judge():
        st.session_state.selected_judge = st.session_state.judge_selector

    def randomize_selection():
        st.session_state.dataset_selector = random.choice(
            ["all"] + df_response_judging["dataset"].dropna().unique().tolist()
        )
        st.session_state.selected_model = random.choice(model_options)
        update_instruction_options()
        st.session_state.selected_instruction = random.choice(
            st.session_state.instruction_options
        )

    def randomize_selection_human_annotations():
        st.session_state.dataset_selector_human_annotations = random.choice(
            ["all"] + df_human_annotations["dataset"].dropna().unique().tolist()
        )
        update_instruction_options()
        st.session_state.selected_instruction_human_annotations = random.choice(
            st.session_state.instruction_options_human_annotations
        )
        st.session_state.selected_output_human_annotations = random.choice(
            df_human_annotations[
                df_human_annotations["instruction"]
                == st.session_state.selected_instruction_human_annotations
            ]["output_2"]
            .dropna()
            .tolist()
        )

    st.title("🦙 AlpacaEval Explorer 🦙")

    st.markdown(
        "###### An interactive tool to analyze and explore the data behind the [AlpacaEval Leaderboard](https://tatsu-lab.github.io/alpaca_eval/) in more depth"
    )

    st.markdown(
        "###### Created and maintained by [Justin Zhao](https://x.com/justinxzhao)"
    )

    col1, col2, col3 = st.columns(3)

    with col1:
        with st.expander("About AlpacaEval"):
            st.markdown(
                """- [AlpacaEval](https://github.com/tatsu-lab/alpaca_eval) is an evaluation benchmark to assess the performance of large language models (LLMs).
- It has high correlation with [Chatbot Arena](https://chat.lmsys.org/), and is a fast and affordable benchmark for chat LLMs that uses LLMs (specifically GPT-4) to estimate response quality.
- LLM responses are assessed in a pairwise fashion (arena), where each model's responses are compared to a reference model's responses. 
- All reference responses are generated by GPT-4-1106. The LLM Judge is also GPT-4-1106.

"""
            )

    with col2:
        with st.expander("About this tool"):
            st.markdown(
                """There are 3 main tabs.
1. Use the **Data explorer** to look at individual pairwise battles between models.
2. Use the **Length bias explorer** to look at how response lengths affect win rates.
3. Use the **Human cross annotations** tab to explore the human cross annotations.
"""
            )

    with col3:
        with st.expander("Motivation"):
            st.markdown(
                """
- Several arena-based benchmarks have demonstrated that a clear ranking among LLMs can be established, but there is a general dearth of analysis and understanding as to why the rankings are the way they are. For example, it's hard to discern how factors like feel and style
are weighed against correctness.
- I created this tool to provide a more interactive and intuitive way to explore the data behind the AlpacaEval leaderboard. It allows users to easily compare responses between models, look at individual battles, and analyze how response lengths affect win rates.
- If you have any feedback on the tool, please reach out.
    """
            )

    outer_tabs = st.tabs(
        [
            "Data explorer",
            "Length bias explorer",
            "Human cross annotations",
        ]
    )

    # Load the data
    df_human_annotations = pd.read_json("data/alpaca_farm_human_crossannotations.json")
    df = pd.read_json("data/model_win_rates.jsonl", lines=True, orient="records")
    # df_responses = pd.read_json("data/df_responses.jsonl", lines=True, orient="records")
    df_response_judging = pd.read_json(
        "data/df_response_judging.jsonl", lines=True, orient="records"
    )

    # Prepare the model selector options
    model_options = df_response_judging["generator_2"].unique().tolist()

    with outer_tabs[1]:
        # Define the preset groups
        presets = {
            "gpt": df[df["model_name"].str.contains("openai|gpt", case=False)][
                "model_name"
            ].tolist(),
            "claude": df[df["model_name"].str.contains("claude", case=False)][
                "model_name"
            ].tolist(),
            "moa": df[df["model_name"].str.contains("moa", case=False)][
                "model_name"
            ].tolist(),
            "llama": df[df["model_name"].str.contains("llama", case=False)][
                "model_name"
            ].tolist(),
            "custom": [],
        }

        # Add radio button for preset groups
        preset_selection = st.radio(
            "Select a preset group of models or choose 'custom' to select manually.",
            options=["custom", "gpt", "claude", "moa", "llama"],
        )

        # Add multiselect for custom model selection
        if preset_selection == "custom":
            selected_models = st.multiselect(
                "Select models to highlight", options=df["model_name"].unique()
            )
        else:
            selected_models = presets[preset_selection]

        st.divider()

        def create_scatter_plot(df, y_column, selected_models, title):
            fig = go.Figure()

            # Add scatter plots for num_words_mean and num_tokens_mean
            fig.add_trace(
                go.Scatter(
                    x=df["num_words_mean"],
                    y=df[y_column],
                    mode="markers",
                    name="words",
                    text=df["model_name"],
                    marker=dict(size=5, color="skyblue"),
                    showlegend=True,
                )
            )
            fig.add_trace(
                go.Scatter(
                    x=df["num_tokens_mean"],
                    y=df[y_column],
                    mode="markers",
                    name="tokens",
                    text=df["model_name"],
                    marker=dict(size=5, color="orange"),
                    showlegend=True,
                    visible="legendonly",  # Make 'words' trace initially visible only in legend
                )
            )

            # Highlight selected models
            if selected_models:
                selected_data = df[df["model_name"].isin(selected_models)]
                fig.add_trace(
                    go.Scatter(
                        x=selected_data["num_words_mean"],
                        y=selected_data[y_column],
                        mode="markers",
                        name="selected words",
                        text=selected_data["model_name"],
                        marker=dict(size=10, color="blue"),
                        showlegend=True,
                    )
                )
                fig.add_trace(
                    go.Scatter(
                        x=selected_data["num_tokens_mean"],
                        y=selected_data[y_column],
                        mode="markers",
                        name="selected tokens",
                        text=selected_data["model_name"],
                        marker=dict(size=10, color="orangered"),
                        showlegend=True,
                        visible="legendonly",  # Make 'selected words' trace initially visible only in legend
                    )
                )

            # Add trendlines
            def add_trendline(fig, x, y, name, color, visibility="legendonly"):
                X = sm.add_constant(df[x])
                model = sm.OLS(df[y], X).fit()
                trendline = model.predict(X)
                fig.add_trace(
                    go.Scatter(
                        x=df[x],
                        y=trendline,
                        mode="lines",
                        name=f"{name} trendline",
                        line=dict(color=color, width=2),
                        visible=visibility,  # Control the initial visibility
                    )
                )
                return model.rsquared

            r_squared_words = add_trendline(
                fig, "num_words_mean", y_column, "words", "blue", visibility=True
            )
            r_squared_tokens = add_trendline(
                fig, "num_tokens_mean", y_column, "tokens", "orangered"
            )

            # Update layout with titles and labels
            fig.update_layout(
                xaxis_title="Mean length",
                yaxis_title=(
                    "Win rate"
                    if y_column == "win_rate"
                    else (
                        "LC Win Rate"
                        if y_column == "length_controlled_winrate"
                        else "Discrete Win Rate"
                    )
                ),
                title=title,
                legend_title="Legend",
            )

            return fig, r_squared_words, r_squared_tokens

        st.markdown("#### Overall win rate")
        y_column1 = "length_controlled_winrate"
        y_column2 = "win_rate"
        y_column3 = "discrete_win_rate"

        fig1, r_squared_words_1, r_squared_tokens_1 = create_scatter_plot(
            df, y_column1, selected_models, "Length-Controlled Win Rate"
        )
        fig2, r_squared_words_2, r_squared_tokens_2 = create_scatter_plot(
            df, y_column2, selected_models, "Win Rate"
        )
        fig3, r_squared_words_3, r_squared_tokens_3 = create_scatter_plot(
            df, y_column3, selected_models, "Discrete Win Rate"
        )

        # Create tabs for each chart
        tab1, tab2, tab3 = st.tabs(["LC Win Rate", "Win Rate", "Discrete Win Rate"])

        with tab1:
            col1, col2 = st.columns([3, 2])
            col1.plotly_chart(fig1)
            col2.markdown("#### Rankings")
            prepped_df = prep_rankings_table(df, "length_controlled_winrate")
            col2.dataframe(
                prepped_df,
                hide_index=True,
            )
            with st.expander("Trendline R²"):
                st.markdown(
                    f"- R² (Words vs {y_column1}): {r_squared_words_1:.2f} \n- R² (Tokens vs {y_column1}): {r_squared_tokens_1:.2f}"
                )

        with tab2:
            col1, col2 = st.columns([3, 2])
            col1.plotly_chart(fig2)
            col2.markdown("#### Rankings")
            prepped_df = prep_rankings_table(df, "win_rate")
            col2.dataframe(
                prepped_df,
                hide_index=True,
            )
            with st.expander("Trendline R²"):
                st.markdown(
                    f"- R² (Words vs {y_column2}): {r_squared_words_2:.2f} \n- R² (Tokens vs {y_column2}): {r_squared_tokens_2:.2f}"
                )

        with tab3:
            col1, col2 = st.columns([3, 2])
            col1.plotly_chart(fig3)
            col2.markdown("#### Rankings")
            prepped_df = prep_rankings_table(df, "discrete_win_rate")
            col2.dataframe(
                prepped_df,
                hide_index=True,
            )
            with st.expander("Trendline R²"):
                st.markdown(
                    f"- R² (Words vs {y_column3}): {r_squared_words_3:.2f}\n- R² (Tokens vs {y_column3}): {r_squared_tokens_3:.2f}"
                )

        st.markdown("#### Length bias in battles")

        df_response_judging_copy = df_response_judging.copy()
        if not selected_models:
            df_response_judging_copy["output_1_num_words"] = df_response_judging_copy[
                "output_1"
            ].apply(lambda x: len(x.split()))
            df_response_judging_copy["output_2_num_words"] = df_response_judging_copy[
                "output_2"
            ].apply(lambda x: len(x.split()))
            df_response_judging_copy["output_num_words_diff"] = (
                df_response_judging_copy["output_1_num_words"]
                - df_response_judging_copy["output_2_num_words"]
            )
            df_response_judging_copy["assigned_preference"] = (
                df_response_judging_copy["preference"]
                .round(0)
                .apply(get_preference_from_rounded_score)
            )
        else:
            df_response_judging_copy = df_response_judging_copy[
                df_response_judging_copy["generator_2"].isin(selected_models)
            ]
            df_response_judging_copy["output_1_num_words"] = df_response_judging_copy[
                "output_1"
            ].apply(lambda x: len(x.split()))
            df_response_judging_copy["output_2_num_words"] = df_response_judging_copy[
                "output_2"
            ].apply(lambda x: len(x.split()))
            df_response_judging_copy["output_num_words_diff"] = (
                df_response_judging_copy["output_1_num_words"]
                - df_response_judging_copy["output_2_num_words"]
            )
            df_response_judging_copy["assigned_preference"] = (
                df_response_judging_copy["preference"]
                .round(0)
                .apply(get_preference_from_rounded_score)
            )

        col1, col2 = st.columns(2)
        fig = px.scatter(
            df_response_judging_copy,
            x="output_1_num_words",
            y="output_2_num_words",
            color="assigned_preference",
            title=f"Pairwise preference based on response length",
            labels={
                "output_1_num_words": f"{fixed_model} (1) number of words",
                "output_2_num_words": "Target model (2) number of words",
            },
            color_discrete_map={
                "[1>2]": "blue",
                "[2>1]": "orangered",
                "[1=2]": "green",
            },
        )
        col1.plotly_chart(fig)

        # Plot of output_num_words_diff histogram, colored by assigned_preference.
        fig = px.histogram(
            df_response_judging_copy,
            x="output_num_words_diff",
            color="assigned_preference",
            title=f"Pairwise preference counts based on difference in response length",
            color_discrete_map={
                "[1>2]": "blue",
                "[2>1]": "orangered",
                "[1=2]": "green",
            },
            range_x=[-500, 500],
            labels={
                "output_num_words_diff": "Length difference in words between gpt4_1106_preview and target model"
            },
        )
        col2.plotly_chart(fig)

        with st.expander("Raw data"):
            st.dataframe(df)

    # Data explorer
    with outer_tabs[0]:
        # Add randomize button at the top of the app
        st.markdown("#### Choose example")
        st.button(
            ":game_die: Randomize!",
            on_click=randomize_selection,
            type="primary",
        )

        left_col, right_col = st.columns([1, 3])

        st.session_state.selected_dataset = left_col.selectbox(
            "Select Dataset",
            ["all"] + df_response_judging["dataset"].dropna().unique().tolist(),
            key="dataset_selector",
            on_change=update_instruction_options,
        )
        update_instruction_options()
        st.session_state.selected_instruction = right_col.selectbox(
            f"Select Instruction ({len(st.session_state.instruction_options)} unique instructions)",
            st.session_state.instruction_options,
            key="instruction_selector",
            on_change=update_instruction,
            index=(
                st.session_state.instruction_options.index(
                    st.session_state.selected_instruction
                )
                if st.session_state.selected_instruction
                in st.session_state.instruction_options
                else 0
            ),
        )

        # All the models.
        all_models_judgings_details = df_response_judging[
            (df_response_judging["generator_1"] == fixed_model)
            & (
                df_response_judging["instruction"]
                == st.session_state.selected_instruction
            )
        ]

        st.divider()

        st.markdown(f"#### Selected instruction")
        st.info(st.session_state.selected_instruction)

        st.divider()

        st.markdown(f"#### Overall Battles")
        all_models_judgings_details["output_1_num_words"] = all_models_judgings_details[
            "output_1"
        ].apply(lambda x: len(x.split()))
        all_models_judgings_details["output_2_num_words"] = all_models_judgings_details[
            "output_2"
        ].apply(lambda x: len(x.split()))
        all_models_judgings_details["output_num_words_diff"] = (
            all_models_judgings_details["output_1_num_words"]
            - all_models_judgings_details["output_2_num_words"]
        )
        all_models_judgings_details["assigned_preference"] = (
            all_models_judgings_details["preference"]
            .round(0)
            .apply(get_preference_from_rounded_score)
        )

        # st.write(all_models_judgings_details)

        col1, col2, col3 = st.columns(3)

        fig = px.histogram(
            all_models_judgings_details,
            x="output_num_words_diff",
            color="assigned_preference",
            title=f"Pairwise preference counts based on difference in response length",
            color_discrete_map={
                "[1>2]": "blue",
                "[2>1]": "orangered",
                "[1=2]": "green",
            },
            range_x=[-500, 500],
            labels={
                "output_num_words_diff": "Difference in number of words between response 1 and 2.",
                "assigned_preference": "Assigned Preference",
            },
        )
        col1.plotly_chart(fig)

        # Plot of assigned preference counts.
        fig = px.histogram(
            all_models_judgings_details,
            x="assigned_preference",
            title=f"Assigned preferences for {fixed_model} vs. all models",
        )
        col2.plotly_chart(fig)

        # Models that are better than the fixed model.
        num_words_for_fixed_model = len(
            all_models_judgings_details.iloc[0]["output_1"].split()
        )
        better_models = all_models_judgings_details[
            all_models_judgings_details["assigned_preference"] == "[2>1]"
        ]

        shorter_models = better_models[
            better_models["output_2_num_words"] <= num_words_for_fixed_model
        ]
        longer_models = better_models[
            better_models["output_2_num_words"] > num_words_for_fixed_model
        ]
        col3.markdown(
            f"##### Models that are better than {fixed_model} ({num_words_for_fixed_model})"
        )
        if shorter_models.size != 0:
            shorter_models_string = ""
            for _, shorter_model in shorter_models.iterrows():
                if shorter_model["generator_2"] != fixed_model:
                    shorter_models_string += f"- {shorter_model['generator_2']} ({shorter_model['output_2_num_words']})\n"
            col3.markdown("**With shorter or equal length responses:**")
            col3.markdown(shorter_models_string)
        else:
            col3.write("None")
        if longer_models.size != 0:
            longer_models_string = ""
            for _, longer_model in longer_models.iterrows():
                if longer_model["generator_2"] != fixed_model:
                    longer_models_string += f"- {longer_model['generator_2']} ({longer_model['output_2_num_words']})\n"
            col3.markdown("**With longer responses:**")
            col3.markdown(longer_models_string)
        else:
            col3.write("None")

        # Judging details.
        st.markdown(f"#### Individual Battle Details")
        judging_details = df_response_judging[
            (df_response_judging["generator_1"] == fixed_model)
            & (df_response_judging["generator_2"] == st.session_state.selected_model)
            & (
                df_response_judging["instruction"]
                == st.session_state.selected_instruction
            )
        ]

        # if not judging_details.empty:
        if not judging_details["preference"].empty:
            preference = get_preference(judging_details["preference"])
            if preference == "[1>2]":
                st.write(
                    f"**{fixed_model}** is better than **{st.session_state.selected_model}**"
                )
            else:
                st.write(
                    f"**{st.session_state.selected_model}** is better than **{fixed_model}**"
                )
            st.write(
                f"- **Score:** {judging_details['preference'].round(2).item()}\n- **Assigned preference:** {preference}"
            )

            with st.expander("Additional information"):
                st.write(
                    judging_details[
                        [
                            "instruction",
                            "time_per_example",
                            "price_per_example",
                            "raw_completion",
                        ]
                    ]
                )

        # Create two columns for model selectors
        st.markdown("#### Responses")
        col1, col2 = st.columns(2)

        with col1:
            st.selectbox(
                "Reference model",
                [fixed_model],
                key="fixed_model",
            )

            # Get the response string for the fixed model
            if st.session_state.selected_instruction:
                preference = get_preference(judging_details["preference"])
                response_details_fixed = df_response_judging[
                    (
                        df_response_judging["instruction"]
                        == st.session_state.selected_instruction
                    )
                    & (df_response_judging["generator_1"] == fixed_model)
                ].iloc[0]

                st.write(
                    f'Number of words: {len(response_details_fixed["output_1"].split())}'
                )

                # Display the response string
                if preference == "[1>2]":
                    st.success(response_details_fixed["output_1"])
                else:
                    st.error(response_details_fixed["output_1"])

        with col2:
            st.session_state.selected_model = st.selectbox(
                "Select Model",
                model_options,
                key="model_selector",
                on_change=update_model,
                index=(
                    model_options.index(st.session_state.selected_model)
                    if st.session_state.selected_model
                    else 0
                ),
            )

            # Get the response string for the selected model
            if (
                st.session_state.selected_model
                and st.session_state.selected_instruction
            ):
                response_details_dynamic = df_response_judging[
                    (
                        df_response_judging["instruction"]
                        == st.session_state.selected_instruction
                    )
                    & (
                        df_response_judging["generator_2"]
                        == st.session_state.selected_model
                    )
                ].iloc[0]

                st.write(
                    f'Number of words: {len(response_details_dynamic["output_2"].split())}'
                )

                # Display the response string
                if preference == "[2>1]":
                    st.success(response_details_dynamic["output_2"])
                else:
                    st.error(response_details_dynamic["output_2"])

    with outer_tabs[2]:
        st.markdown(
            """The original [AlpacaFarm paper](https://arxiv.org/abs/2305.14387) includes a release of 20K human preferences between a given and reference model on the AlpacaFarm evaluation set. 2.5K of these are cross-annotations (4 humans annotating the same 650 examples). This tab allows you to explore the **human cross-annotations** in more detail."""
        )

        st.markdown("#### Choose example")
        st.button(
            ":game_die: Randomize!",
            on_click=randomize_selection_human_annotations,
            type="primary",
            key="randomize_button_human_annotations",
        )

        left_col, right_col = st.columns([1, 3])

        st.session_state.selected_dataset_human_annotations = left_col.selectbox(
            "Select Dataset",
            ["all"] + df_human_annotations["dataset"].dropna().unique().tolist(),
            key="dataset_selector_human_annotations",
            on_change=update_instruction_options_human_annotations,
        )
        update_instruction_options_human_annotations()
        st.session_state.selected_instruction_human_annotations = right_col.selectbox(
            f"Select Instruction ({len(st.session_state.instruction_options_human_annotations)} unique instructions)",
            st.session_state.instruction_options_human_annotations,
            key="instruction_selector_human_annotations",
            on_change=update_instruction,
            index=(
                st.session_state.instruction_options_human_annotations.index(
                    st.session_state.selected_instruction_human_annotations
                )
                if st.session_state.selected_instruction_human_annotations
                in st.session_state.instruction_options_human_annotations
                else 0
            ),
        )

        st.divider()

        st.markdown(f"#### Selected instruction")
        st.info(st.session_state.selected_instruction_human_annotations)

        st.divider()

        # Need an output column?

        st.markdown("#### Responses")
        col1, col2 = st.columns(2)

        with col1:
            st.selectbox(
                "Output 1 (reference)",
                df_human_annotations.loc[
                    df_human_annotations["instruction"]
                    == st.session_state.selected_instruction_human_annotations
                ]["output_1"]
                .unique()
                .tolist(),
                key="output_selector_human_annotations_fuxed",
                index=0,
                # label_visibility="collapsed",
            )

            # Get the response string for the fixed model
            if st.session_state.selected_instruction_human_annotations:
                response_details_fixed = df_human_annotations[
                    (
                        df_human_annotations["instruction"]
                        == st.session_state.selected_instruction_human_annotations
                    )
                ].iloc[0]

                st.write(
                    f'Number of words: {len(response_details_fixed["output_1"].split())}'
                )

                # Display the response string
                st.info(response_details_fixed["output_1"])

        with col2:
            st.session_state.selected_output_human_annotations = st.selectbox(
                "Output 2",
                df_human_annotations.loc[
                    df_human_annotations["instruction"]
                    == st.session_state.selected_instruction_human_annotations
                ]["output_2"]
                .dropna()
                .tolist(),
                key="output_selector_human_annotations",
                index=0,
                # label_visibility="collapsed",
            )

            # Get the response string for the selected model
            if (
                st.session_state.selected_output_human_annotations
                and st.session_state.selected_instruction_human_annotations
            ):
                response_details_dynamic = df_human_annotations[
                    (
                        df_human_annotations["instruction"]
                        == st.session_state.selected_instruction_human_annotations
                    )
                    & (
                        df_human_annotations["output_2"]
                        == st.session_state.selected_output_human_annotations
                    )
                ].iloc[0]

                st.write(
                    f'Number of words: {len(response_details_dynamic["output_2"].split())}'
                )
                st.info(response_details_dynamic["output_2"])

        st.divider()

        # Judging details.
        st.markdown(f"#### Human Judging")

        col1, col2 = st.columns(2)

        with col1:

            judging_details = df_human_annotations[
                (df_human_annotations["output_1"] == response_details_fixed["output_1"])
                & (
                    df_human_annotations["output_2"]
                    == response_details_dynamic["output_2"]
                )
            ]
            judging_details["assigned_preference"] = judging_details[
                "preference"
            ].apply(get_preference_from_rounded_score)
            is_unanimous_value = is_unanimous(judging_details["preference"])
            st.write("**Unanimous?** ", is_unanimous_value)

            # Draw a histogram of preference.
            fig = px.histogram(
                judging_details,
                x="assigned_preference",
            )
            fig.update_layout(xaxis_title="Preference")
            st.plotly_chart(fig)

            with st.expander("Data details"):
                st.dataframe(
                    judging_details[["annotator_index", "assigned_preference"]],
                    hide_index=True,
                )

        # Generate the heatmap figure
        with col2:
            agreement_matrix = create_agreement_heatmap(df_human_annotations)
            # st.write(
            #     f"**Overall interannotator agreement:** {agreement_matrix.mean().mean():.3f}"
            # )
            with st.expander(
                f"**Overall interannotator agreement:** {agreement_matrix.mean().mean():.3f}"
            ):
                st.pyplot(plt)


if __name__ == "__main__":
    app()