import gradio as gr
import pandas as pd

# Import configurations and informational texts
from src.config import (
    CITATIONS,
    INFO_BENCHMARK_TASK,
    INFO_GOTO_AI,
    INFO_SCORE_CALCULATION,
    INTRODUCTION_TEXT,
    TITLE,
    file_path,
    hidden_tabs,
    model_types,
)

# Import function to load leaderboard tables
from src.populate import load_tables


FIXED_COL_WIDTHS = {"Model": "260px", "Type": "90px", "Size": "70px", "Total": "90px"}


def column_widths_for(columns: list[str]) -> list[str]:
    return [FIXED_COL_WIDTHS.get(col, "130px") for col in columns]


def make_filter_fn(df: pd.DataFrame):
    def filter_table(search: str, sizes: list[str]) -> pd.DataFrame:
        filtered = df
        if sizes:
            filtered = filtered[filtered["Size"].isin(sizes)]
        if search:
            filtered = filtered[filtered["Model"].str.contains(search, case=False, na=False)]
        return filtered

    return filter_table


# Create a Gradio application with block-based UI
# 'Blocks()' is used to group multiple components in a single interface
demo = gr.Blocks()
with demo:
    gr.HTML(TITLE)  # Display the main title of the application
    gr.Markdown(INTRODUCTION_TEXT, elem_classes="markdown-text")  # Display introductory text

    # Create tabs to display different leaderboard tables
    with gr.Tabs(elem_classes="tab-buttons") as tabs:
        tables = load_tables(file_path)  # Load leaderboard data from file
        for model_type in model_types:
            with gr.TabItem(model_type, elem_id="llm-benchmark-tab-table", id=model_type):
                for i, t in enumerate(tables):  # Loop through the tables to create tabs
                    if (model_type, t["name"]) in hidden_tabs:
                        continue
                    with gr.TabItem(t["name"], elem_id="llm-benchmark-tab-table", id=i):
                        table_df = t["table"][t["table"]["Type"] == model_type]
                        table_df = table_df.drop(columns=t["hidden_col"], errors="ignore")
                        table_df = table_df.dropna(axis=1, how="all")

                        sizes = sorted(table_df["Size"].dropna().unique().tolist())

                        with gr.Row():
                            search_box = gr.Textbox(label="Search model", placeholder="Search by model name")
                            size_filter = gr.CheckboxGroup(choices=sizes, value=sizes, label="Model sizes")

                        table = gr.Dataframe(
                            value=table_df,
                            interactive=False,
                            wrap=True,
                            column_widths=column_widths_for(list(table_df.columns)),
                        )

                        filter_fn = make_filter_fn(table_df)
                        search_box.change(filter_fn, inputs=[search_box, size_filter], outputs=table)
                        size_filter.change(filter_fn, inputs=[search_box, size_filter], outputs=table)

    # Add additional informational sections using Accordion
    with gr.Row():
        with gr.Accordion("📚 Benchmark Tasks", open=False):
            gr.Markdown(INFO_BENCHMARK_TASK, elem_classes="markdown-text")

    with gr.Row():
        with gr.Accordion("🧮 Score Calculation", open=False):
            gr.Markdown(INFO_SCORE_CALCULATION, elem_classes="markdown-text")

    with gr.Row():
        with gr.Accordion("🤝 About GoTo-AI", open=False):
            gr.Markdown(INFO_GOTO_AI, elem_classes="markdown-text")

    with gr.Row():
        with gr.Accordion("📝 Citations", open=False):
            gr.Markdown(CITATIONS, elem_classes="markdown-text")

# Run the application
demo.launch()