"""Standalone probe for the PRIMO public benchmark.

Loads one embedding submission that spans every dataset (rows keyed by
``dataset_id`` + ``sample_id``), then scores each TASK = (dataset, target): a
dataset is embedded once and reused across all its tasks. Per task it fits a
fixed probe—linear for scalar targets and multi-output ridge for response
matrices—over repeated frozen CV folds or a fixed transfer split. Turning those
predictions into scores + leaderboard aggregates is the job of ``scoring.py``;
this module only reads data and runs the probe.

Failures are split by who caused them: a bad submission raises/records a
``SubmissionError`` (the submitter fixes it); anything on our side, such as a
failed Hugging Face fetch or an unexpected bug, raises ``EvaluatorError`` so it
is never silently charged against the submitter.

Reads the public ``datasets.yaml`` manifest (what the submitter embeds) plus the
private ``tasks.yaml`` registry, each task's private ``labels.csv``, and hidden
perturbation ``targets.npz`` matrices where applicable. No monorepo imports, so
it runs unchanged inside a public Hugging Face Space.

Deps: numpy, pandas, scikit-learn, pyyaml. ``huggingface_hub`` is used only by
the fetch helpers / CLI, imported lazily.
"""

import argparse
import logging
import os
import time
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd
import yaml
from scoring import (
    METRICS,
    TaskScore,
    category_means,
    compute_residual_sample_spearman,
    compute_target_centered_sample_spearman,
)
from sklearn.linear_model import LogisticRegressionCV, RidgeCV
from sklearn.model_selection import StratifiedKFold
from sklearn.preprocessing import StandardScaler

SAMPLE_ID = "sample_id"
DATASET_ID = "dataset_id"
TASK_ID = "task_id"
LABEL = "label"
FOLD = "fold"
FOLD_PREFIX = f"{FOLD}_"
SPLIT = "split"
SPLIT_TRAIN = "train"
SPLIT_TEST = "test"
PAIR_ID = "pair_id"
CONTROL_ID = "control_id"
PERTURBED_ID = "perturbed_id"
PAIRING = "pairing"
PAIRED = "paired"
CARTESIAN = "cartesian"
PERTURBATION = "perturbation"
RESIDUAL_SPEARMAN = "residual_spearman"
CENTERED_SPEARMAN = "centered_spearman"
ORG = "ScientaLab"
PUBLIC_REPO = f"{ORG}/primo"
LABELS_REPO = f"{ORG}/primo-labels"
# Keep the established leaderboard history available until it is migrated.
RESULTS_REPO = "ScientaLab/primo-results"
MANIFEST_FILENAME = "datasets.yaml"
TASKS_FILENAME = "tasks.yaml"
LABELS_FILENAME = "labels.csv"
TARGETS_FILENAME = "targets.npz"

RIDGE_ALPHAS = np.logspace(-3.0, 6.0, 19)
PERTURBATION_TEST_Z_CLIP = 20.0
LOGREG_CS = 10
INNER_CV = 5
MAX_ITER = 5000
RANDOM_STATE = 0

FETCH_ATTEMPTS = 3
FETCH_BACKOFF = 1.0

NPZ_ID_KEYS = ("sample_ids", "sample_id", "ids")
NPZ_EMB_KEYS = ("embeddings", "embedding", "emb", "X")
NPZ_DSID_KEYS = ("dataset_ids", "dataset_id")


class SubmissionError(ValueError):
    """A submission the evaluator cannot score (bad format, missing samples...)."""


class EvaluatorError(RuntimeError):
    """A failure on our side (data fetch, unexpected bug), not the submitter's.

    Raised instead of being recorded as a per-dataset skip, so a transient
    Hugging Face hiccup or an evaluator bug never silently costs a submitter
    their coverage.
    """


@dataclass
class TaskOutcome:
    """What happened to one scoreable task: scored / missing / invalid."""

    task_id: str
    dataset_id: str
    status: str
    score: TaskScore | None = None
    reason: str | None = None


@dataclass(frozen=True)
class PerturbationTargets:
    """Hidden response vectors and their aligned gene identifiers."""

    delta: np.ndarray
    gene_ids: np.ndarray


def _norm_id(value: object) -> str:
    """Normalise an id so int/str/float spellings of the same id join."""
    if isinstance(value, float) and value.is_integer():
        return str(int(value))
    return str(value).strip()


def load_tasks_registry(path: str | Path) -> list[dict]:
    """Read the private ``tasks.yaml`` registry (one entry per task)."""
    with open(path) as handle:
        data = yaml.safe_load(handle)
    tasks = data.get("tasks", []) if isinstance(data, dict) else (data or [])
    ids = [_norm_id(t[TASK_ID]) for t in tasks]
    if len(set(ids)) != len(ids):
        raise EvaluatorError(f"duplicate task_id in the registry: {ids}")
    return tasks


def _fold_columns(labels: pd.DataFrame) -> list[str]:
    """Return the ordered external-CV fold columns in private labels."""
    repeated = [column for column in labels if column.startswith(FOLD_PREFIX)]
    invalid = [
        column for column in repeated if not column.removeprefix(FOLD_PREFIX).isdigit()
    ]
    if invalid:
        raise EvaluatorError(f"labels.csv has invalid repeat-fold columns: {invalid}")
    first = [FOLD] if FOLD in labels.columns else []
    return [
        *first,
        *sorted(repeated, key=lambda column: int(column.removeprefix(FOLD_PREFIX))),
    ]


def load_labels(path: str | Path) -> pd.DataFrame:
    """Read private labels with repeated CV folds or one transfer split."""
    df = pd.read_csv(path)
    is_perturbation = PAIR_ID in df.columns
    required = {SAMPLE_ID, PAIR_ID} if is_perturbation else {SAMPLE_ID, LABEL}
    missing = required - set(df.columns)
    if missing:
        raise EvaluatorError(f"labels.csv missing columns: {sorted(missing)}")
    has_folds = bool(_fold_columns(df))
    has_split = SPLIT in df.columns
    pairing_values = (
        set(df[PAIRING].astype(str)) if is_perturbation and PAIRING in df else set()
    )
    is_cartesian = pairing_values == {CARTESIAN}
    if sum((has_folds, has_split, is_cartesian)) != 1:
        raise EvaluatorError(
            "labels.csv needs exactly one evaluation assignment: folds, transfer "
            "split, or Cartesian pairing."
        )
    if is_cartesian:
        missing = {CONTROL_ID, PERTURBED_ID, PAIRING} - set(df.columns)
        if missing:
            raise EvaluatorError(
                f"Cartesian labels.csv missing columns: {sorted(missing)}"
            )
    return df


def load_targets(path: str | Path) -> PerturbationTargets:
    """Read and validate one task's hidden perturbation responses."""
    with np.load(path, allow_pickle=False) as data:
        missing = {"delta", "gene_ids"} - set(data.files)
        if missing:
            raise EvaluatorError(f"targets.npz missing arrays: {sorted(missing)}")
        delta = np.asarray(data["delta"], dtype=float)
        gene_ids = np.asarray(data["gene_ids"]).astype(str)
    if delta.ndim != 2:
        raise EvaluatorError(f"targets delta must be 2D, got shape {delta.shape}")
    if gene_ids.ndim != 1 or len(gene_ids) != delta.shape[1]:
        raise EvaluatorError("targets gene_ids must align with delta columns")
    if not np.isfinite(delta).all():
        raise EvaluatorError("targets delta contains NaN or inf values")
    return PerturbationTargets(delta=delta, gene_ids=gene_ids)


def load_submission(path: str | Path) -> dict[str, pd.DataFrame]:
    """Load a multi-dataset submission into ``{dataset_id: raw_block_frame}``."""
    ext = Path(path).suffix.lower()
    if ext == ".npz":
        frame = _npz_to_frame(path)
    elif ext in (".parquet", ".pq"):
        frame = pd.read_parquet(path)
    elif ext in (".tsv", ".txt"):
        frame = pd.read_csv(path, sep="\t")
    elif ext == ".csv":
        frame = pd.read_csv(path)
    else:
        raise SubmissionError(
            f"unsupported submission type '{ext}'. Use .csv/.tsv/.parquet/.npz"
        )
    return _split_by_dataset(frame)


def _npz_to_frame(path: str | Path) -> pd.DataFrame:
    data = np.load(path, allow_pickle=True)
    ds_key = next((k for k in NPZ_DSID_KEYS if k in data), None)
    id_key = next((k for k in NPZ_ID_KEYS if k in data), None)
    emb_key = next((k for k in NPZ_EMB_KEYS if k in data), None)
    if ds_key is None or id_key is None or emb_key is None:
        raise SubmissionError(
            f".npz needs a dataset-id array {NPZ_DSID_KEYS}, a sample-id array "
            f"{NPZ_ID_KEYS}, and an embedding array {NPZ_EMB_KEYS}; "
            f"found {list(data.keys())}"
        )
    emb = np.asarray(data[emb_key])
    if emb.ndim != 2:
        raise SubmissionError(f"embeddings must be 2D, got shape {emb.shape}")
    frame = pd.DataFrame(emb)
    frame.insert(0, SAMPLE_ID, np.asarray(data[id_key]))
    frame.insert(0, DATASET_ID, np.asarray(data[ds_key]))
    return frame


def _split_by_dataset(frame: pd.DataFrame) -> dict[str, pd.DataFrame]:
    """Group a submission by ``dataset_id`` into per-dataset blocks (raw frames)."""
    lower = {str(c).lower(): c for c in frame.columns}
    if DATASET_ID not in lower:
        raise SubmissionError(
            f"submission needs a '{DATASET_ID}' column; found {list(frame.columns)}"
        )
    ds_col = lower[DATASET_ID]
    blocks: dict[str, pd.DataFrame] = {}
    for raw_ds, group in frame.groupby(ds_col, sort=False):
        dataset_id = _norm_id(raw_ds)
        if dataset_id in blocks:
            raise SubmissionError(
                f"dataset_id '{dataset_id}' appears under multiple spellings"
            )
        blocks[dataset_id] = group.drop(columns=[ds_col]).dropna(axis=1, how="all")
    if not blocks:
        raise SubmissionError("empty submission")
    return blocks


def _to_embedding_frame(frame: pd.DataFrame) -> pd.DataFrame:
    lower = {str(c).lower(): c for c in frame.columns}
    if SAMPLE_ID not in lower:
        raise SubmissionError(
            f"submission needs a '{SAMPLE_ID}' column; found {list(frame.columns)}"
        )
    id_col = lower[SAMPLE_ID]
    ids = [_norm_id(v) for v in frame[id_col]]
    emb = frame.drop(columns=[id_col])
    non_numeric = [c for c in emb.columns if not pd.api.types.is_numeric_dtype(emb[c])]
    if non_numeric:
        raise SubmissionError(
            f"embedding columns must be numeric; non-numeric: {non_numeric[:5]}"
        )
    if emb.shape[1] == 0:
        raise SubmissionError("submission has no embedding columns")
    out = emb.astype(float)
    out.index = ids
    if out.index.has_duplicates:
        dups = out.index[out.index.duplicated()].unique().tolist()
        raise SubmissionError(f"duplicate sample_ids in submission: {dups[:5]}")
    return out


def _align(labels: pd.DataFrame, emb: pd.DataFrame) -> np.ndarray:
    """Return the embedding matrix in labels order, or raise if a sample is missing."""
    lab_ids = [_norm_id(v) for v in labels[SAMPLE_ID]]
    present = set(emb.index)
    missing = [s for s in lab_ids if s not in present]
    if missing:
        raise SubmissionError(
            f"{len(missing)}/{len(lab_ids)} labelled samples missing from "
            f"submission, e.g. {missing[:5]}"
        )
    matrix = emb.loc[lab_ids].to_numpy(dtype=float)
    if not np.isfinite(matrix).all():
        raise SubmissionError("submission contains NaN or inf values")
    return matrix


def _inner_cv(y_train: np.ndarray) -> StratifiedKFold:
    counts = np.unique(y_train, return_counts=True)[1]
    splits = max(2, min(INNER_CV, int(counts.min())))
    return StratifiedKFold(n_splits=splits, shuffle=True, random_state=RANDOM_STATE)


def _fit_predict(
    task_type: str,
    x_train: np.ndarray,
    y_train: np.ndarray,
    x_test: np.ndarray,
    classes: np.ndarray,
) -> np.ndarray:
    """Fit the fixed linear probe on train, return predictions for test."""
    scaler = StandardScaler().fit(x_train)
    x_train = scaler.transform(x_train)
    x_test = scaler.transform(x_test)
    if task_type == "classification":
        if len(np.unique(y_train)) < 2:
            raise SubmissionError("a training fold has a single class")
        scoring = "roc_auc" if len(classes) == 2 else "roc_auc_ovr_weighted"
        clf = LogisticRegressionCV(
            Cs=LOGREG_CS,
            cv=_inner_cv(y_train),
            scoring=scoring,
            solver="lbfgs",
            max_iter=MAX_ITER,
            random_state=RANDOM_STATE,
        ).fit(x_train, y_train)
        proba = clf.predict_proba(x_test)
        col = {c: i for i, c in enumerate(clf.classes_)}
        out = np.zeros((x_test.shape[0], len(classes)))
        for j, c in enumerate(classes):
            if c in col:
                out[:, j] = proba[:, col[c]]
        return out
    reg = RidgeCV(alphas=RIDGE_ALPHAS).fit(x_train, y_train)
    return reg.predict(x_test)


def _fold_score(
    metric: str,
    task_type: str,
    y: np.ndarray,
    folds: np.ndarray,
    matrix: np.ndarray,
    classes: np.ndarray,
) -> float:
    """Return the mean native metric across one external CV partition."""
    scores = []
    for fold in sorted(np.unique(folds)):
        test = folds == fold
        predictions = _fit_predict(
            task_type, matrix[~test], y[~test], matrix[test], classes
        )
        try:
            scores.append(METRICS[metric](y[test], predictions, classes))
        except ValueError as error:
            raise SubmissionError(
                f"CV fold {fold} is not scoreable: {error}"
            ) from error
    if not np.isfinite(scores).all():
        raise SubmissionError("a CV fold produced a non-finite score")
    return float(np.mean(scores))


def _repeated_cv_scores(
    metric: str,
    task_type: str,
    y: np.ndarray,
    labels: pd.DataFrame,
    matrix: np.ndarray,
    classes: np.ndarray,
) -> tuple[float, ...]:
    """Return one mean fold metric for every external CV repeat."""
    scores = [
        _fold_score(metric, task_type, y, labels[column].to_numpy(), matrix, classes)
        for column in _fold_columns(labels)
    ]
    return tuple(float(score) for score in scores)


def _repeated_cv_score(
    metric: str,
    task_type: str,
    y: np.ndarray,
    labels: pd.DataFrame,
    matrix: np.ndarray,
    classes: np.ndarray,
) -> float:
    """Return the aggregate repeated-CV score for legacy callers."""
    return float(
        np.mean(_repeated_cv_scores(metric, task_type, y, labels, matrix, classes))
    )


def _transfer_predict(
    task_type: str,
    y: np.ndarray,
    split: np.ndarray,
    matrix: np.ndarray,
    classes: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """Fit once on the train cohort, predict the test cohort (no pooling)."""
    train, test = split == SPLIT_TRAIN, split == SPLIT_TEST
    if not train.any() or not test.any():
        raise EvaluatorError("transfer labels.csv has an empty train or test side")
    preds = _fit_predict(task_type, matrix[train], y[train], matrix[test], classes)
    return y[test], preds


def _select_degs(delta_train: np.ndarray, n_top: int) -> np.ndarray:
    """Select the strongest mean absolute responses using training rows only."""
    n_keep = min(int(n_top), delta_train.shape[1])
    if n_keep <= 0:
        raise EvaluatorError("perturbation task has no response genes to select")
    effect = np.abs(np.mean(delta_train, axis=0))
    return np.argsort(effect, kind="stable")[-n_keep:]


def _mean_scores(scores: list[float]) -> float:
    """Average finite fold scores."""
    finite = np.asarray(scores)[np.isfinite(scores)]
    return float(np.mean(finite)) if len(finite) else float("nan")


def _perturbation_fold(
    matrix: np.ndarray,
    delta: np.ndarray,
    train: np.ndarray,
    test: np.ndarray,
    n_top: int,
    target_centered: bool,
) -> float:
    """Fit and score one leakage-safe perturbation fold."""
    if not train.any() or not test.any():
        raise EvaluatorError("perturbation fold has an empty train or test side")
    genes = _select_degs(delta[train], n_top)
    y_train = delta[train][:, genes]
    y_test = delta[test][:, genes]
    scaler = StandardScaler().fit(matrix[train])
    x_train = scaler.transform(matrix[train])
    x_test = np.clip(
        scaler.transform(matrix[test]),
        -PERTURBATION_TEST_Z_CLIP,
        PERTURBATION_TEST_Z_CLIP,
    )
    predictions = RidgeCV(alphas=RIDGE_ALPHAS).fit(x_train, y_train).predict(x_test)
    if target_centered:
        correlation = compute_target_centered_sample_spearman(y_test, predictions)
    else:
        training_mean = np.mean(y_train, axis=0)
        correlation = compute_residual_sample_spearman(
            y_test, predictions, training_mean
        )
    if not np.isfinite(correlation):
        raise SubmissionError("a perturbation fold has a degenerate Spearman score")
    return float(0.5 * (1.0 + np.clip(correlation, -1.0, 1.0)))


def _paired_perturbation_scores(
    matrix: np.ndarray,
    delta: np.ndarray,
    labels: pd.DataFrame,
    n_top: int,
) -> tuple[float, ...]:
    """Score three frozen subject-grouped CV partitions."""
    repeat_scores = []
    for column in _fold_columns(labels):
        fold_metrics = []
        folds = labels[column].to_numpy()
        for fold in sorted(np.unique(folds)):
            test = folds == fold
            fold_metrics.append(
                _perturbation_fold(matrix, delta, ~test, test, n_top, False)
            )
        repeat_scores.append(_mean_scores(fold_metrics))
    return tuple(repeat_scores)


def _transfer_perturbation_scores(
    matrix: np.ndarray,
    delta: np.ndarray,
    labels: pd.DataFrame,
    n_top: int,
    target_centered: bool,
) -> tuple[tuple[float, ...], int]:
    """Train every response-decoding component on the source disease only."""
    split = labels[SPLIT].astype(str).to_numpy()
    train, test = split == SPLIT_TRAIN, split == SPLIT_TEST
    metrics = _perturbation_fold(matrix, delta, train, test, n_top, target_centered)
    return (metrics,), int(test.sum())


def _cartesian_partitions(
    control_ids: np.ndarray, perturbed_ids: np.ndarray
) -> list[list[tuple[np.ndarray, np.ndarray]]]:
    """Build six 3-fold specimen-isolated partitions for a 3x3 response grid."""
    from itertools import permutations

    controls = np.unique(control_ids)
    perturbed = np.unique(perturbed_ids)
    if len(controls) != 3 or len(perturbed) != 3:
        raise EvaluatorError(
            "Cartesian perturbation evaluation requires exactly three controls "
            "and three perturbed specimens."
        )
    repeats = []
    for assignment in permutations(perturbed):
        folds = []
        for control_id, perturbed_id in zip(controls, assignment):
            test = (control_ids == control_id) & (perturbed_ids == perturbed_id)
            train = (control_ids != control_id) & (perturbed_ids != perturbed_id)
            if test.sum() != 1 or train.sum() != 4:
                raise EvaluatorError("Cartesian labels do not form a complete 3x3 grid")
            folds.append((train, test))
        repeats.append(folds)
    return repeats


def _cartesian_perturbation_scores(
    matrix: np.ndarray,
    delta: np.ndarray,
    labels: pd.DataFrame,
    n_top: int,
) -> tuple[float, ...]:
    """Evaluate all control-to-perturbed permutations without specimen leakage."""
    repeats = _cartesian_partitions(
        labels[CONTROL_ID].astype(str).to_numpy(),
        labels[PERTURBED_ID].astype(str).to_numpy(),
    )
    repeat_scores = []
    for folds in repeats:
        fold_metrics = []
        for train, test in folds:
            fold_metrics.append(
                _perturbation_fold(matrix, delta, train, test, n_top, False)
            )
        repeat_scores.append(_mean_scores(fold_metrics))
    return tuple(repeat_scores)


def _score_perturbation_task(
    task: dict,
    labels: pd.DataFrame,
    emb: pd.DataFrame,
    targets: PerturbationTargets,
) -> TaskScore:
    """Decode one hidden response matrix from submitted baseline embeddings."""
    if len(labels) != len(targets.delta):
        raise EvaluatorError("labels.csv and targets.npz have different row counts")
    matrix = _align(labels, emb)
    n_top = int(task.get("n_top_de_genes", 0))
    pairing = str(task.get(PAIRING, labels.get(PAIRING, pd.Series([PAIRED])).iloc[0]))
    if SPLIT in labels:
        repeat_scores, n_samples = _transfer_perturbation_scores(
            matrix,
            targets.delta,
            labels,
            n_top,
            bool(task.get("target_centered", False)),
        )
    elif pairing == CARTESIAN:
        repeat_scores = _cartesian_perturbation_scores(
            matrix, targets.delta, labels, n_top
        )
        n_samples = len(labels)
    else:
        repeat_scores = _paired_perturbation_scores(
            matrix, targets.delta, labels, n_top
        )
        n_samples = len(labels)
    return TaskScore(
        task_id=_norm_id(task[TASK_ID]),
        dataset_id=_norm_id(task[DATASET_ID]),
        category=str(task["category"]),
        metric=str(task.get("metric", RESIDUAL_SPEARMAN)),
        score=float(np.mean(repeat_scores)),
        n_samples=n_samples,
        repeat_scores=repeat_scores,
    )


def score_task(
    task: dict,
    labels: pd.DataFrame,
    emb: pd.DataFrame,
    targets: PerturbationTargets | None = None,
) -> TaskScore:
    """Score one task: repeated fold-wise CV, or a fixed transfer split.

    Private biology (disease, tissue, and target names) never leaks in.
    """
    metric = task["metric"]
    if task.get("task_type") == PERTURBATION:
        if metric not in {CENTERED_SPEARMAN, RESIDUAL_SPEARMAN}:
            raise EvaluatorError(
                f"perturbation task {task.get(TASK_ID)} needs metric "
                f"'{CENTERED_SPEARMAN}'"
            )
        if targets is None:
            raise EvaluatorError("perturbation task is missing targets.npz")
        return _score_perturbation_task(task, labels, emb, targets)
    if metric not in METRICS:
        raise EvaluatorError(f"unknown metric '{metric}' for task {task.get(TASK_ID)}")
    task_type = task["task_type"]
    matrix = _align(labels, emb)
    y = labels[LABEL].to_numpy()
    classes = np.unique(y) if task_type == "classification" else np.array([])

    if SPLIT in labels.columns:
        y_eval, preds = _transfer_predict(
            task_type, y, labels[SPLIT].to_numpy(), matrix, classes
        )
        try:
            score = METRICS[metric](y_eval, preds, classes)
        except ValueError as error:
            raise SubmissionError(f"transfer task not scoreable: {error}") from error
        repeat_scores = (score,)
    else:
        repeat_scores = _repeated_cv_scores(
            task_type=task_type,
            metric=metric,
            y=y,
            labels=labels,
            matrix=matrix,
            classes=classes,
        )
        score = float(np.mean(repeat_scores))
        y_eval = y

    return TaskScore(
        task_id=_norm_id(task[TASK_ID]),
        dataset_id=_norm_id(task[DATASET_ID]),
        category=str(task["category"]),
        metric=metric,
        score=score,
        n_samples=int(len(y_eval)),
        repeat_scores=repeat_scores,
    )


def fetch_manifest(token: str | None = None) -> list[dict]:
    """Download the public ``datasets.yaml`` registry of opaque dataset ids."""
    from huggingface_hub import hf_hub_download

    path = hf_hub_download(
        PUBLIC_REPO, MANIFEST_FILENAME, repo_type="dataset", token=token
    )
    with open(path) as handle:
        data = yaml.safe_load(handle)
    if isinstance(data, dict):
        return data.get("datasets", [])
    return data or []


def manifest_ids(manifest: list[dict]) -> set[str]:
    """The canonical set of valid dataset ids from the manifest."""
    return {_norm_id(entry["id"]) for entry in manifest}


def scoreable_tasks(tasks: list[dict], valid_ids: set[str]) -> list[dict]:
    """Registry tasks whose dataset is in the public manifest; orphans are skipped.

    Shared by ``score_all`` and the leaderboard so both agree on what full
    coverage means. An orphan (a task on a dataset not yet published) is logged
    and dropped rather than making full coverage unreachable for everyone.
    """
    kept, orphans = [], []
    for task in tasks:
        if _norm_id(task[DATASET_ID]) in valid_ids:
            kept.append(task)
        else:
            orphans.append(_norm_id(task[TASK_ID]))
    if orphans:
        logging.warning("skipping tasks on unpublished datasets: %s", orphans)
    return kept


def fetch_tasks_registry(token: str | None = None) -> list[dict]:
    """Download the private ``tasks.yaml`` registry from the labels repo."""
    from huggingface_hub import hf_hub_download

    path = hf_hub_download(
        LABELS_REPO, TASKS_FILENAME, repo_type="dataset", token=token
    )
    return load_tasks_registry(path)


def fetch_task_labels(task_id: str, token: str | None = None) -> pd.DataFrame:
    """Download a task's private ``labels.csv`` from the labels repo."""
    from huggingface_hub import hf_hub_download

    path = hf_hub_download(
        LABELS_REPO, f"{task_id}/{LABELS_FILENAME}", repo_type="dataset", token=token
    )
    return load_labels(path)


def fetch_task_targets(task_id: str, token: str | None = None) -> PerturbationTargets:
    """Download a task's private compressed response matrix."""
    from huggingface_hub import hf_hub_download

    path = hf_hub_download(
        LABELS_REPO, f"{task_id}/{TARGETS_FILENAME}", repo_type="dataset", token=token
    )
    return load_targets(path)


def _retry(fn):
    """Call ``fn``, retrying transient failures with linear backoff."""
    last: Exception | None = None
    for attempt in range(FETCH_ATTEMPTS):
        try:
            return fn()
        except Exception as error:  # noqa: BLE001
            last = error
            if attempt < FETCH_ATTEMPTS - 1:
                time.sleep(FETCH_BACKOFF * (attempt + 1))
    raise last


def _task_outcome(
    task: dict,
    emb,
    block_error,
    fetch_labels,
    fetch_targets,
    token,
) -> TaskOutcome:
    """Score one task, mapping each failure to the right category."""
    task_id = _norm_id(task[TASK_ID])
    dataset_id = _norm_id(task[DATASET_ID])
    if emb is None and block_error is None:
        return TaskOutcome(task_id, dataset_id, "missing")
    if block_error is not None:
        return TaskOutcome(task_id, dataset_id, "invalid", reason=block_error)
    try:
        labels = _retry(lambda: fetch_labels(task_id, token))
    except Exception as error:  # noqa: BLE001
        raise EvaluatorError(
            f"could not load evaluation data for a task: {error}"
        ) from error
    targets = None
    if task.get("task_type") == PERTURBATION:
        try:
            targets = _retry(lambda: fetch_targets(task_id, token))
        except Exception as error:  # noqa: BLE001
            raise EvaluatorError(
                f"could not load perturbation targets for a task: {error}"
            ) from error
    try:
        score = score_task(task, labels, emb, targets)
    except SubmissionError as error:
        return TaskOutcome(task_id, dataset_id, "invalid", reason=str(error))
    if not np.isfinite(score.score):
        return TaskOutcome(
            task_id,
            dataset_id,
            "invalid",
            reason="degenerate score (constant predictions)",
        )
    return TaskOutcome(task_id, dataset_id, "scored", score=score)


def score_all(
    path: str | Path,
    token: str | None = None,
    *,
    datasets: list[dict] | None = None,
    tasks: list[dict] | None = None,
    fetch_labels=None,
    fetch_targets=None,
) -> dict:
    """Score a submission per task and roll it up into a blind result.

    Unknown dataset ids and bad files raise/record a ``SubmissionError``; a
    failed fetch or internal error raises ``EvaluatorError`` (our side, not the
    submitter's). A dataset is embedded once and reused across all its tasks.
    """
    blocks = load_submission(path)
    if datasets is None:
        try:
            datasets = fetch_manifest(token)
        except Exception as error:  # noqa: BLE001
            raise EvaluatorError(
                f"could not load the dataset manifest: {error}"
            ) from error
    valid = manifest_ids(datasets)

    unknown = sorted(set(blocks) - valid)
    if unknown:
        raise SubmissionError(f"unknown dataset_id(s) not in the benchmark: {unknown}")

    if tasks is None:
        try:
            tasks = fetch_tasks_registry(token)
        except Exception as error:  # noqa: BLE001
            raise EvaluatorError(
                f"could not load the task registry: {error}"
            ) from error
    fetch_labels = fetch_labels or fetch_task_labels
    fetch_targets = fetch_targets or fetch_task_targets

    emb_by_ds, block_error = {}, {}
    for dataset_id, block in blocks.items():
        try:
            emb_by_ds[dataset_id] = _to_embedding_frame(block)
        except SubmissionError as error:
            block_error[dataset_id] = str(error)

    outcomes = [
        _task_outcome(
            task,
            emb_by_ds.get(_norm_id(task[DATASET_ID])),
            block_error.get(_norm_id(task[DATASET_ID])),
            fetch_labels,
            fetch_targets,
            token,
        )
        for task in scoreable_tasks(tasks, valid)
    ]
    return _summarize(outcomes)


def _summarize(outcomes: list[TaskOutcome]) -> dict:
    """Combine per-task outcomes into blind per-category numbers + coverage."""
    scored = [o.score for o in outcomes if o.status == "scored"]
    dataset_status = _dataset_status(outcomes)
    n_scored, n_total = len(scored), len(outcomes)
    n_ds_scored = sum(1 for status in dataset_status.values() if status == "scored")
    return {
        "per_task": scored,
        "categories": category_means(scored),
        "invalid": [
            {"task_id": o.task_id, "dataset_id": o.dataset_id, "reason": o.reason}
            for o in outcomes
            if o.status == "invalid"
        ],
        "missing": sorted({o.dataset_id for o in outcomes if o.status == "missing"}),
        "incomplete": sorted(
            ds for ds, status in dataset_status.items() if status == "incomplete"
        ),
        "coverage": n_scored / n_total if n_total else 0.0,
        "n_scored": n_scored,
        "n_total": n_total,
        "full_coverage": n_total > 0 and n_scored == n_total,
        "n_datasets_scored": n_ds_scored,
        "n_datasets_total": len(dataset_status),
    }


def _dataset_status(outcomes: list[TaskOutcome]) -> dict[str, str]:
    """Blind per-dataset STATUS only (no score): scored / missing / incomplete."""
    statuses: dict[str, set[str]] = defaultdict(set)
    for outcome in outcomes:
        statuses[outcome.dataset_id].add(outcome.status)
    out = {}
    for dataset_id in sorted(statuses):
        seen = statuses[dataset_id]
        if seen == {"missing"}:
            out[dataset_id] = "missing"
        elif seen == {"scored"}:
            out[dataset_id] = "scored"
        else:
            out[dataset_id] = "incomplete"
    return out


def _cli() -> None:
    parser = argparse.ArgumentParser(description="Score a PRIMO submission locally.")
    parser.add_argument("--submission", required=True, help="CSV/TSV/Parquet/NPZ file")
    parser.add_argument("--token", default=None, help="HF token (else env HF_TOKEN)")
    args = parser.parse_args()

    token = args.token or os.environ.get("HF_TOKEN")
    result = score_all(args.submission, token)
    print(
        f"scored     : {result['n_datasets_scored']}/{result['n_datasets_total']} "
        f"datasets, {result['n_scored']}/{result['n_total']} tasks  "
        f"(full_coverage={result['full_coverage']})"
    )
    for category, stats in sorted(result["categories"].items()):
        print(
            f"  {category:20s} {stats['metric']} = {stats['mean']:.4f} "
            f"(n={stats['n_tasks']})"
        )
    for bad in result["invalid"]:
        print(f"  INVALID [{bad['dataset_id']}]: {bad['reason']}")
    if result["missing"]:
        print(f"  missing    : {result['missing']}")
    if result["incomplete"]:
        print(f"  incomplete : {result['incomplete']}")


if __name__ == "__main__":
    _cli()