"""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 linear probe, pooled out-of-fold over the frozen CV folds or fit once on a
fixed train/test split (transfer), and produces predictions. 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 and each task's private ``labels.csv``. 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
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"
SPLIT = "split"
SPLIT_TRAIN = "train"
SPLIT_TEST = "test"

ORG = "ScientaLab"
PUBLIC_REPO = f"{ORG}/primo"
LABELS_REPO = f"{ORG}/primo-labels"
RESULTS_REPO = f"{ORG}/primo-results"
MANIFEST_FILENAME = "datasets.yaml"
TASKS_FILENAME = "tasks.yaml"
LABELS_FILENAME = "labels.csv"

RIDGE_ALPHAS = np.logspace(-3.0, 6.0, 19)
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


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 load_labels(path: str | Path) -> pd.DataFrame:
    """Read a private ``labels.csv`` (sample_id, label, and one of fold/split)."""
    df = pd.read_csv(path)
    missing = {SAMPLE_ID, LABEL} - set(df.columns)
    if missing:
        raise EvaluatorError(f"labels.csv missing columns: {sorted(missing)}")
    if (FOLD in df.columns) == (SPLIT in df.columns):
        raise EvaluatorError(
            f"labels.csv needs exactly one of '{FOLD}' (CV) or '{SPLIT}' (transfer)."
        )
    return df


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 _oof_predict(
    task_type: str,
    y: np.ndarray,
    folds: np.ndarray,
    matrix: np.ndarray,
    classes: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """Pooled out-of-fold predictions over the frozen CV folds."""
    if task_type == "classification":
        oof = np.zeros((len(y), len(classes)))
    else:
        oof = np.zeros(len(y))
    for fold in sorted(np.unique(folds)):
        test = folds == fold
        oof[test] = _fit_predict(
            task_type, matrix[~test], y[~test], matrix[test], classes
        )
    return y, oof


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 score_task(task: dict, labels: pd.DataFrame, emb: pd.DataFrame) -> TaskScore:
    """Score one task: pooled out-of-fold CV, or a fixed train/test transfer split.

    Reads only ``task_id``, ``dataset_id``, ``metric``, ``task_type`` and
    ``category``. Private biology (disease/tissue/target) never leaks in.
    """
    metric = task["metric"]
    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
    else:
        y_eval, preds = _oof_predict(
            task_type, y, labels[FOLD].to_numpy(), matrix, classes
        )
        score = METRICS[metric](y_eval, preds, classes)

    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)),
    )


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 _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, 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
    try:
        score = score_task(task, labels, emb)
    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,
) -> 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

    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,
            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()