"""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,
    Tag,
    TagPositiveConfirmation,
    TagSuggestionRejection,
)
from ...models.tag import image_tag

# Auto-apply sources whose tags are PROVISIONAL: they never train a head (or seed
# a CCIP reference) unless the operator confirms them (milestone 139). Keeping
# auto-applied predictions out of training is what makes them "soft" — a misfire
# can't reinforce itself, so the retraction sweep can actually drop it.
_AUTO_SOURCES = ("head_auto", "ccip_auto", "ml_auto", "presentation_auto")


def _hygiene_excluded_ids(session: Session) -> set[int]:
    """Ids of images carrying ANY system tag (wip / banner / editor
    screenshot — milestone #128). These images are excluded from OTHER
    concepts' head training entirely: not positives (a rough wip tagged as a
    character drags that head toward 'generic sketch') and not sampled or
    rejection negatives (a wip OF character X is not evidence against X) —
    simply absent. A system tag's OWN head trains on them unchanged; that is
    what makes auto-flagging banners/editor screenshots work.

    Item-level by design: a wip-tagged process video contributes (or
    withholds) ALL its sampled frames, though some may show the finished
    piece. Operator call 2026-07-03: with enough clean data this washes out —
    no per-frame handling.
    """
    return set(
        session.execute(
            select(image_tag.c.image_record_id)
            .join(Tag, Tag.id == image_tag.c.tag_id)
            .where(Tag.is_system.is_(True))
        ).scalars().all()
    )


def _ids_with_tag(session: Session, tag_id: int) -> list[int]:
    """Image ids that count as POSITIVES for this tag's head: human-applied
    (manual / accepted) tags PLUS any auto-applied tag the operator explicitly
    confirmed (TagPositiveConfirmation). Unconfirmed auto-applied tags are
    EXCLUDED — they are provisional and must not train the head that judges
    them (milestone 139)."""
    confirmed = (
        select(TagPositiveConfirmation.image_record_id)
        .where(TagPositiveConfirmation.tag_id == tag_id)
    )
    return [
        r[0] for r in session.execute(
            select(image_tag.c.image_record_id)
            .where(image_tag.c.tag_id == tag_id)
            .where(
                image_tag.c.source.not_in(_AUTO_SOURCES)
                | image_tag.c.image_record_id.in_(confirmed)
            )
        ).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),
    }