leaderboard.py

"""Table building for the PRIMO front-end: persisted results + registry -> frames.

Everything that turns stored scores into what the user sees lives here, kept
apart from Gradio and from any I/O so it unit-tests without a network
(``app.py`` fetches and renders, ``boards.py`` decides what a board *is*).

Two tables per board, deliberately different:

``ranked_table``   the leaderboard. Only models that covered EVERY task of the
                   board are listed -- a partial run cannot win by skipping the
                   hard cohorts.
``per_task_table`` one row per task, one column per model. Keeps partial
                   submissions, so a newcomer who covered three cohorts sees
                   their numbers instead of vanishing. Never a ranking.

Our own baselines are entries like any other: they are labelled, and they rank
where their score puts them. A baseline pinned to the bottom would hide the one
result worth publishing -- a foundation model losing to a PCA.

Task metadata is read defensively: a registry written before diseases were
recorded yields blank cells rather than breaking the page.
"""

import math
from dataclasses import dataclass

import pandas as pd
from boards import Board, label, metric_label
from evaluator import _norm_id
from results import IS_BASELINE, MODEL_NAME, display_name, with_baseline_flag
from scoring import TaskScore, category_means, sort_key

EMPTY_RANKED_COLUMNS = ["Model"]
EMPTY_PER_TASK_COLUMNS = ["Task"]
SCORE_DECIMALS = 3


@dataclass(frozen=True)
class TopModel:
    """One line of a home card's mini-ranking; ``score`` is None when unrankable."""

    name: str
    score: float | None
    is_baseline: bool


def _round(value: float | None) -> float | None:
    """Round for display; anything non-finite becomes a blank cell, not ``-inf``."""
    if value is None or not math.isfinite(value):
        return None
    return round(value, SCORE_DECIMALS)


def latest_only(df: pd.DataFrame) -> pd.DataFrame:
    """Keep each entry's most recent submission, so nobody can shop for a lucky run.

    An entry is ``(name, is_baseline)``, not just the name. Sharing one namespace
    would let somebody who submits a model called ``pca-50`` bury the published
    ``pca-50`` baseline simply by submitting after it.
    """
    flagged = with_baseline_flag(df)
    if flagged.empty:
        return flagged
    latest = flagged.groupby([MODEL_NAME, IS_BASELINE])["submitted_at"].transform("max")
    return flagged[flagged["submitted_at"] == latest]


def scores_from_rows(rows: pd.DataFrame, by_id: dict[str, dict]) -> list[TaskScore]:
    """Rebuild one submission's task scores from persisted rows + the registry.

    Metadata (dataset_id, category, metric) is joined from the registry by
    ``task_id``; a row whose task has retired from the registry is dropped.
    """
    out = []
    for _, r in rows.iterrows():
        task_id = _norm_id(r["task_id"])
        meta = by_id.get(task_id)
        if meta is None:
            continue
        out.append(
            TaskScore(
                task_id=task_id,
                dataset_id=_norm_id(meta["dataset_id"]),
                category=str(meta["category"]),
                metric=str(meta["metric"]),
                score=float(r["score"]),
                n_samples=0,
            )
        )
    return out


def _board_registry(by_id: dict[str, dict], board: Board) -> dict[str, dict]:
    return {tid: meta for tid, meta in by_id.items() if tid in board.task_ids}


def _entries(df: pd.DataFrame, by_id: dict[str, dict], board: Board) -> list[dict]:
    """Per-category means for every model that covered the whole board."""
    scoped = _board_registry(by_id, board)
    entries = []
    for (model, is_baseline, submitted), rows in latest_only(df).groupby(
        [MODEL_NAME, IS_BASELINE, "submitted_at"]
    ):
        if not board.task_ids.issubset({_norm_id(t) for t in rows["task_id"]}):
            continue
        categories = category_means(scores_from_rows(rows, scoped))
        entries.append(
            {
                "model_name": model,
                "submitted_at": submitted,
                "categories": categories,
                "rank": sort_key(categories),
                "is_baseline": bool(is_baseline),
            }
        )
    return sorted(entries, key=lambda e: -e["rank"])


def ranked_table(
    df: pd.DataFrame, by_id: dict[str, dict], board: Board
) -> pd.DataFrame:
    """Rank the models that covered every task of ``board``.

    One column per task category, each in its native metric (AUROC or Pearson);
    two metrics never share a column. ``Mean`` averages those columns -- it does
    cross metrics, which is why the pages call it a tie-break rather than a
    score -- and is dropped when the board holds a single category, where it
    would just repeat it.
    """
    entries = _entries(df, by_id, board)
    if not entries:
        return pd.DataFrame(columns=EMPTY_RANKED_COLUMNS)

    columns = sorted({cat for e in entries for cat in e["categories"]})
    metric_of: dict[str, str] = {}
    for entry in entries:
        for cat, stats in entry["categories"].items():
            metric_of.setdefault(cat, stats["metric"])

    rows = []
    for position, entry in enumerate(entries, start=1):
        categories = entry["categories"]
        model = display_name(entry["model_name"], entry["is_baseline"])
        row = {"Rank": position, "Model": model}
        if len(columns) > 1:
            row["Mean"] = _round(entry["rank"])
        for cat in columns:
            header = f"{label(cat)} ({metric_label(metric_of[cat])})"
            row[header] = _round(categories[cat]["mean"]) if cat in categories else None
        rows.append(row)
    return pd.DataFrame(rows)


RESERVED_COLUMNS = frozenset({"Task", "Family", "Area", "Metric", "Best"})


def per_task_table(
    df: pd.DataFrame, by_id: dict[str, dict], board: Board
) -> pd.DataFrame:
    """One row per task, one column per model -- read across to pick a model.

    Unlike the ranked table this keeps partial submissions: a model that covered
    only some tasks appears with gaps rather than vanishing. Models are ordered
    by their mean over the board, so the strongest column comes first.

    Model names become COLUMN names here, which is why ``RESERVED_COLUMNS``
    exists: a model called ``Task`` would overwrite every task title with its own
    score. The Submit form refuses those names rather than this table renaming
    them, so what a submitter types is what the leaderboard shows.

    ``Best`` breaks a tie towards the model with the better mean over the board,
    because ``present`` is built in ``models`` order -- which is that ranking.
    """
    scoped = _board_registry(by_id, board)
    if df.empty or not scoped:
        return pd.DataFrame(columns=EMPTY_PER_TASK_COLUMNS)

    scores: dict[tuple[str, str], float] = {}
    for _, r in latest_only(df).iterrows():
        task_id = _norm_id(r["task_id"])
        if task_id in scoped:
            model = display_name(str(r[MODEL_NAME]), bool(r[IS_BASELINE]))
            scores[(task_id, model)] = float(r["score"])
    if not scores:
        return pd.DataFrame(columns=EMPTY_PER_TASK_COLUMNS)

    by_model: dict[str, list[float]] = {}
    for (_, model), score in scores.items():
        by_model.setdefault(model, []).append(score)
    models = sorted(by_model, key=lambda m: -sum(by_model[m]) / len(by_model[m]))

    rows = []
    for task_id, task in sorted(scoped.items()):
        row = {
            "Task": str(task.get("title") or task_id),
            "Family": label(str(task.get("category", ""))),
            "Area": str(task.get("therapeutic_area", "")),
            "Metric": metric_label(str(task.get("metric", ""))),
        }
        present = {m: scores[(task_id, m)] for m in models if (task_id, m) in scores}
        for model in models:
            row[model] = _round(present.get(model))
        row["Best"] = max(present, key=present.get) if present else None
        rows.append(row)
    return pd.DataFrame(rows)


def top_models(
    df: pd.DataFrame, by_id: dict[str, dict], board: Board, limit: int
) -> list[TopModel]:
    """The board's best full-coverage entries -- the card teaser, baselines included."""
    return [
        TopModel(
            name=str(e["model_name"]),
            score=_round(e["rank"]),
            is_baseline=e["is_baseline"],
        )
        for e in _entries(df, by_id, board)[:limit]
    ]


TASK_COLUMNS = [
    "Task",
    "Predicts",
    "Family",
    "Area",
    "Disease",
    "Tissue",
    "Patients",
    "Metric",
]


def tasks_table(tasks: list[dict]) -> pd.DataFrame:
    """The Tasks page: one row per task, scannable. Provenance is never a column."""
    if not tasks:
        return pd.DataFrame(columns=TASK_COLUMNS)
    rows = []
    for task in sorted(tasks, key=lambda t: str(t.get("task_id", ""))):
        rows.append(
            {
                "Task": str(task.get("title") or task.get("task_id", "")),
                "Predicts": str(task.get("description") or ""),
                "Family": label(str(task.get("category", ""))),
                "Area": str(task.get("therapeutic_area", "")),
                "Disease": ", ".join(str(d) for d in task.get("diseases") or []),
                "Tissue": str(task.get("tissue", "")),
                "Patients": task.get("n_samples"),
                "Metric": metric_label(str(task.get("metric", ""))),
            }
        )
    return pd.DataFrame(rows, columns=TASK_COLUMNS)


REPOSITORY_OF_PREFIX = {"GSE": "NCBI GEO", "E-MTAB": "EMBL-EBI ArrayExpress"}


def source_repositories(tasks: list[dict]) -> list[str]:
    """The public archives the cohorts come from.

    Only the archive is named, never the accession: publishing ``GSE193677`` would
    let anyone download that study and recover the labels by matching expression.

    The registry's ``license`` field is deliberately NOT surfaced. It is one
    hand-written line per dataset, unvalidated, and several cohorts pool three
    studies under it -- so publishing it would state a rights position nobody has
    checked. About says what the archives actually say instead.
    """
    repositories: set[str] = set()
    for task in tasks:
        for source in task.get("sources") or []:
            for prefix, repository in REPOSITORY_OF_PREFIX.items():
                if str(source).upper().startswith(prefix):
                    repositories.add(repository)
    return sorted(repositories)