util/render.py

"""Build per-pillar HTML leaderboard tables for the VANTAGE-Bench leaderboard.

Primary API
-----------
build_all_html_tables(filtered_models, global_ranks) -> dict[str, str]
    Returns one hand-rendered HTML <table> string per pillar key in
    config.PILLARS ('overall', 'spatial', 'st', 'temporal', 'semantic'),
    all sharing the Overall tab's visual style (striping, hover, badges,
    bold column-max, row data-id for the click bridge in app.py).

make_radar_svg(model) -> str
    4-axis radar chart SVG (220 × 160 px) for the model detail side panel.
    Axes: Semantic · Spatial · Sp-Temp · Temporal (clockwise from top).
"""

from __future__ import annotations

import math

from .config import PILLARS, TASKS, TASK_METRIC_LABELS
from .data import ModelRecord

_MISSING = "—"

# Short task key (config.TASKS) → JSON score field in ModelRecord.scores.
_TASK_JSON_FIELD: dict[str, str] = {
    "loc":      "2d_localization",
    "ground":   "2d_referring_expressions",
    "pointing": "2d_spatial_pointing",
    "sot":      "single_object_tracking",
    "temploc":  "temporal_localization",
    "dvc":      "dense_video_captioning",
    "ev":       "event_verification",
    "vqa":      "video_qa",
}

# Overall tab: ordered (display_label, json_field) pairs.
# Task columns ordered to align with pillar super-header groups injected by JS
# (see _RESIZE_JS in app.py): Spatial (3) · Spatio-Temp (1) · Temporal (2) · Semantic (2).
_OVERALL_SCORE_COLS: list[tuple[str, str]] = [
    # Spatial
    ("Obj Loc",   "2d_localization"),
    ("Ref Exp",   "2d_referring_expressions"),
    ("Pointing",  "2d_spatial_pointing"),
    # Spatio-Temporal
    ("SOT",       "single_object_tracking"),
    # Temporal
    ("Temp Loc",  "temporal_localization"),
    ("DVC",       "dense_video_captioning"),
    # Semantic
    ("Event Ver", "event_verification"),
    ("VQA",       "video_qa"),
]

# Pillar group spec for the Overall-tab super-header row.
# Each entry: (colspan, label). Two empty leading cells cover the # and Name columns.
OVERALL_PILLAR_GROUPS: list[tuple[int, str]] = [
    (2, ""),              # # + Name
    (3, "Spatial"),
    (1, "Spatio-Temp"),
    (2, "Temporal"),
    (2, "Semantic"),
]

# Pillar aggregate score column shown immediately after Name on each pillar tab.
# Reuses the same pre-computed pillar field from the model record — no recomputation.
_PILLAR_AGGREGATE_COL: dict[str, tuple[str, str]] = {
    "spatial":  ("Spatial",  "spatial"),
    "st":       ("Sp-Temp",  "spatio_temporal"),
    "temporal": ("Temporal", "temporal"),
    "semantic": ("Semantic", "semantic"),
}


# -- Shared primitives -----------------------------------------------------


def _sort_by_rank(
    models: list[ModelRecord], rank_map: dict[str, int]
) -> list[ModelRecord]:
    sentinel = float("inf")
    return sorted(models, key=lambda m: (rank_map.get(m.id, sentinel), m.name))


def _column_maxes(
    models: list[ModelRecord], json_fields: list[str]
) -> dict[str, float]:
    """Maximum score per field across the given model set."""
    maxes: dict[str, float] = {}
    for f in json_fields:
        vals = [m.scores[f] for m in models if f in m.scores]
        if vals:
            maxes[f] = max(vals)
    return maxes


def _model_html(m: ModelRecord, rank_one_id: str | None) -> str:
    """HTML cell for the Model column: name + inline badges, org sub-line, type badges."""
    name_part = f"<b>{m.name}</b>" if m.id == rank_one_id else m.name
    if m.model_url:
        name_core = f'<a href="{m.model_url}" target="_blank" class="mc-link">{name_part}</a>'
    else:
        name_core = name_part
    verified_badge = '<span class="b b-verified">✓</span>' if m.verified else ""
    name_html = f'<span class="mc-name-row">{name_core}{verified_badge}</span>'

    badges: list[str] = []
    if m.result_type == "ensemble":
        badges.append('<span class="b b-ensemble">system / pipeline</span>')
    else:
        badges.append('<span class="b b-single">single</span>')
    if m.type == "open":
        badges.append('<span class="b b-open">open</span>')
    else:
        badges.append('<span class="b b-prop">prop.</span>')
    if m.is_new:
        badges.append('<span class="b b-new">new</span>')
    badge_html = "".join(badges)
    return (
        f'<span data-n="{m.name}" class="mc">'
        f'{name_html}'
        f'<span class="mc-org">{m.organization}</span>'
        f'<span class="mc-badges">{badge_html}</span>'
        f'</span>'
    )


def _score_cols_for_pillar(
    pillar: str,
    selected_tasks: list[str] | None = None,
) -> list[tuple[str, str]]:
    """Return [(display_label, json_field)] for one pillar.

    selected_tasks: JSON field names to include (Gradio column-toggle).
    None means show all. Only applies to non-overall pillars.
    """
    if pillar == "overall":
        return list(_OVERALL_SCORE_COLS)
    task_keys = PILLARS[pillar]
    if selected_tasks is not None:
        task_keys = [tk for tk in task_keys if _TASK_JSON_FIELD[tk] in selected_tasks]
    cols: list[tuple[str, str]] = []
    # Lead each pillar tab with its pre-computed pillar aggregate score.
    agg = _PILLAR_AGGREGATE_COL.get(pillar)
    if agg is not None:
        cols.append(agg)
    for tk in task_keys:
        json_field = _TASK_JSON_FIELD[tk]
        metric = TASK_METRIC_LABELS.get(json_field, "")
        label = f"{TASKS[tk]} ({metric})" if metric else TASKS[tk]
        cols.append((label, json_field))
    return cols


# -- Primary API -----------------------------------------------------------


def build_overall_html_table(
    filtered_models: list[ModelRecord],
    rank_map: dict[str, int],
) -> str:
    """Hand-rendered HTML <table> for the Overall tab.

    Used instead of gr.Dataframe because Gradio's DataFrame component
    does not cleanly support multi-level (grouped) column headers.
    Renders a two-row header:
        # · Name · Overall span both rows (rowspan=2)
        Spatial (×3) | Spatio-Temp (×1) | Temporal (×2) | Semantic (×2)
        Obj Loc | Ref Exp | Pointing | SOT | Temp Loc | DVC | Event Ver | VQA

    The Overall column sits between Name and the pillar groups as a
    standalone (non-grouped) column showing each model's stored overall
    score from ModelRecord.scores["overall"].

    Preserves: striping, hover, bold column-max, model badges, rank order,
    scroll behavior.
    """
    sorted_models = _sort_by_rank(filtered_models, rank_map)
    rank_one_id = next(
        (m.id for m in sorted_models if rank_map.get(m.id) == 1), None
    )

    score_cols = _OVERALL_SCORE_COLS
    json_fields = [f for _, f in score_cols] + ["overall"]
    col_max = _column_maxes(filtered_models, json_fields)
    overall_max = col_max.get("overall")

    # Render the table directly (no .table-wrap wrapper, no resize handle):
    # the gr.HTML element itself acts as the rectangular container, with
    # styling on .lb-table-overall in css.py. This avoids Gradio's
    # rounded/clipped DataFrame-shell visuals.
    parts: list[str] = ['<table class="lb-overall-table">']

    # Column-width control via <colgroup> (avoids nth-child collisions across
    # the two header rows). The Spatio-Temporal column gets a wider class
    # because its lone task (SOT) sits under the long "Spatio-Temp" pillar
    # super-header — we want that header to fit on one line.
    _ST_FIELD = "single_object_tracking"
    parts.append('<colgroup>')
    parts.append('<col class="col-rank">')
    parts.append('<col class="col-name">')
    parts.append('<col class="col-overall">')
    for _, field in score_cols:
        cls = "col-score col-score-st" if field == _ST_FIELD else "col-score"
        parts.append(f'<col class="{cls}">')
    parts.append('</colgroup>')

    # Header: two rows. # / Name / Overall span both via rowspan=2.
    parts.append('<thead>')
    parts.append('<tr class="lb-group-row">')
    parts.append('<th class="lb-corner" rowspan="2">#</th>')
    parts.append('<th class="lb-corner" rowspan="2">Name</th>')
    parts.append('<th class="lb-corner lb-corner-num" rowspan="2">Overall</th>')
    for span, label in [(3, "Spatial"), (1, "Spatio-Temp"),
                        (2, "Temporal"), (2, "Semantic")]:
        parts.append(f'<th class="lb-group" colspan="{span}">{label}</th>')
    parts.append('</tr>')

    parts.append('<tr class="lb-task-row">')
    for label, _ in score_cols:
        parts.append(f'<th>{label}</th>')
    parts.append('</tr>')
    parts.append('</thead>')

    # Body
    # Total column count: # + Name + Overall + 8 task scores = 11.
    total_cols = 3 + len(score_cols)
    parts.append('<tbody>')
    if not sorted_models:
        parts.append(
            f'<tr><td colspan="{total_cols}" class="lb-empty-row">'
            f'No models match — adjust the filters.</td></tr>'
        )
    for i, m in enumerate(sorted_models, 1):
        parts.append(f'<tr data-id="{m.id}">')
        parts.append(f'<td class="lb-rank">{i}</td>')
        parts.append(f'<td class="lb-name">{_model_html(m, rank_one_id)}</td>')

        # Overall cell — bold if it's the column max.
        ov = m.scores.get("overall")
        if ov is None:
            parts.append(f'<td class="lb-score lb-overall">{_MISSING}</td>')
        else:
            is_max = overall_max is not None and ov == overall_max
            cls = "lb-score lb-overall lb-max" if is_max else "lb-score lb-overall"
            parts.append(f'<td class="{cls}">{ov:.2f}</td>')

        for _, f in score_cols:
            v = m.scores.get(f)
            if v is None:
                parts.append(f'<td class="lb-score">{_MISSING}</td>')
            else:
                m_val = col_max.get(f)
                is_max = m_val is not None and v == m_val
                cls = "lb-score lb-max" if is_max else "lb-score"
                parts.append(f'<td class="{cls}">{v:.2f}</td>')
        parts.append('</tr>')
    parts.append('</tbody>')
    parts.append('</table>')
    return "".join(parts)


def build_pillar_html_table(
    pillar: str,
    filtered_models: list[ModelRecord],
    rank_map: dict[str, int],
) -> str:
    """Hand-rendered HTML <table> for a pillar tab (Spatial / Sp-Temp /
    Temporal / Semantic) — mirrors build_overall_html_table's markup and
    CSS classes so every tab shares one visual style.

    Single-row header (no pillar super-header grouping needed on these
    tabs). The pillar's own aggregate score is the first score column and
    gets the same purple-outlined "headline" treatment as the Overall
    column on the Overall tab (`lb-agg` mirrors `lb-overall`).
    """
    sorted_models = _sort_by_rank(filtered_models, rank_map)
    rank_one_id = next(
        (m.id for m in sorted_models if rank_map.get(m.id) == 1), None
    )

    score_cols = _score_cols_for_pillar(pillar)
    json_fields = [f for _, f in score_cols]
    col_max = _column_maxes(filtered_models, json_fields)
    agg_field = _PILLAR_AGGREGATE_COL[pillar][1] if pillar in _PILLAR_AGGREGATE_COL else None

    parts: list[str] = ['<table class="lb-pillar-table">']

    parts.append('<colgroup>')
    parts.append('<col class="col-rank">')
    parts.append('<col class="col-name">')
    for _, field in score_cols:
        cls = "col-score col-agg" if field == agg_field else "col-score"
        parts.append(f'<col class="{cls}">')
    parts.append('</colgroup>')

    parts.append('<thead>')
    parts.append('<tr class="lb-header-row">')
    parts.append('<th class="lb-corner">#</th>')
    parts.append('<th class="lb-corner">Name</th>')
    for label, field in score_cols:
        if field == agg_field:
            parts.append(f'<th class="lb-agg-th">{label}</th>')
        else:
            parts.append(f'<th>{label}</th>')
    parts.append('</tr>')
    parts.append('</thead>')

    total_cols = 2 + len(score_cols)
    parts.append('<tbody>')
    if not sorted_models:
        parts.append(
            f'<tr><td colspan="{total_cols}" class="lb-empty-row">'
            f'No models match — adjust the filters.</td></tr>'
        )
    for i, m in enumerate(sorted_models, 1):
        parts.append(f'<tr data-id="{m.id}">')
        parts.append(f'<td class="lb-rank">{i}</td>')
        parts.append(f'<td class="lb-name">{_model_html(m, rank_one_id)}</td>')

        for _, f in score_cols:
            is_agg = f == agg_field
            v = m.scores.get(f)
            base_cls = "lb-score lb-agg" if is_agg else "lb-score"
            if v is None:
                parts.append(f'<td class="{base_cls}">{_MISSING}</td>')
            else:
                m_val = col_max.get(f)
                is_max = m_val is not None and v == m_val
                cls = f"{base_cls} lb-max" if is_max else base_cls
                parts.append(f'<td class="{cls}">{v:.2f}</td>')
        parts.append('</tr>')
    parts.append('</tbody>')
    parts.append('</table>')
    return "".join(parts)


def build_all_html_tables(
    filtered_models: list[ModelRecord],
    global_ranks: dict[str, dict[str, int]],
) -> dict[str, str]:
    """Return {pillar_key: html_table_str} for every tab, Overall included.

    Replaces build_all_tables + build_overall_html_table as the single
    entry point once every tab renders as a hand-built HTML table instead
    of a mix of gr.HTML (Overall) and gr.Dataframe (the four pillars).
    """
    result: dict[str, str] = {
        "overall": build_overall_html_table(filtered_models, global_ranks["overall"]),
    }
    for pillar in PILLARS:
        if pillar == "overall":
            continue
        result[pillar] = build_pillar_html_table(
            pillar, filtered_models, global_ranks[pillar]
        )
    return result


# -- Radar SVG ------------------------------------------------------------


def make_radar_svg(m: ModelRecord) -> str:
    """4-axis radar chart SVG for the model detail side panel.

    Axes (clockwise from top): Semantic · Spatial · Sp-Temp · Temporal.
    Canvas: 220 × 160 px.  Scores assumed in [0, 100].
    """
    W, H = 220, 160
    cx, cy = W / 2, 82.0   # slightly below centre for top-label room
    chart_r = 50.0          # polygon radius
    label_r = 67.0          # label ring radius

    labels = ["Semantic", "Spatial", "Sp-Temp", "Temporal"]
    fields = ["semantic", "spatial", "spatio_temporal", "temporal"]
    vals: list[float] = []
    for f in fields:
        raw = m.scores.get(f)
        v = float(raw) / 100.0 if raw is not None else 0.0
        vals.append(max(0.0, min(1.0, v)))

    N = 4
    step = 2 * math.pi / N
    off = -math.pi / 2     # index 0 points straight up

    parts: list[str] = [
        f'<svg xmlns="http://www.w3.org/2000/svg" '
        f'width="{W}" height="{H}" viewBox="0 0 {W} {H}">'
    ]

    # Grid rings (4 concentric)
    for g in range(1, 5):
        rg = chart_r * g / 4
        parts.append(
            f'<circle cx="{cx}" cy="{cy}" r="{rg:.1f}" '
            f'fill="none" stroke="#e5e7eb" stroke-width="0.5"/>'
        )

    # Axis lines
    for i in range(N):
        a = off + i * step
        x2 = cx + chart_r * math.cos(a)
        y2 = cy + chart_r * math.sin(a)
        parts.append(
            f'<line x1="{cx}" y1="{cy}" '
            f'x2="{x2:.1f}" y2="{y2:.1f}" '
            f'stroke="#e5e7eb" stroke-width="0.5"/>'
        )

    # Data polygon
    poly_pts = " ".join(
        f"{cx + chart_r * v * math.cos(off + i * step):.1f},"
        f"{cy + chart_r * v * math.sin(off + i * step):.1f}"
        for i, v in enumerate(vals)
    )
    parts.append(
        f'<polygon points="{poly_pts}" '
        f'fill="rgba(37,99,235,0.12)" stroke="#2563eb" '
        f'stroke-width="1.5" stroke-linejoin="round"/>'
    )

    # Labels
    for i, lbl in enumerate(labels):
        a = off + i * step
        lx = cx + label_r * math.cos(a)
        ly = cy + label_r * math.sin(a)
        parts.append(
            f'<text x="{lx:.1f}" y="{ly:.1f}" '
            f'text-anchor="middle" dominant-baseline="middle" '
            f'font-size="9" fill="#6b7280" font-family="sans-serif">'
            f'{lbl}</text>'
        )

    parts.append("</svg>")
    return "".join(parts)