"""Shared data-selection + validated-metric helpers for the heads trainer.

Born in the head-vs-centroid eval harness (#1130, tag_eval.py) that proved the
"frozen embedding + small trained head (with negatives)" spine; the harness was
retired 2026-07-02 (operator: the tagging system is proven, the eval isn't
needed) and these survivors moved here — they ARE the heads' production data
pipeline (heads.py trains and scores with them nightly).

numpy/scikit-learn are imported lazily inside the functions that need them so
the API worker (base image, no ML stack) can import this module.
"""

from __future__ import annotations

from typing import Any

from sqlalchemy import func, select
from sqlalchemy.orm import Session

from ...models import ImageRecord, TagSuggestionRejection
from ...models.tag import image_tag


def _ids_with_tag(session: Session, tag_id: int) -> list[int]:
    return [
        r[0] for r in session.execute(
            select(image_tag.c.image_record_id).where(image_tag.c.tag_id == tag_id)
        ).all()
    ]


def _rejected_ids(session: Session, tag_id: int) -> list[int]:
    return [
        r[0] for r in session.execute(
            select(TagSuggestionRejection.image_record_id)
            .where(TagSuggestionRejection.tag_id == tag_id)
        ).all()
    ]


def _sample_unlabeled(session: Session, exclude: set[int], limit: int) -> list[int]:
    """Random image ids (with an embedding) NOT carrying the tag. Concepts are
    sparse, so an untagged image is almost always a true negative."""
    stmt = (
        select(ImageRecord.id)
        .where(ImageRecord.siglip_embedding.is_not(None))
        .order_by(func.random())
        .limit(limit)
    )
    if exclude:
        stmt = stmt.where(ImageRecord.id.not_in(exclude))
    return [r[0] for r in session.execute(stmt).all()]


def _load_embeddings(session: Session, ids: list[int]) -> dict[int, Any]:
    import numpy as np

    out: dict[int, Any] = {}
    if not ids:
        return out
    # Chunk the IN list to stay well under psycopg's parameter ceiling.
    for i in range(0, len(ids), 2000):
        chunk = ids[i:i + 2000]
        for rid, emb in session.execute(
            select(ImageRecord.id, ImageRecord.siglip_embedding)
            .where(ImageRecord.id.in_(chunk))
            .where(ImageRecord.siglip_embedding.is_not(None))
        ).all():
            out[rid] = np.asarray(emb, dtype=np.float32)
    return out


def _l2norm(X, np):
    n = np.linalg.norm(X, axis=1, keepdims=True)
    n[n == 0] = 1.0
    return X / n


def _metrics_from_scores(y, scores, np) -> dict[str, float]:
    from sklearn.metrics import average_precision_score, precision_recall_curve

    ap = float(average_precision_score(y, scores))
    prec, rec, thr = precision_recall_curve(y, scores)
    f1 = (2 * prec * rec) / np.clip(prec + rec, 1e-9, None)
    best = int(np.argmax(f1))
    # thr has len = len(prec)-1; map best index safely.
    t = float(thr[min(best, len(thr) - 1)]) if len(thr) else 0.5
    return {
        "ap": round(ap, 4),
        "precision": round(float(prec[best]), 4),
        "recall": round(float(rec[best]), 4),
        "f1": round(float(f1[best]), 4),
        "threshold": round(t, 4),
    }


def _safe_folds(y, folds, np) -> int:
    minority = int(min(np.bincount(y)))
    return max(2, min(folds, minority))


def _auto_apply_point(y, scores, target, np) -> dict | None:
    """The auto-apply operating point: the threshold that yields the MOST recall
    while holding precision >= target. This answers 'could this concept fire
    without a human, and how much would it catch?' Returns None if no threshold
    reaches the precision target (concept not auto-apply-ready)."""
    from sklearn.metrics import precision_recall_curve

    prec, rec, thr = precision_recall_curve(y, scores)
    best = None  # (threshold, precision, recall) maximizing recall s.t. prec>=target
    for i in range(len(thr)):  # thr[i] corresponds to prec[i], rec[i]
        if prec[i] >= target and (best is None or rec[i] > best[2]):
            best = (float(thr[i]), float(prec[i]), float(rec[i]))
    if best is None:
        return None
    return {
        "target": round(float(target), 4),
        "threshold": round(best[0], 4),
        "precision": round(best[1], 4),
        "recall": round(best[2], 4),
    }