import pandas as pd
import gradio as gr
import plotly.express as px
import plotly.graph_objects as go

df = pd.read_csv("data/csv/models_data.csv")


filter_mapping = {
    "all": "all",
    "🟢 Pre-trained": "🟢",
    "🟩 Continuously pre-trained": "🟩",
    "🔶 Fine-tuned on domain-specific data": "🔶",
    "💬 Chat-models (RLHF, DPO, IFT, ...)": "💬"
}


def filter_items(df, query):
    if query == "all":
        return df
    filter_value = filter_mapping[query]
    return df[df["T"].str.contains(filter_value, na=False)]


def create_scatter_plot(df, x_col, y_col, title, x_title, y_title):
    fig = px.scatter(df, x=x_col, y=y_col, color='Model', title=title)
    fig.add_trace(
        go.Scatter(
            x=[0, 100],
            y=[0, 100],
            mode="lines",
            name="y=x line",
            line=dict(color='black', dash='dash')
        )
    )
    
    fig.update_layout(
        xaxis_title=x_title,
        yaxis_title=y_title,
        xaxis=dict(range=[0, 100]),
        yaxis=dict(range=[0, 100]),
        legend_title_text='Model'
    )
    fig.update_traces(marker=dict(size=10), selector=dict(mode='markers'))
    return fig


with gr.Blocks(css="custom.css") as demo:
    with gr.Column():
        gr.Markdown(
            """<div style="text-align: center;"><h1>🐰 RABBITS:</h1>
            <h2 style='color: #6aa84f;'>R</span>obust <span style='color: #6aa84f;'>A</span>ssessment of <span style='color: #6aa84f;'>B</span>iomedical <span style='color: #6aa84f;'>B</span>enchmarks <span style='color: #6aa84f;'>I</span>nvolving drug <span style='color: #6aa84f;'>T</span>erm <span style='color: #6aa84f;'>S</span>ubstitutions for Language Models</h2></div>"""
        )
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown(
                    """<p class='markdown-text'>Robust language models are crucial in the medical domain and the RABBITS project tests the robustness of LLMs by evaluating their handling of synonyms, specifically brand and generic drug names. We assessed 16 open-source language models from Hugging Face using systematic synonym substitution on MedQA and MedMCQA tasks. Our results show a consistent decline in performance across all model sizes, highlighting challenges in synonym comprehension. Additionally, we discovered significant dataset contamination by identifying overlaps between MedQA, MedMCQA test sets, and the Dolma 1.6 dataset using an 8-gram analysis. This highlights the need to improve model robustness and address contamination in open-source datasets.</p>"""
                )
            with gr.Column(scale=1):
                gr.Image(value="workflow.svg").style(width=256, height=256)

    with gr.Tabs(elem_classes="tab-buttons"):
        with gr.TabItem("🔍 Evaluation table"):
            with gr.Column():
                with gr.Accordion("➡️ Filter by Column", open=False):
                    shown_columns = gr.CheckboxGroup(
                        choices=df.columns.tolist(),
                        value=df.columns.tolist(),
                        label="Select Columns",
                        interactive=True,
                    )
                with gr.Row():
                    search_bar = gr.Textbox(
                        placeholder="🔍 Search for your model and press ENTER...",
                        show_label=False,
                        elem_id="search-bar"
                    )
                    filter_columns = gr.Radio(
                        label="⏚ Filter model types",
                        choices=[
                            "all", 
                            "🟢 Pre-trained", 
                            "🟩 Continuously pre-trained", 
                            "🔶 Fine-tuned on domain-specific data", 
                            "💬 Chat-models (RLHF, DPO, IFT, ...)"
                        ],
                        value="all",
                        elem_id="filter-columns",
                    )
                leaderboard_df = gr.Dataframe(
                    value=df,
                    headers="keys",
                    datatype=["html" if col == "Model" else "str" for col in df.columns],
                    interactive=False,
                    elem_id="leaderboard-table"
                )

                def update_leaderboard(search_query):
                    filtered_df = df[df["Model"].str.contains(search_query, case=False)]
                    return filtered_df

                search_bar.submit(
                    update_leaderboard,
                    inputs=search_bar,
                    outputs=leaderboard_df
                )

                def filter_update(query):
                    filtered_df = filter_items(df, query)
                    return filtered_df

                filter_columns.change(
                    filter_update,
                    inputs=filter_columns,
                    outputs=leaderboard_df
                )

                shown_columns.change(
                    lambda cols: df[cols], 
                    inputs=shown_columns, 
                    outputs=leaderboard_df
                )

        with gr.TabItem("📊 Evaluation Plots"):
            with gr.Column():
                with gr.Row():
                    scatter1 = gr.Plot(
                        value=create_scatter_plot(df, "medmcqa_orig_filtered", "medmcqa_g2b",
                                                "MedMCQA: Orig vs G2B", "medmcqa_orig_filtered", "medmcqa_g2b"),
                        elem_id="scatter1"
                    )
                    scatter2 = gr.Plot(
                        value=create_scatter_plot(df, "medqa_4options_orig_filtered", "medqa_4options_g2b",
                                                "MedQA: Orig vs G2B", "medqa_4options_orig_filtered", "medqa_4options_g2b"),
                        elem_id="scatter2"
                    )
                with gr.Row():
                    scatter3 = gr.Plot(
                        value=create_scatter_plot(df, "b4bqa", "b4b",
                                                "b4bqa vs b4b", "b4bqa", "b4b"),
                        elem_id="scatter3"
                    )

        with gr.TabItem("📝 About"):
            gr.Markdown(
                """<div style="text-align: center;">
                <h2>About RABBITS LLM Leaderboard</h2>
                <p>This leaderboard ...</p>
                <p>It is designed to ...</p>
                </div>""",
                elem_classes="markdown-text"
            )

        with gr.TabItem("🚀 Submit Here!"):
            gr.Markdown(
                """<div style="text-align: center;">
                <h2>Submit Your Model Results</h2>
                <p>If you have new model results that you would like to add to the leaderboard, please follow the submission guidelines below:</p>
                <ul>
                    <li>COMING SOON</li>
                </ul>
                <p>COMING SOON</p>
                </div>""",
                elem_classes="markdown-text"
            )

if __name__ == "__main__":
    demo.launch()