app.py

"""Gradio front-end for the PRIMO public benchmark.

Six pages: Home (a grid of boards), Leaderboard (one board at a time), Tasks
("what is actually being tested?"), Submit ("how do I enter?"), Contribute
("what is missing, and how do I add it?"), About ("can I trust this?"). A tab is
addressable as ``?tab=contribute``, which is what the open cards on Home link to.

Upload one embedding file spanning every dataset (rows keyed by ``dataset_id``
+ ``sample_id``); a fixed linear probe scores each task (a dataset may carry
several hidden targets). Results roll up into BOARDS -- the whole modality, one
therapeutic area, one task family -- and each board ranks the models that
covered all of its tasks, one column per category in its native metric (AUROC or
Pearson), plus a ``Mean`` of those columns that orders the rows and is labelled
as the cross-metric average it is.

Disclosure policy -- what the public pages may show:
  per task   disease, tissue, area, what is predicted, class names, n, metric
  aggregated the public archives the cohorts sit in, and their licences
  never      study accessions, dataset_id -> cohort, hub keys, per-sample labels
Naming the accession behind ``d002`` would put every label one GEO download
away, so ``sources`` is collapsed to its archive (NCBI GEO / EMBL-EBI
ArrayExpress) and ``citation`` is never rendered at all.

Failures are surfaced by who owns them: a bad file tells the submitter exactly
what to fix; an evaluator-side failure says "our side, please retry" and logs
the traceback rather than blaming the submission.

A page load is ONE fetch: ``_init`` pulls the registry, the boards and the
persisted results once and threads that state into every tab, rather than each
tab fetching for itself.

The board dropdown takes ``allow_custom_value``: its choices are only filled once
the registry loads, so a browser holding a value from an earlier version of the
page -- or any load where the registry is briefly unreachable -- would otherwise
be refused by Gradio's own preprocessing, before ``by_slug`` gets the chance to
fall back to the hero board.
"""

import os
import traceback
from datetime import datetime, timezone
from html import escape
from pathlib import Path

import gradio as gr
import pandas as pd
from boards import Board, build_boards, by_slug
from evaluator import (
    EvaluatorError,
    SubmissionError,
    _norm_id,
    fetch_manifest,
    fetch_tasks_registry,
    manifest_ids,
    score_all,
    scoreable_tasks,
)
from home import banner_html, home_html
from leaderboard import (
    RESERVED_COLUMNS,
    per_task_table,
    ranked_table,
    source_repositories,
    tasks_table,
)
from results import (
    BASELINE_TAG,
    IS_BASELINE,
    OWNER,
    RESULT_COLUMNS,
    append_results,
    append_submission,
    owner_of,
    read_results,
)

PAGES_DIR = Path(__file__).parent / "pages"
STYLE = (Path(__file__).parent / "style.css").read_text()

TOKEN = os.environ.get("HF_TOKEN")
TAB_IDS = ("home", "leaderboard", "tasks", "submit", "contribute", "about")

DEEP_NAVY = "#050A3C"
PANEL = "#0A1049"
LINE = "#1E2775"
CYAN = "#16B3C0"
PAPER = "#F3F8F8"
MUTED = "#9FB3B7"

THEME = gr.themes.Base(
    primary_hue=gr.themes.colors.cyan,
    secondary_hue=gr.themes.colors.blue,
    neutral_hue=gr.themes.colors.slate,
    font=(
        gr.themes.GoogleFont("Funnel Sans"),
        "ui-sans-serif",
        "system-ui",
        "sans-serif",
    ),
).set(
    color_accent=CYAN,
    border_color_accent=CYAN,
    body_background_fill=DEEP_NAVY,
    body_background_fill_dark=DEEP_NAVY,
    background_fill_primary=DEEP_NAVY,
    background_fill_primary_dark=DEEP_NAVY,
    background_fill_secondary=PANEL,
    background_fill_secondary_dark=PANEL,
    block_background_fill=DEEP_NAVY,
    block_background_fill_dark=DEEP_NAVY,
    panel_background_fill=PANEL,
    block_label_background_fill=PANEL,
    block_label_text_color=MUTED,
    body_text_color=PAPER,
    body_text_color_dark=PAPER,
    body_text_color_subdued=MUTED,
    body_text_color_subdued_dark=MUTED,
    border_color_primary=LINE,
    border_color_primary_dark=LINE,
    block_border_color=LINE,
    input_background_fill=PANEL,
    input_background_fill_dark=PANEL,
    button_primary_background_fill=CYAN,
    button_primary_text_color=DEEP_NAVY,
    button_secondary_background_fill=PANEL,
    table_border_color=LINE,
    table_text_color=PAPER,
    table_even_background_fill=DEEP_NAVY,
    table_odd_background_fill=PANEL,
    link_text_color=CYAN,
)


def _page_text(name: str) -> str:
    """One ``pages/<tab>.md`` per prose tab, named after the tab it fills.

    Editing the site's words never means touching Python. Home has no file --
    it is the generated board grid.
    """
    return (PAGES_DIR / f"{name}.md").read_text()


def _registry_by_id() -> dict[str, dict]:
    """Scoreable tasks keyed by task_id (dataset present in the public manifest)."""
    datasets = manifest_ids(fetch_manifest(TOKEN))
    registry = scoreable_tasks(fetch_tasks_registry(TOKEN), datasets)
    return {_norm_id(task["task_id"]): task for task in registry}


PageState = tuple[dict[str, dict], list[Board], pd.DataFrame]


def _page_state() -> PageState:
    """Registry, boards and persisted results -- one fetch per render.

    A failed fetch yields empty structures so the page renders a "come back
    later" state instead of a stack trace. Every tab is built from one of these,
    threaded through rather than refetched, so a page load is one round trip.
    """
    try:
        by_id = _registry_by_id()
        return by_id, build_boards(by_id), read_results(TOKEN)
    except Exception:  # noqa: BLE001
        traceback.print_exc()
        return {}, [], pd.DataFrame(columns=RESULT_COLUMNS)


COLUMN_WIDTHS = {
    "Rank": "70px",
    "Model": "230px",
    "Mean": "110px",
    "Task": "300px",
    "Family": "150px",
    "Area": "170px",
    "Metric": "100px",
    "Best": "170px",
}
DEFAULT_WIDTH = "160px"


def _score_columns(df: pd.DataFrame) -> list[str]:
    return [c for c in df.select_dtypes("number").columns if c != "Rank"]


def _styled(df: pd.DataFrame, label: str, axis: int, pinned: int) -> gr.DataFrame:
    """An MTEB-style table: best value in bold, ids pinned, searchable.

    ``axis=0`` bolds the best model per column (the ranked table); ``axis=1``
    bolds the best model per row (the per-task table, read across). The label
    names the table, so the copy file needs one block per tab, not per table.
    """
    if df.empty:
        return gr.DataFrame(df, label=label, interactive=False)
    scores = _score_columns(df)
    labels = [c for c in df.columns if c not in scores]
    styler = df.style.format("{:.3f}", subset=scores, na_rep="n/a").highlight_max(
        subset=scores, axis=axis, props="font-weight: 700"
    )
    if labels:
        styler = styler.format(na_rep="n/a", subset=labels)
    return gr.DataFrame(
        styler,
        label=label,
        interactive=False,
        wrap=True,
        pinned_columns=pinned,
        column_widths=[COLUMN_WIDTHS.get(c, DEFAULT_WIDTH) for c in df.columns],
        show_search="filter",
        show_copy_button=True,
        show_fullscreen_button=True,
    )


def _board_header(board: Board | None) -> str:
    """The board title strip. Registry strings are escaped: they are data, not code."""
    if board is None:
        return (
            '<div class="primo-board-head"><h2>No board available</h2>'
            "<p>The task registry could not be loaded. Please retry shortly.</p></div>"
        )
    return (
        '<div class="primo-board-head primo-accent"><h2>'
        f'<span class="primo-code primo-code-modality">{escape(board.code)}</span>'
        f"{escape(board.name)}</h2>"
        f"<p>{escape(board.blurb)} {board.n_tasks} tasks · "
        f"{board.n_cohorts} cohorts · {board.n_patients:,} patients · "
        f"{escape(board.modality)}</p></div>"
    )


RANKED_LABEL = "🏆 Ranked"
PER_TASK_LABEL = "🔍 Per task"


def _board_page(slug: str | None, state: PageState | None = None):
    """Header + both tables for one board.

    ``state`` lets a caller that has already fetched pass it in; the first render
    builds every tab from one fetch that way.
    """
    by_id, boards, df = _page_state() if state is None else state
    board = by_slug(boards, slug)
    if board is None:
        empty = pd.DataFrame()
        return (
            _board_header(None),
            _styled(empty, RANKED_LABEL, 0, 0),
            _styled(empty, PER_TASK_LABEL, 1, 0),
        )
    return (
        _board_header(board),
        _styled(ranked_table(df, by_id, board), RANKED_LABEL, 0, 2),
        _styled(per_task_table(df, by_id, board), PER_TASK_LABEL, 1, 1),
    )


def _board_choices(boards: list[Board]) -> list[tuple[str, str]]:
    return [(f"{b.code} · {b.name} · {b.group}", b.slug) for b in boards]


def _about_text(by_id: dict[str, dict]) -> str:
    """Methodology + the archives the cohorts live in, never their accessions.

    The placeholder is substituted, not ``.format``-ed: a page of prose is free
    to contain a brace, and a stray one must not blow up the tab. An unreachable
    registry leaves ``by_id`` empty and the sentence falls back to a generic one.
    """
    repositories = source_repositories(list(by_id.values()))
    return _page_text("about").replace(
        "{repositories}", " and ".join(repositories) or "public archives"
    )


def _summary(result: dict, model_name: str) -> str:
    lines = [
        f"**{model_name}**: covered {result['n_datasets_scored']}/"
        f"{result['n_datasets_total']} datasets",
        "",
    ]
    for category, stats in sorted(result["categories"].items()):
        lines.append(f"- **{category}** ({stats['metric']}) = {stats['mean']:.3f}")
    if result["full_coverage"]:
        lines.append("\n✅ **full coverage.** You are ranked on every board.")
    else:
        lines.append(
            "\n⚠️ **partial coverage.** You are ranked on the boards whose tasks you "
            "covered in full, and your scores always appear in each board's "
            "**per-task** table."
        )
        if result["missing"]:
            lines.append(f"- missing from file: {result['missing']}")
        if result["incomplete"]:
            lines.append(f"- could not score: {result['incomplete']}")
    return "\n".join(lines)


def _claimed_by(model: str) -> str:
    """Who already owns this model name, or ``""``.

    A results fetch that fails leaves the name free: a Hugging Face hiccup must
    not block a submission, and the worst case is the collision we had before.
    """
    try:
        return owner_of(read_results(TOKEN), model)
    except Exception:  # noqa: BLE001
        traceback.print_exc()
        return ""


def evaluate(
    submission_path: str,
    model_name: str,
    email: str,
    paper_link: str,
    hf_model_link: str,
    notes: str,
    slug: str | None,
    profile: gr.OAuthProfile | None,
):
    def _refuse(message: str):
        return (message, *_board_page(slug))

    if profile is None:
        return _refuse("Please sign in with Hugging Face to submit.")
    if not submission_path:
        return _refuse("Please upload a submission file.")
    if not model_name or not model_name.strip():
        return _refuse("Please enter a model name.")
    if not email or not email.strip():
        return _refuse("Please enter a contact email.")
    model = model_name.strip()
    if BASELINE_TAG in model.lower():
        return _refuse(
            f"`{BASELINE_TAG}` is reserved for our reference submissions. Please "
            "pick another model name."
        )
    if model in RESERVED_COLUMNS:
        return _refuse(
            f"`{model}` is a column of the per-task table. Please pick another "
            "model name."
        )
    claimed = _claimed_by(model)
    if claimed and claimed != profile.username:
        return _refuse(
            f"The model name `{model}` is already taken by @{claimed}, and the "
            "board keeps each name's latest submission. Please pick another name."
        )
    try:
        result = score_all(submission_path, TOKEN)
    except SubmissionError as error:
        return _refuse(f"❌ {error}")
    except EvaluatorError as error:
        traceback.print_exc()
        return _refuse(
            "⚠️ We couldn't evaluate your submission. This is on our side, not your "
            f"file. Please try again in a moment.\n\n`{error}`"
        )
    except Exception as error:  # noqa: BLE001
        traceback.print_exc()
        return _refuse(f"⚠️ Unexpected evaluation error (our side): {error}")

    summary = _summary(result, model)
    submitted_at = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
    rows = [
        {
            "model_name": model,
            "task_id": task.task_id,
            "score": round(float(task.score), 4),
            "submitted_at": submitted_at,
            IS_BASELINE: False,
            OWNER: profile.username,
        }
        for task in result["per_task"]
    ]
    meta = {
        "model_name": model,
        "submitted_at": submitted_at,
        OWNER: profile.username,
        "email": email.strip(),
        "paper_link": (paper_link or "").strip(),
        "hf_model_link": (hf_model_link or "").strip(),
        "notes": (notes or "").strip(),
    }
    try:
        append_results(rows, TOKEN)
        append_submission(meta, TOKEN)
    except Exception as error:  # noqa: BLE001
        traceback.print_exc()
        summary += f"\n\n⚠️ scored, but the leaderboard was not saved: {error}"
    return (summary, *_board_page(slug))


def _landing_tab(params: dict, has_board: bool) -> str:
    """Which tab a visitor lands on: ``?tab=`` wins, then ``?board=``, else Home.

    An unknown ``?tab=`` falls through to Home rather than selecting nothing,
    which would render the Space with every panel collapsed.
    """
    tab = params.get("tab")
    if tab in TAB_IDS:
        return tab
    return "leaderboard" if has_board else "home"


def _init(request: gr.Request):
    """Render every tab from one fetch, landing where the query params ask."""
    state = _page_state()
    by_id, boards, df = state
    params = dict(request.query_params) if request else {}
    board = by_slug(boards, params.get("board"))
    selected = _landing_tab(params, bool(params.get("board") and board))
    return (
        gr.Tabs(selected=selected),
        gr.Dropdown(
            choices=_board_choices(boards), value=board.slug if board else None
        ),
        banner_html(boards),
        home_html(boards, df, by_id),
        tasks_table(list(by_id.values())),
        _about_text(by_id),
        *_board_page(board.slug if board else None, state),
    )


def build_demo() -> gr.Blocks:
    with gr.Blocks(
        title="PRIMO Benchmark", theme=THEME, css=STYLE, fill_width=True
    ) as demo:
        gr.Markdown(
            "# 🧬 PRIMO: Patient Representations in Multi-Omics\n\n"
            "**Does your model's patient embedding predict a real clinical "
            "outcome?** We score it on cohorts whose labels you never see."
        )
        banner = gr.HTML()
        with gr.Tabs() as tabs:
            with gr.Tab("Home", id="home"):
                home = gr.HTML()
            with gr.Tab("Leaderboard", id="leaderboard"):
                board_sel = gr.Dropdown(
                    label="Board",
                    choices=[],
                    allow_custom_value=True,
                    info="A slice of the benchmark: one modality, area or task family.",
                )
                board_head = gr.HTML()
                gr.Markdown(_page_text("leaderboard"))
                ranked = gr.DataFrame(label=RANKED_LABEL, interactive=False)
                per_task = gr.DataFrame(label=PER_TASK_LABEL, interactive=False)
            with gr.Tab("Tasks", id="tasks"):
                gr.Markdown(_page_text("tasks"))
                tasks_df = gr.DataFrame(interactive=False, wrap=True)
            with gr.Tab("Submit", id="submit"):
                gr.Markdown(_page_text("submit"))
                gr.LoginButton()
                with gr.Row():
                    with gr.Column():
                        model_tb = gr.Textbox(
                            label="Model name",
                            placeholder="e.g. eva-rna-v1",
                            info="Shown on the leaderboard.",
                        )
                        email_tb = gr.Textbox(
                            label="Email address",
                            placeholder="you@lab.org",
                            info="Contact for this submission, kept private.",
                        )
                        notes_tb = gr.Textbox(
                            label="Training data / notes (optional)",
                            placeholder="e.g. pretrained on atlas X",
                            info="About the model or its training data.",
                        )
                    with gr.Column():
                        paper_tb = gr.Textbox(
                            label="Paper link (optional)",
                            placeholder="https://arxiv.org/abs/...",
                        )
                        hf_tb = gr.Textbox(
                            label="Hugging Face model link (optional)",
                            placeholder="https://huggingface.co/...",
                        )
                        file_in = gr.File(
                            label="Submission (.csv / .tsv / .parquet / .npz)",
                            type="filepath",
                        )
                run_btn = gr.Button("Evaluate", variant="primary")
                result_md = gr.Markdown()
            with gr.Tab("Contribute", id="contribute"):
                gr.Markdown(_page_text("contribute"))
            with gr.Tab("About", id="about"):
                about_md = gr.Markdown()

        board_view = [board_head, ranked, per_task]
        run_btn.click(
            evaluate,
            [file_in, model_tb, email_tb, paper_tb, hf_tb, notes_tb, board_sel],
            [result_md, *board_view],
        )
        board_sel.change(_board_page, board_sel, board_view)
        demo.load(
            _init,
            None,
            [tabs, board_sel, banner, home, tasks_df, about_md, *board_view],
        )
    return demo


if __name__ == "__main__":
    build_demo().launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)