from __future__ import annotations

import gradio as gr
import pandas as pd


ALL_OPTION = "All"
CATEGORY_HEADER_LABELS = {
    "Art & Entertainment": "Art &\nEntert.",
    "Art And Entertainment": "Art &\nEntert.",
    "Culture & Tradition": "Culture &\nTrad.",
    "Culture And Tradition": "Culture &\nTrad.",
    "Geography & Nature": "Geogr. &\nNature",
    "Geography And Nature": "Geogr. &\nNature",
    "History & Society": "History &\nSociety",
    "History And Society": "History &\nSociety",
    "Image Understanding": "Image\nUnd.",
    "Visual Reasoning": "Visual\nReas.",
}


def score_cell_style(value: object) -> str:
    if not isinstance(value, (int, float)) or pd.isna(value) or value < 50:
        return ""

    ratio = min((float(value) - 50.0) / 50.0, 1.0)
    if ratio < 0.2:
        background = "#edf7c2"
        color = "#0f172a"
    elif ratio < 0.4:
        background = "#d7ef7a"
        color = "#0f172a"
    elif ratio < 0.6:
        background = "#b5d94c"
        color = "#0f172a"
    elif ratio < 0.8:
        background = "#7fae2e"
        color = "#ffffff"
    else:
        background = "#4f7f18"
        color = "#ffffff"
    return f"background-color: {background}; color: {color}; font-weight: 600;"


def format_metric_value(value: object) -> str:
    if isinstance(value, (int, float)) and pd.notna(value):
        return f"{float(value):.2f}"
    return ""


def style_leaderboard(frame: pd.DataFrame, display_columns: list[str]):
    metric_columns = [column for column in display_columns if column not in {"Org", "Model", "Size"}]
    styled_frame = frame[display_columns].copy()
    return styled_frame.style.format({column: format_metric_value for column in metric_columns}, na_rep="").map(
        score_cell_style, subset=metric_columns
    )


def visible_columns_for_category(
    category_columns: list[str],
    subcategory_groups: dict[str, list[tuple[str, str]]],
    selected_category: str,
) -> tuple[list[str], dict[str, str]]:
    base_columns = ["Org", "Model", "Size", "Overall"]
    if selected_category == "Overall":
        renamed_columns = {
            column: CATEGORY_HEADER_LABELS[column]
            for column in category_columns
            if column in CATEGORY_HEADER_LABELS
        }
        return base_columns + category_columns, renamed_columns

    subcategory_columns = subcategory_groups.get(selected_category, [])
    display_columns = base_columns.copy()
    renamed_columns: dict[str, str] = {}

    for column_key, column_label in subcategory_columns:
        if column_label == "Unknown":
            display_columns.append(selected_category)
            renamed_columns[selected_category] = selected_category
            continue
        display_columns.append(column_key)
        renamed_columns[column_key] = column_label

    return display_columns, renamed_columns


def filter_leaderboard(
    leaderboard_df: pd.DataFrame,
    display_columns: list[str],
    search_query: str,
    selected_split: str,
    selected_types: list[str],
    selected_size_ranges: list[str],
    selected_families: list[str],
) -> pd.DataFrame:
    filtered = leaderboard_df.drop(columns=["Rank"], errors="ignore").copy()
    selected_types = selected_filter_values(selected_types)
    selected_size_ranges = selected_filter_values(selected_size_ranges)
    selected_families = selected_filter_values(selected_families)

    query = (search_query or "").strip().lower()
    if query:
        filtered = filtered[filtered["Model"].str.lower().str.contains(query, na=False)]

    if selected_split:
        filtered = filtered[filtered["Split"] == selected_split]

    if selected_types:
        filtered = filtered[filtered["Type"].isin(selected_types)]

    if selected_size_ranges:
        filtered = filtered[filtered["Size range"].isin(selected_size_ranges)]

    if selected_families:
        filtered = filtered[filtered["Family"].isin(selected_families)]

    metric_columns = [column for column in display_columns if column not in {"Split", "Org", "Model", "Size"}]
    filtered = filtered.sort_values(by=metric_columns, ascending=False).reset_index(drop=True)
    return filtered


def selected_filter_values(values: list[str] | None) -> list[str]:
    if not values:
        return []
    return [value for value in values if value != ALL_OPTION]


def render_leaderboard_filters(leaderboard_df: pd.DataFrame):
    type_options = ["Open-weight", "Proprietary", "Random"]
    size_options = [
        size_range
        for size_range in ["1-10B", "10-20B", "20-40B", "30-100B", ">100B"]
        if size_range in set(leaderboard_df["Size range"].dropna().astype(str))
    ]
    family_options = sorted(leaderboard_df["Family"].dropna().astype(str).unique().tolist())
    available_types = [option for option in type_options if option in set(leaderboard_df["Type"].dropna().astype(str))]

    with gr.Row(elem_id="leaderboard-search-row"):
        search_box = gr.Textbox(
            value="",
            placeholder="Search by model name...",
            show_label=False,
            scale=3,
            elem_id="leaderboard-search",
        )

    filter_label_style = (
        "color: #181416; font-size: 0.86rem; font-weight: 700; "
        "line-height: 1.2; margin: 0 0 8px;"
    )

    with gr.Row(elem_id="leaderboard-filter-row"):
        with gr.Column(elem_classes=["povisle-filter-field"]):
            gr.HTML(f'<div style="{filter_label_style}">Model Type</div>')
            type_filter = gr.Dropdown(
                choices=[ALL_OPTION, *available_types],
                value=[ALL_OPTION],
                multiselect=True,
                allow_custom_value=False,
                label="Model Type",
                show_label=False,
                elem_id="type-filter",
            )
        with gr.Column(elem_classes=["povisle-filter-field"]):
            gr.HTML(f'<div style="{filter_label_style}">Size</div>')
            size_filter = gr.Dropdown(
                choices=[ALL_OPTION, *size_options],
                value=[ALL_OPTION],
                multiselect=True,
                allow_custom_value=False,
                label="Size",
                show_label=False,
                elem_id="size-filter",
            )
        with gr.Column(elem_classes=["povisle-filter-field"]):
            gr.HTML(f'<div style="{filter_label_style}">Family</div>')
            family_filter = gr.Dropdown(
                choices=[ALL_OPTION, *family_options],
                value=[ALL_OPTION],
                multiselect=True,
                allow_custom_value=False,
                label="Family",
                show_label=False,
                elem_id="family-filter",
            )

    return search_box, type_filter, size_filter, family_filter


def render_leaderboard(
    leaderboard_df: pd.DataFrame,
    category_columns: list[str],
    subcategory_groups: dict[str, list[tuple[str, str]]],
    selected_category: str,
    search_query: str,
    selected_split: str,
    selected_types: list[str],
    selected_size_ranges: list[str],
    selected_families: list[str],
) -> tuple[object, list[str]]:
    display_columns, renamed_columns = visible_columns_for_category(
        category_columns,
        subcategory_groups,
        selected_category,
    )
    filtered = filter_leaderboard(
        leaderboard_df,
        display_columns,
        search_query,
        selected_split,
        selected_types,
        selected_size_ranges,
        selected_families,
    )
    styled = style_leaderboard(
        filtered.rename(columns=renamed_columns),
        [renamed_columns.get(column, column) for column in display_columns],
    )
    return styled


def update_leaderboard(
    leaderboard_df: pd.DataFrame,
    category_columns: list[str],
    subcategory_groups: dict[str, list[tuple[str, str]]],
    selected_category: str,
    search_query: str,
    selected_split: str,
    selected_types: list[str],
    selected_size_ranges: list[str],
    selected_families: list[str],
):
    styled = render_leaderboard(
        leaderboard_df,
        category_columns,
        subcategory_groups,
        selected_category,
        search_query,
        selected_split,
        selected_types,
        selected_size_ranges,
        selected_families,
    )
    return styled


def update_split_leaderboards(
    leaderboard_df: pd.DataFrame,
    category_columns: list[str],
    subcategory_groups: dict[str, list[tuple[str, str]]],
    selected_category: str,
    search_query: str,
    selected_types: list[str],
    selected_size_ranges: list[str],
    selected_families: list[str],
):
    return (
        render_leaderboard(
            leaderboard_df,
            category_columns,
            subcategory_groups,
            selected_category,
            search_query,
            "Test",
            selected_types,
            selected_size_ranges,
            selected_families,
        ),
        render_leaderboard(
            leaderboard_df,
            category_columns,
            subcategory_groups,
            selected_category,
            search_query,
            "Validation",
            selected_types,
            selected_size_ranges,
            selected_families,
        ),
    )


def render_leaderboard_category_tab(
    leaderboard_df: pd.DataFrame,
    category_columns: list[str],
    subcategory_groups: dict[str, list[tuple[str, str]]],
    search_box,
    type_filter,
    size_filter,
    family_filter,
) -> None:
    with gr.Row(elem_id="leaderboard-category-row"):
        category_view_select = gr.Radio(
            choices=["Overall", *[category for category in category_columns if subcategory_groups.get(category)]],
            value="Overall",
            show_label=False,
            container=True,
            elem_id="category-view-select",
        )
    initial_test_table_value = render_leaderboard(
        leaderboard_df,
        category_columns,
        subcategory_groups,
        "Overall",
        "",
        "Test",
        [],
        [],
        [],
    )
    initial_validation_table_value = render_leaderboard(
        leaderboard_df,
        category_columns,
        subcategory_groups,
        "Overall",
        "",
        "Validation",
        [],
        [],
        [],
    )
    with gr.Tabs(elem_id="leaderboard-category-split-tabs", elem_classes=["leaderboard-split-tabs"]):
        with gr.Tab("Test"):
            test_leaderboard_table = gr.Dataframe(
                value=initial_test_table_value,
                max_height=760,
                interactive=False,
                wrap=True,
                buttons=[],
                elem_id="leaderboard-table-test",
                elem_classes=["leaderboard-table"],
            )
        with gr.Tab("Validation"):
            validation_leaderboard_table = gr.Dataframe(
                value=initial_validation_table_value,
                max_height=760,
                interactive=False,
                wrap=True,
                buttons=[],
                elem_id="leaderboard-table-validation",
                elem_classes=["leaderboard-table"],
            )

    search_box.submit(
        fn=lambda selected_category, search_query, selected_type, selected_size, selected_family: update_split_leaderboards(
            leaderboard_df,
            category_columns,
            subcategory_groups,
            selected_category,
            search_query,
            selected_type,
            selected_size,
            selected_family,
        ),
        inputs=[category_view_select, search_box, type_filter, size_filter, family_filter],
        outputs=[test_leaderboard_table, validation_leaderboard_table],
    )

    for component in (type_filter, size_filter, family_filter, category_view_select):
        component.change(
            fn=lambda selected_category, search_query, selected_type, selected_size, selected_family: update_split_leaderboards(
                leaderboard_df,
                category_columns,
                subcategory_groups,
                selected_category,
                search_query,
                selected_type,
                selected_size,
                selected_family,
            ),
            inputs=[category_view_select, search_box, type_filter, size_filter, family_filter],
            outputs=[test_leaderboard_table, validation_leaderboard_table],
        )