app.py

"""DNA Foundation Model Benchmark Leaderboard — Gradio Blocks app."""

import gradio as gr
import pandas as pd

from src.constants import CLASSIFICATION_METRICS
from src.data import apply_filters, deduplicate, load_data
from src.plots import (
    bar_plot_per_category,
    bar_plot_per_task,
    bubble_context_length,
    efficiency_table,
    heatmap_variants,
    ranking_bar,
    scatter_speed,
    strip_plot,
    top_n_variants,
)

# ---------------------------------------------------------------------------
# Startup: load & deduplicate once
# ---------------------------------------------------------------------------

print("Loading benchmark data from HF Hub …")
try:
    _LOADED_DF = load_data()
    _RAW_DF = deduplicate(_LOADED_DF)
    print(f"  Loaded {len(_RAW_DF)} rows, {_RAW_DF['model_alias'].nunique()} models.")
except Exception as exc:
    print(f"  WARNING: Could not load data — {exc}")
    _LOADED_DF = pd.DataFrame()
    _RAW_DF = pd.DataFrame()

_METRIC_CHOICES = list(CLASSIFICATION_METRICS.values())
_METRIC_KEYS = list(CLASSIFICATION_METRICS.keys())

_CATEGORY_CHOICES = sorted(_RAW_DF["task_category"].unique().tolist()) if not _RAW_DF.empty else []

_SPEED_COL_META = [
    ("test_throughput_seq_s",   "Throughput (seq/s)"),
    ("test_embedding_time_s",   "Embedding Time (s)"),
    ("vram_model_mb",           "VRAM (MB)"),
]
_SPEED_CHOICES = [lbl for col, lbl in _SPEED_COL_META if not _RAW_DF.empty and col in _RAW_DF.columns]
_SPEED_COL_MAP = {lbl: col for col, lbl in _SPEED_COL_META}

_DATASET_CHOICES = (
    sorted(_RAW_DF["dataset_name"].dropna().unique().tolist())
    if not _RAW_DF.empty and "dataset_name" in _RAW_DF.columns
    else []
)


def _metric_key(label: str) -> str:
    if not label or label not in _METRIC_CHOICES:
        return _METRIC_KEYS[0]
    return _METRIC_KEYS[_METRIC_CHOICES.index(label)]


def _numeric_range(df: pd.DataFrame, col: str) -> tuple[float, float]:
    if col not in df.columns or df.empty:
        return 0.0, 1.0
    lo, hi = float(df[col].min()), float(df[col].max())
    if lo == hi:
        return lo, lo + 1.0
    return lo, hi


# ---------------------------------------------------------------------------
# Render helpers
# ---------------------------------------------------------------------------


def render_leaderboard(df: pd.DataFrame, metric_label: str, top_n: int = 20):
    key = _metric_key(metric_label)
    df = top_n_variants(df, key, int(top_n))
    return ranking_bar(df, key), strip_plot(df, key)


def render_per_category(df: pd.DataFrame, metric_label: str, top_n: int = 20, datasets: list[str] | None = None):
    if datasets and "dataset_name" in df.columns:
        df = df[df["dataset_name"].isin(datasets)]
    key = _metric_key(metric_label)
    df = top_n_variants(df, key, int(top_n))
    return bar_plot_per_category(df, key)


def render_per_task(df: pd.DataFrame, metric_label: str, category: str, top_n: int = 20, datasets: list[str] | None = None):
    if datasets and "dataset_name" in df.columns:
        df = df[df["dataset_name"].isin(datasets)]
    key = _metric_key(metric_label)
    df = top_n_variants(df, key, int(top_n))
    return bar_plot_per_task(df, key, category)


def render_speed(
    df: pd.DataFrame,
    metric_label: str,
    speed_label: str | None,
    log_scale: bool,
    top_n: int = 20,
):
    key = _metric_key(metric_label)
    df = top_n_variants(df, key, int(top_n))
    col = _SPEED_COL_MAP.get(speed_label) if speed_label else None
    return scatter_speed(df, key, speed_col=col, log_scale=log_scale), efficiency_table(df, key)


def render_context(df: pd.DataFrame, metric_label: str, top_n: int = 20):
    key = _metric_key(metric_label)
    df = top_n_variants(df, key, int(top_n))
    return bubble_context_length(df, key)


def render_variant_heatmap(
    df: pd.DataFrame,
    metric_label: str,
    model_alias: list[str] | str,
    aggregate_by_dataset: bool = False,
):
    key = _metric_key(metric_label)
    return heatmap_variants(df, key, model_alias, aggregate_by_dataset)


def update_category_choices(df: pd.DataFrame, datasets: list[str] | None = None):
    if datasets and "dataset_name" in df.columns:
        df = df[df["dataset_name"].isin(datasets)]
    cats = sorted(df["task_category"].unique().tolist()) if not df.empty else []
    return gr.update(choices=cats, value=cats[0] if cats else None)


def apply_and_preview(
    df: pd.DataFrame,
    prod_only: bool,
    exclude_models: list[str],
    params_min: float,
    params_max: float,
    dim_min: float,
    dim_max: float,
    ctx_min: float,
    ctx_max: float,
    pooling_strategies: list[str] | None = None,
) -> tuple[pd.DataFrame, pd.DataFrame, str]:
    filters = {
        "subsample_train": "full" if prod_only else None,
        "exclude_models": exclude_models or [],
        "model_params_min": params_min,
        "model_params_max": params_max,
        "embedding_dim_min": dim_min,
        "embedding_dim_max": dim_max,
        "max_context_size_min": ctx_min,
        "max_context_size_max": ctx_max,
        "pooling_strategies": pooling_strategies or [],
    }
    result = apply_filters(df, filters)
    preview = result.head(50) if not result.empty else pd.DataFrame()
    return result, preview, f"{len(result)} rows active"


def _hash_label(row: pd.Series) -> str:
    """Human-readable label for an embedding config hash dropdown entry."""

    def _s(val) -> str:
        return "" if val is None or (isinstance(val, float) and pd.isna(val)) else str(val)

    category = _s(row.get("task_category"))
    task = _s(row.get("task_name"))
    pool = _s(row.get("pooling_strategy"))
    strategy = _s(row.get("layer_selection_strategy"))
    idx = row.get("layer_index")
    layer = strategy + (f"({int(idx)})" if idx is not None and not pd.isna(idx) else "")
    h = _s(row.get("embedding_config_hash"))[:8]

    parts = [f"{category}/{task}" if category else task, pool, layer]
    return "  ·  ".join(p for p in parts if p) + f"  ({h})"


def get_hash_choices_for_model(loaded_df: pd.DataFrame, model: str) -> list[tuple[str, str]]:
    if not model or loaded_df.empty or "embedding_config_hash" not in loaded_df.columns:
        return []
    sub = loaded_df[loaded_df["model_alias"] == model].drop_duplicates("embedding_config_hash")
    return sorted(
        [(_hash_label(row), row["embedding_config_hash"]) for _, row in sub.iterrows()],
        key=lambda x: x[0],
    )


def get_head_types_for_hash(loaded_df: pd.DataFrame, model: str, key: str) -> list[str]:
    if loaded_df.empty or "head_type" not in loaded_df.columns:
        return []
    sub = loaded_df
    if model:
        sub = sub[sub["model_alias"] == model]
    if key and "embedding_config_hash" in loaded_df.columns:
        sub = sub[sub["embedding_config_hash"] == key]
    return sorted(sub["head_type"].dropna().unique().tolist())


_DEFAULT_INSP_COLS = ["kept", "run_id", "run_at", "subsample_train", "task_name", "mcc_test", "accuracy_test"]


def build_inspector_table(
    loaded_df: pd.DataFrame,
    raw_df: pd.DataFrame,
    model: str,
    key: str,
    head_type: str,
    selected_cols: list[str] | None = None,
    task_cats: list[str] | None = None,
    tasks: list[str] | None = None,
    pooling_strategies: list[str] | None = None,
) -> pd.DataFrame:
    if loaded_df.empty:
        return pd.DataFrame()
    masks = []
    if model:
        masks.append(loaded_df["model_alias"] == model)
    if key and "embedding_config_hash" in loaded_df.columns:
        masks.append(loaded_df["embedding_config_hash"] == key)
    if head_type and "head_type" in loaded_df.columns:
        masks.append(loaded_df["head_type"] == head_type)
    if task_cats and "task_category" in loaded_df.columns:
        masks.append(loaded_df["task_category"].isin(task_cats))
    if tasks and "task_name" in loaded_df.columns:
        masks.append(loaded_df["task_name"].isin(tasks))
    if pooling_strategies and "pooling_strategy" in loaded_df.columns:
        masks.append(loaded_df["pooling_strategy"].isin(pooling_strategies))
    if not masks:
        return pd.DataFrame()
    combined = masks[0]
    for m in masks[1:]:
        combined = combined & m
    sub = loaded_df[combined].copy()
    kept_ids = set(raw_df["run_id"]) if "run_id" in raw_df.columns else set()
    sub["kept"] = sub["run_id"].isin(kept_ids).map({True: "✓", False: "✗"})
    cols = selected_cols if selected_cols else _DEFAULT_INSP_COLS
    return sub[[c for c in cols if c in sub.columns or c == "kept"]].reset_index(drop=True)


def exclude_run_ids(df: pd.DataFrame, run_ids_text: str) -> tuple[pd.DataFrame, pd.DataFrame, str]:
    ids = [r.strip() for r in run_ids_text.split(",") if r.strip()]
    result = apply_filters(df, {"exclude_run_ids": ids})
    preview = result.head(50) if not result.empty else pd.DataFrame()
    return result, preview, f"{len(result)} rows active"


def _filter_status_banner(filtered_df: pd.DataFrame, raw_df: pd.DataFrame) -> str:
    if filtered_df.empty or raw_df.empty:
        return ""
    if len(filtered_df) < len(raw_df):
        return (
            f"⚠️ **Active scope filters** — {len(filtered_df):,} / {len(raw_df):,} rows visible. "
            "Go to 🔬 Scope / Filter to review."
        )
    return ""


# ---------------------------------------------------------------------------
# Build UI
# ---------------------------------------------------------------------------

all_models = sorted(_RAW_DF["model_alias"].unique().tolist()) if not _RAW_DF.empty else []
params_min_v, params_max_v = _numeric_range(_RAW_DF, "model_params")
dim_min_v, dim_max_v = _numeric_range(_RAW_DF, "embedding_dim")
ctx_min_v, ctx_max_v = _numeric_range(_RAW_DF, "max_context_size")

_insp_default_model = all_models[0] if all_models else None
_insp_default_keys = get_hash_choices_for_model(_LOADED_DF, _insp_default_model)
_insp_default_heads = get_head_types_for_hash(_LOADED_DF, _insp_default_model, None)
_ALL_INSP_COLS = ["kept"] + _LOADED_DF.columns.tolist() if not _LOADED_DF.empty else _DEFAULT_INSP_COLS
_all_task_cats = (
    sorted(_LOADED_DF["task_category"].dropna().unique().tolist())
    if not _LOADED_DF.empty and "task_category" in _LOADED_DF.columns
    else []
)
_all_tasks = (
    sorted(_LOADED_DF["task_name"].dropna().unique().tolist())
    if not _LOADED_DF.empty and "task_name" in _LOADED_DF.columns
    else []
)
_all_pooling_strategies = (
    sorted(_LOADED_DF["pooling_strategy"].dropna().unique().tolist())
    if not _LOADED_DF.empty and "pooling_strategy" in _LOADED_DF.columns
    else []
)

with gr.Blocks(title="DNA Benchmark Leaderboard", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# 🧬 DNA Foundation Model Benchmark Leaderboard\n"
        "Compare DNA language models across genomics classification tasks."
    )

    filter_banner = gr.Markdown("")

    # ── Global controls ───────────────────────────────────────────────────────
    with gr.Row():
        metric_dd = gr.Dropdown(choices=_METRIC_CHOICES, value=_METRIC_CHOICES[0], label="Metric", scale=2)
        top_n_sl = gr.Slider(minimum=5, maximum=50, step=5, value=20, label="Top N variants", scale=1)

    # ── Shared state ──────────────────────────────────────────────────────────
    loaded_df_state = gr.State(_LOADED_DF)
    raw_df_state = gr.State(_RAW_DF)
    filtered_df_state = gr.State(_RAW_DF)

    # ── Tabs ──────────────────────────────────────────────────────────────────
    with gr.Tabs() as tabs:
        # ── Tab 1: Leaderboard ────────────────────────────────────────────────
        with gr.Tab("🏅 Leaderboard"):
            leaderboard_fig = gr.Plot(label="Variant ranking (best at top)")
            violin_fig = gr.Plot(label="Per-task scores by variant (rank 1 at top)")

            gr.Markdown("---\n### Drill into model variants")
            model_dd = gr.Dropdown(
                choices=all_models,
                value=[all_models[0]] if all_models else [],
                multiselect=True,
                label="Select model(s) — single model: all variants as rows; multiple: best variant per family",
            )
            agg_dataset_cb = gr.Checkbox(
                label="Aggregate by dataset (show mean per benchmark dataset instead of per task)",
                value=False,
            )
            variant_violin_fig = gr.Plot(label="Variant comparison (metric per task)")

        # ── Tab 2: Per-Category Performance ───────────────────────────────────
        with gr.Tab("📊 Per-Category Performance"):
            category_fig = gr.Plot(label="Mean performance per task category")

            gr.Markdown("---\n### Drill into individual tasks")
            category_dd = gr.Dropdown(
                choices=_CATEGORY_CHOICES,
                value=_CATEGORY_CHOICES[0] if _CATEGORY_CHOICES else None,
                label="Select category",
            )
            dataset_dd = gr.Dropdown(
                choices=_DATASET_CHOICES,
                value=[],
                multiselect=True,
                label="Filter by dataset",
            )
            per_task_fig = gr.Plot(label="Performance per task within selected category")

        # ── Tab 3: Speed vs Performance ───────────────────────────────────────
        with gr.Tab("⚡ Speed vs Performance"):
            gr.Markdown(
                "Each point = one variant (mean metric across all tasks vs mean speed across tasks). "
                "**Pareto frontier** = variants not dominated on both axes (black outline)."
            )
            with gr.Row():
                speed_metric_radio = gr.Radio(
                    choices=_SPEED_CHOICES,
                    value=_SPEED_CHOICES[0] if _SPEED_CHOICES else None,
                    label="X-axis metric",
                    scale=3,
                )
                log_scale_cb = gr.Checkbox(value=True, label="Log scale on X-axis", scale=1)
            speed_fig = gr.Plot(label="Performance vs efficiency tradeoff")
            gr.Markdown("### Efficiency ranking")
            efficiency_tbl = gr.Dataframe(
                label="Variants sorted by performance — speed and derived efficiency metrics",
                interactive=False,
            )

        # ── Tab 4: Context Length ─────────────────────────────────────────────
        with gr.Tab("🔬 Context Length"):
            context_fig = gr.Plot(label="Performance vs task sequence length")

        # ── Tab 5: Scope / Filter ─────────────────────────────────────────────
        with gr.Tab("🔬 Scope / Filter"):
            prod_only_cb = gr.Checkbox(
                label="Production runs only (subsample_train IS NULL)",
                value=False,
                info="Currently returns no data — all runs use subsample_train=0.05.",
            )
            exclude_models_ms = gr.Dropdown(choices=all_models, multiselect=True, label="Exclude models", value=[])
            pooling_filter_ms = gr.Dropdown(
                choices=_all_pooling_strategies,
                multiselect=True,
                label="Pooling strategy",
                value=[],
                info="Leave empty to include all pooling strategies.",
            )
            with gr.Row():
                params_min_sl = gr.Slider(
                    minimum=params_min_v, maximum=params_max_v, value=params_min_v, label="Model params (min)"
                )
                params_max_sl = gr.Slider(
                    minimum=params_min_v, maximum=params_max_v, value=params_max_v, label="Model params (max)"
                )
            with gr.Row():
                dim_min_sl = gr.Slider(
                    minimum=dim_min_v, maximum=dim_max_v, value=dim_min_v, label="Embedding dim (min)"
                )
                dim_max_sl = gr.Slider(
                    minimum=dim_min_v, maximum=dim_max_v, value=dim_max_v, label="Embedding dim (max)"
                )
            with gr.Row():
                ctx_min_sl = gr.Slider(
                    minimum=ctx_min_v, maximum=ctx_max_v, value=ctx_min_v, label="Max context size (min bp)"
                )
                ctx_max_sl = gr.Slider(
                    minimum=ctx_min_v, maximum=ctx_max_v, value=ctx_max_v, label="Max context size (max bp)"
                )
            apply_btn = gr.Button("Apply Filters", variant="primary")

            gr.Markdown("---")
            run_ids_tb = gr.Textbox(label="Exclude run_ids (comma-separated)", placeholder="run_id_1, run_id_2, …")
            exclude_runs_btn = gr.Button("Exclude run_ids")

            gr.Markdown("---")
            settings_count = gr.Label(label="Active rows")
            settings_preview = gr.Dataframe(label="Preview (first 50 rows)", interactive=False)

            gr.Markdown(
                "---\n### Raw Run Inspector\nSelect a model and embedding config to audit which runs were kept or dropped by dedup."
            )
            with gr.Row():
                insp_model_dd = gr.Dropdown(choices=all_models, value=_insp_default_model, label="Model", scale=2)
                insp_cache_dd = gr.Dropdown(
                    choices=_insp_default_keys, value=None, label="Embedding config hash", scale=3
                )
                insp_head_dd = gr.Dropdown(choices=_insp_default_heads, value=None, label="Head type", scale=2)
            with gr.Row():
                insp_task_cat_ms = gr.Dropdown(
                    choices=_all_task_cats, value=None, multiselect=True, label="Task category", scale=2
                )
                insp_task_ms = gr.Dropdown(choices=_all_tasks, value=None, multiselect=True, label="Task", scale=3)
                insp_pooling_ms = gr.Dropdown(
                    choices=_all_pooling_strategies, value=None, multiselect=True, label="Pooling strategy", scale=2
                )
            insp_cols_dd = gr.Dropdown(
                choices=_ALL_INSP_COLS,
                value=_DEFAULT_INSP_COLS,
                multiselect=True,
                label="Columns to display",
            )
            insp_table = gr.Dataframe(
                label="Runs for selection  (✓ = kept by dedup, ✗ = dropped)",
                interactive=False,
            )

    # ── Event wiring ──────────────────────────────────────────────────────────

    _leaderboard_triggers = [filtered_df_state, metric_dd, top_n_sl]
    _context_triggers = [filtered_df_state, metric_dd, top_n_sl]
    _speed_triggers = [filtered_df_state, metric_dd, speed_metric_radio, log_scale_cb, top_n_sl]

    for trigger in [filtered_df_state, metric_dd, top_n_sl]:
        trigger.change(render_leaderboard, _leaderboard_triggers, [leaderboard_fig, violin_fig])
        trigger.change(render_context, _context_triggers, context_fig)

    for trigger in [filtered_df_state, metric_dd, speed_metric_radio, log_scale_cb, top_n_sl]:
        trigger.change(render_speed, _speed_triggers, [speed_fig, efficiency_tbl])

    # Category bar chart: reacts to df, metric, top_n, AND dataset filter
    _cat_triggers = [filtered_df_state, metric_dd, top_n_sl, dataset_dd]
    for trigger in [filtered_df_state, metric_dd, top_n_sl, dataset_dd]:
        trigger.change(render_per_category, _cat_triggers, category_fig)

    # Per-task drill-down: reacts to category/dataset dropdowns OR global controls
    _per_task_triggers = [filtered_df_state, metric_dd, category_dd, top_n_sl, dataset_dd]
    for trigger in [filtered_df_state, metric_dd, category_dd, top_n_sl, dataset_dd]:
        trigger.change(render_per_task, _per_task_triggers, per_task_fig)

    # Variant heatmap: reacts to model selector, metric, filtered data, or aggregation toggle
    _variant_triggers = [filtered_df_state, metric_dd, model_dd, agg_dataset_cb]
    for trigger in [filtered_df_state, metric_dd, model_dd, agg_dataset_cb]:
        trigger.change(render_variant_heatmap, _variant_triggers, variant_violin_fig)

    # Update category dropdown choices when global filter OR dataset filter changes
    for trigger in [filtered_df_state, dataset_dd]:
        trigger.change(update_category_choices, [filtered_df_state, dataset_dd], category_dd)

    # Settings — Apply Filters
    apply_btn.click(
        apply_and_preview,
        inputs=[
            raw_df_state,
            prod_only_cb,
            exclude_models_ms,
            params_min_sl,
            params_max_sl,
            dim_min_sl,
            dim_max_sl,
            ctx_min_sl,
            ctx_max_sl,
            pooling_filter_ms,
        ],
        outputs=[filtered_df_state, settings_preview, settings_count],
    ).then(
        _filter_status_banner,
        inputs=[filtered_df_state, raw_df_state],
        outputs=filter_banner,
    )

    # Settings — Exclude run_ids
    exclude_runs_btn.click(
        exclude_run_ids,
        inputs=[filtered_df_state, run_ids_tb],
        outputs=[filtered_df_state, settings_preview, settings_count],
    ).then(
        _filter_status_banner,
        inputs=[filtered_df_state, raw_df_state],
        outputs=filter_banner,
    )

    # Inspector — model updates hash choices; model+hash each update head_type choices independently
    insp_model_dd.change(
        lambda df, m: gr.update(choices=get_hash_choices_for_model(df, m), value=None),
        [loaded_df_state, insp_model_dd],
        insp_cache_dd,
    )
    insp_model_dd.change(
        lambda df, m, k: gr.update(choices=get_head_types_for_hash(df, m, k), value=None),
        [loaded_df_state, insp_model_dd, insp_cache_dd],
        insp_head_dd,
    )
    insp_cache_dd.change(
        lambda df, m, k: gr.update(choices=get_head_types_for_hash(df, m, k), value=None),
        [loaded_df_state, insp_model_dd, insp_cache_dd],
        insp_head_dd,
    )
    _insp_inputs = [
        loaded_df_state,
        raw_df_state,
        insp_model_dd,
        insp_cache_dd,
        insp_head_dd,
        insp_cols_dd,
        insp_task_cat_ms,
        insp_task_ms,
        insp_pooling_ms,
    ]
    for trigger in [
        insp_model_dd,
        insp_cache_dd,
        insp_head_dd,
        insp_cols_dd,
        insp_task_cat_ms,
        insp_task_ms,
        insp_pooling_ms,
    ]:
        trigger.change(build_inspector_table, _insp_inputs, insp_table)

    # Trigger resize on tab switch so Plotly reflows to correct container width
    tabs.select(fn=None, js="() => setTimeout(() => window.dispatchEvent(new Event('resize')), 100)")

    # Initial render on load
    demo.load(render_leaderboard, inputs=_leaderboard_triggers, outputs=[leaderboard_fig, violin_fig])
    demo.load(render_per_category, inputs=_cat_triggers, outputs=category_fig)
    demo.load(render_per_task, inputs=_per_task_triggers, outputs=per_task_fig)
    demo.load(render_speed, inputs=_speed_triggers, outputs=[speed_fig, efficiency_tbl])
    demo.load(render_context, inputs=_context_triggers, outputs=context_fig)
    demo.load(render_variant_heatmap, inputs=_variant_triggers, outputs=variant_violin_fig)


if __name__ == "__main__":
    demo.launch(share=False, server_name="0.0.0.0")