"""Frozen paper score definitions, including full-vocabulary worst-class F1.

Copied from the paper artifact's wristharbench/scoring.py; no baseline code is
needed. Macro F1 averages classes observed in truth or prediction; worst-class
F1 considers the full dataset vocabulary, including absent classes.
"""

from __future__ import annotations

from typing import Any


def from_confusion(confusion: Any) -> dict[str, Any]:
    import numpy as np

    cm = np.asarray(confusion, dtype=float)
    if cm.ndim != 2 or cm.shape[0] != cm.shape[1] or not np.isfinite(cm).all() or (cm < 0).any() or cm.sum() <= 0:
        raise ValueError("Expected a nonempty square finite nonnegative confusion matrix")
    truth, predicted = cm.sum(1), cm.sum(0)
    tp = np.diag(cm)
    denominator = truth + predicted
    per_class = np.divide(2 * tp, denominator, out=np.zeros(len(tp)), where=denominator > 0)
    recall = np.divide(tp, truth, out=np.zeros(len(tp)), where=truth > 0)
    return {"macro_f1": float(per_class[denominator > 0].mean()),
            "accuracy": float(tp.sum() / cm.sum()),
            "balanced_accuracy": float(recall[truth > 0].mean()),
            "worst_class_f1": float(per_class.min())}


def score(y_true: Any, y_pred: Any, n_classes: int) -> dict[str, Any]:
    import numpy as np
    from sklearn.metrics import confusion_matrix

    truth, prediction = np.asarray(y_true), np.asarray(y_pred)
    if n_classes < 1 or truth.ndim != 1 or truth.shape != prediction.shape or not len(truth):
        raise ValueError("Expected aligned nonempty labels and a positive global class count")
    if not np.issubdtype(truth.dtype, np.integer) or not np.issubdtype(prediction.dtype, np.integer):
        raise ValueError("Labels must be encoded integers")
    if min(truth.min(), prediction.min()) < 0 or max(truth.max(), prediction.max()) >= n_classes:
        raise ValueError("Labels fall outside the global class vocabulary")
    cm = confusion_matrix(truth, prediction, labels=np.arange(n_classes))
    return {**from_confusion(cm), "confusion_matrix": cm.tolist()}