FabledCurator/backend/app/services/ml/heads.py

"""Production heads: train + score the per-concept classifiers (#114).

The eval harness (#1130) proved the spine, then retired 2026-07-02 once the
tagging system was accepted; this is the production form.
- TRAIN (sync, ml worker — needs scikit-learn): for every general/character tag
  with enough labelled positives, fit a logistic-regression head on the FROZEN
  SigLIP embeddings (positives + negatives = rejections + sampled unlabeled),
  derive an honest suggest threshold + earned-auto-apply point from CROSS-
  VALIDATED scores, and upsert a TagHead row. Uses the eval-proven data loaders
  + metric helpers (training_data.py) so heads match the measured numbers.
- SCORE (async, API worker — numpy via pgvector, NO scikit-learn): score one
  image's embedding against all current heads → the suggestions the rail shows,
  REPLACING Camie predictions + per-tag centroids.

scikit-learn is imported lazily inside the train path so the API worker can still
import this module to enqueue training + to score (scoring needs only numpy).
"""

from __future__ import annotations

import logging
from datetime import UTC, datetime
from typing import Any

from sqlalchemy import delete, exists, func, select
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Session

from ...models import (
    HeadAutoApplyRun,
    HeadTrainingRun,
    ImageRecord,
    ImageRegion,
    MLSettings,
    PresentationReview,
    Tag,
    TagHead,
    TagKind,
    TagPositiveConfirmation,
    TagSuggestionRejection,
)
from ...models.tag import CHROME_SYSTEM_TAGS, PROCESS_SYSTEM_TAGS, image_tag
from .training_data import (
    _AUTO_SOURCES,
    _applied_or_rejected,
    _auto_apply_point,
    _hygiene_excluded_ids,
    _ids_with_tag,
    _l2norm,
    _load_embeddings,
    _metrics_from_scores,
    _rejected_ids,
    _safe_folds,
    _sample_unlabeled,
)

log = logging.getLogger(__name__)

DEFAULT_NEG_RATIO = 3
DEFAULT_CV_FOLDS = 5
MIN_POSITIVES_FLOOR = 8     # hard floor; settings.head_min_positives can raise it
_UNLABELED_POOL = 4000
_EXAMPLES_MIN = 8           # need at least this many embedded +/- to fit a head

# Auto-apply / match confidence operating range. Every graduated auto-apply or
# CCIP-match threshold the operator can set lives in this band, and the head
# precision target is clamped to it: below 0.5 "auto-apply" is meaningless, and
# 1.0 is unachievable so 0.999 is the ceiling. One source shared by the service
# clamp (_normalize_params) and the API validator (ml_admin._validate).
AUTO_APPLY_THRESHOLD_MIN = 0.5
AUTO_APPLY_THRESHOLD_MAX = 0.999

# Only these tag kinds get heads (the surfaced suggestion categories).
_HEAD_KINDS = (TagKind.general, TagKind.character)
# tag.kind -> the suggestion category the rail groups under.
_CATEGORY = {TagKind.general: "general", TagKind.character: "character"}

# System-tag (wip/banner/editor screenshot) heads surface as suggestions at
# this FLAT confidence floor instead of their auto-derived (precision-tuned)
# suggest threshold. The auto threshold is high, so it hides the borderline /
# false-positive guesses — which are exactly the cases the operator wants to
# SEE and REJECT to sharpen these heads (hard-negative mining: "negatively
# reinforce what isn't a system tag"). Operator-set 0.65 (2026-07-03): high
# enough not to spam near-zero scores, low enough to surface real mistakes.
# Content-tag heads keep their own thresholds; the typed-dropdown's
# threshold_override still overrides everything (show-all mode).
_SYSTEM_TAG_SUGGEST_FLOOR = 0.65


def _sigmoid(z, np):
    """Logistic sigmoid 1/(1+e^-z): the head score→probability transform. One home
    for what was inlined at every scoring site (suggest, both sweeps, retract)."""
    return 1.0 / (1.0 + np.exp(-z))


def _conflict_scores(Xn, Wc, bc, np):
    """The presentation conflict signal (#141): per row, the MAX content-head
    probability and WHICH head produced it. Shared by the system-tag sweep's guard-2
    and the soft-wip audit — both ask "does this ALSO look like real content?"."""
    cprobs = _sigmoid(Xn @ Wc.T + bc, np)
    return cprobs.max(axis=1), cprobs.argmax(axis=1)


def _insert_presentation_review(
    session, *, image_record_id, tag_id, conflict_tag_id, conflict_score, mode,
):
    """Single-source the ring-loud PresentationReview row shape so the two writers
    (system-tag sweep guard-2 + soft-wip audit) can't drift on columns or `mode` —
    they share the (image_record_id, tag_id) composite PK, so a divergent `mode`
    would be a silent first-writer-wins bug."""
    session.execute(
        pg_insert(PresentationReview)
        .values(
            image_record_id=image_record_id, tag_id=tag_id,
            conflict_tag_id=conflict_tag_id, conflict_score=conflict_score,
            mode=mode,
        )
        .on_conflict_do_nothing()
    )


class HeadTrainingAlreadyRunning(Exception):
    """Raised by start_head_training_run when a run is already in flight."""


def start_head_training_run(session: Session, params: dict[str, Any]) -> int:
    """Create a HeadTrainingRun (status='running') + dispatch the ml-queue task.
    Returns the run id. One training run at a time (light guard)."""
    existing = session.execute(
        select(HeadTrainingRun.id).where(HeadTrainingRun.status == "running")
    ).scalar_one_or_none()
    if existing is not None:
        raise HeadTrainingAlreadyRunning(existing)
    norm = _normalize_params(session, params)
    run = HeadTrainingRun(
        params=norm, status="running", last_progress_at=datetime.now(UTC)
    )
    session.add(run)
    session.flush()
    run_id = run.id
    from ...tasks.ml import train_heads as _task
    _task.delay(run_id)
    return run_id


def _settings(session: Session) -> MLSettings:
    return MLSettings.load_sync(session)


def _normalize_params(session: Session, params: dict[str, Any] | None) -> dict[str, Any]:
    params = params or {}
    s = _settings(session)
    try:
        min_pos = max(MIN_POSITIVES_FLOOR, int(params.get("min_positives", s.head_min_positives)))
    except (TypeError, ValueError):
        min_pos = max(MIN_POSITIVES_FLOOR, s.head_min_positives)
    try:
        neg_ratio = max(1, int(params.get("neg_ratio", DEFAULT_NEG_RATIO)))
    except (TypeError, ValueError):
        neg_ratio = DEFAULT_NEG_RATIO
    try:
        cv_folds = max(2, int(params.get("cv_folds", DEFAULT_CV_FOLDS)))
    except (TypeError, ValueError):
        cv_folds = DEFAULT_CV_FOLDS
    try:
        precision_target = min(max(float(params.get("precision_target", s.head_auto_apply_precision)), AUTO_APPLY_THRESHOLD_MIN), AUTO_APPLY_THRESHOLD_MAX)
    except (TypeError, ValueError):
        precision_target = s.head_auto_apply_precision
    return {
        "min_positives": min_pos,
        "neg_ratio": neg_ratio,
        "cv_folds": cv_folds,
        "precision_target": round(precision_target, 4),
    }


def _embedder_version(session: Session) -> str:
    return _settings(session).embedder_model_version


def _eligible_tag_ids(session: Session, min_pos: int) -> list[int]:
    """Concept tags (general/character) with >= min_pos POSITIVE images — the set
    that gets a head. Counts human-applied + operator-confirmed tags only;
    unconfirmed auto-applied predictions do NOT count toward eligibility (they
    don't train the head — milestone 139), so a concept can't graduate on its own
    guesses."""
    confirmed = exists().where(
        TagPositiveConfirmation.image_record_id == image_tag.c.image_record_id,
        TagPositiveConfirmation.tag_id == image_tag.c.tag_id,
    )
    rows = session.execute(
        select(Tag.id)
        .join(image_tag, image_tag.c.tag_id == Tag.id)
        .where(Tag.kind.in_(_HEAD_KINDS))
        .where(image_tag.c.source.not_in(_AUTO_SOURCES) | confirmed)
        .group_by(Tag.id)
        .having(func.count(image_tag.c.image_record_id) >= min_pos)
    ).all()
    return [r[0] for r in rows]


def _head_fingerprints(session: Session, tag_ids: list[int]) -> dict[int, str]:
    """Per-tag training-data fingerprint: (positive count, latest positive
    created_at) + (rejection count, latest rejected_at). It moves whenever a tag
    gains/loses a positive or a rejection — the incremental-retrain change
    detector (#1317 p2). A newly-added positive/rejection always has the latest
    timestamp, so even a remove-one-add-one (unchanged count) is caught. The
    sampled-unlabeled negative pool + the hygiene set drift GLOBALLY and are
    reconciled by the nightly full run, not captured here."""
    if not tag_ids:
        return {}
    pos = session.execute(
        select(
            image_tag.c.tag_id,
            func.count(image_tag.c.image_record_id),
            func.max(image_tag.c.created_at),
        )
        .where(image_tag.c.tag_id.in_(tag_ids))
        .group_by(image_tag.c.tag_id)
    ).all()
    pos_map = {t: (c, m) for t, c, m in pos}
    rej = session.execute(
        select(
            TagSuggestionRejection.tag_id,
            func.count(),
            func.max(TagSuggestionRejection.rejected_at),
        )
        .where(TagSuggestionRejection.tag_id.in_(tag_ids))
        .group_by(TagSuggestionRejection.tag_id)
    ).all()
    rej_map = {t: (c, m) for t, c, m in rej}
    # Confirmations promote an auto-applied tag to a positive (milestone 139), so
    # a confirm must move the fingerprint too — else a manual Retrain right after
    # confirming wouldn't fold the tag in (the nightly full run would).
    conf = session.execute(
        select(TagPositiveConfirmation.tag_id, func.count())
        .where(TagPositiveConfirmation.tag_id.in_(tag_ids))
        .group_by(TagPositiveConfirmation.tag_id)
    ).all()
    conf_map = dict(conf)
    out = {}
    for t in tag_ids:
        pc, pm = pos_map.get(t, (0, None))
        rc, rm = rej_map.get(t, (0, None))
        out[t] = f"{pc}:{pm}:{rc}:{rm}:{conf_map.get(t, 0)}"
    return out


def _heads_needing_retrain(
    session: Session, eligible: list[int], embedding_version: str,
    fps: dict[int, str], full: bool,
) -> list[int]:
    """The eligible tag_ids to (re)fit: no head yet, a head trained in a DIFFERENT
    embedding space (a model swap), or a changed training-data fingerprint.
    full=True forces every eligible tag. sklearn-free (train_head itself needs
    scikit-learn) so the incremental decision is unit-testable on its own."""
    if full:
        return list(eligible)
    existing = {
        tag_id: (fp, ev)
        for tag_id, fp, ev in session.execute(
            select(
                TagHead.tag_id, TagHead.train_fingerprint,
                TagHead.embedding_version,
            )
        ).all()
    }
    out = []
    for tag_id in eligible:
        prev = existing.get(tag_id)
        if (
            prev is None
            or prev[1] != embedding_version
            or prev[0] != fps.get(tag_id)
        ):
            out.append(tag_id)
    return out


def train_all_heads(
    session: Session, params: dict[str, Any], run: HeadTrainingRun | None = None
) -> dict[str, int]:
    """(Re)train eligible concept heads, INCREMENTALLY by default (#1317 p2):
    refit only the tags whose training data changed since last fit, so a manual
    Retrain click is fast. `params["full"]=True` (the nightly run) refits every
    head to reconcile sampled-negative + hygiene drift. Prunes heads whose tag is
    no longer eligible. Commits per head so a SIGKILL leaves trained heads durable.
    Returns {n_trained, n_skipped} (n_skipped = unchanged + too-few-examples)."""
    import numpy as np

    cfg = _normalize_params(session, params)
    embedding_version = _embedder_version(session)
    full = bool((params or {}).get("full"))
    eligible = _eligible_tag_ids(session, cfg["min_positives"])
    eligible_set = set(eligible)
    # Computed once per run, not per head — the hygiene set is identical for
    # every non-system concept.
    hygiene = _hygiene_excluded_ids(session)
    fps = _head_fingerprints(session, eligible)
    to_train = set(
        _heads_needing_retrain(session, eligible, embedding_version, fps, full)
    )
    trained = 0
    failed = 0
    for i, tag_id in enumerate(eligible):
        if tag_id not in to_train:
            continue
        try:
            ok = train_head(
                session, tag_id, embedding_version, cfg, np, hygiene=hygiene
            )
        except Exception:
            log.exception("train_head failed for tag %d", tag_id)
            ok = False
        if ok:
            # Stamp the fingerprint we trained against so an unchanged tag is
            # skipped on the next incremental run.
            head = session.get(TagHead, tag_id)
            if head is not None:
                head.train_fingerprint = fps.get(tag_id)
        session.commit()
        trained += int(ok)
        failed += int(not ok)
        if run is not None and i % 10 == 0:
            run.last_progress_at = datetime.now(UTC)
            session.commit()
    # Retire heads whose concept dropped out of the eligible set (lost its
    # positives, or the tag was re-kinded) so stale heads can't keep suggesting.
    if eligible_set:
        session.execute(delete(TagHead).where(TagHead.tag_id.not_in(eligible_set)))
    else:
        session.execute(delete(TagHead))
    session.commit()
    # n_skipped = unchanged (not attempted) + failed-to-fit (too few examples).
    return {
        "n_trained": trained,
        "n_skipped": (len(eligible) - len(to_train)) + failed,
    }


def head_training_ids(
    session: Session, tag_id: int, cfg: dict, hygiene: set[int] | None = None,
) -> tuple[list[int], list[int]] | None:
    """Select (pos_ids, neg_ids) for one head. Split out of train_head and
    kept sklearn-free so the hygiene exclusion is testable in the CI env
    (sklearn only exists in the ml image). Returns None when the concept has
    too few usable positives.

    Training hygiene (#128): images carrying a system tag are ABSENT from
    every other concept's training — dropped as positives AND kept out of
    the rejection/sampled negative pool (see _hygiene_excluded_ids). A system
    tag's own head trains on them unfiltered: its positives ARE the hygiene
    images."""
    tag = session.get(Tag, tag_id)
    if tag is not None and tag.is_system:
        hygiene = set()
    elif hygiene is None:
        hygiene = _hygiene_excluded_ids(session)

    pos_ids = [i for i in _ids_with_tag(session, tag_id) if i not in hygiene]
    if len(pos_ids) < cfg["min_positives"]:
        return None

    pos_set = set(pos_ids)
    rejected = [
        i for i in _rejected_ids(session, tag_id)
        if i not in pos_set and i not in hygiene
    ]
    want_neg = max(len(pos_ids) * cfg["neg_ratio"], _EXAMPLES_MIN * 4)
    sampled = _sample_unlabeled(
        session, pos_set | set(rejected) | hygiene,
        min(_UNLABELED_POOL, want_neg),
    )
    return pos_ids, rejected + [i for i in sampled if i not in pos_set]


def train_head(
    session: Session, tag_id: int, embedding_version: str, cfg: dict, np,
    hygiene: set[int] | None = None,
) -> bool:
    """Fit + upsert one head. Returns True if a head was written, False if the
    concept had too few usable examples to train (the row is then removed)."""
    from sklearn.linear_model import LogisticRegression
    from sklearn.model_selection import StratifiedKFold, cross_val_predict

    ids = head_training_ids(session, tag_id, cfg, hygiene)
    if ids is None:
        session.execute(delete(TagHead).where(TagHead.tag_id == tag_id))
        return False
    pos_ids, neg_ids = ids
    emb = _load_embeddings(session, pos_ids + neg_ids)
    pos = [emb[i] for i in pos_ids if i in emb]
    neg = [emb[i] for i in neg_ids if i in emb]
    if len(pos) < _EXAMPLES_MIN or len(neg) < _EXAMPLES_MIN:
        session.execute(delete(TagHead).where(TagHead.tag_id == tag_id))
        return False

    X = np.vstack(pos + neg).astype(np.float32)
    y = np.array([1] * len(pos) + [0] * len(neg))
    Xn = _l2norm(X, np)

    clf = LogisticRegression(max_iter=1000, class_weight="balanced")
    cv = StratifiedKFold(
        n_splits=_safe_folds(y, cfg["cv_folds"], np), shuffle=True, random_state=0
    )
    # Honest thresholds from out-of-fold scores; deployable weights from a final
    # fit on ALL the data.
    cv_probs = cross_val_predict(clf, Xn, y, cv=cv, method="predict_proba")[:, 1]
    metrics = _metrics_from_scores(y, cv_probs, np)
    auto = _auto_apply_point(y, cv_probs, cfg["precision_target"], np)
    clf.fit(Xn, y)

    head = session.get(TagHead, tag_id)
    if head is None:
        head = TagHead(tag_id=tag_id)
        session.add(head)
    head.embedding_version = embedding_version
    head.weights = clf.coef_[0].astype(np.float32).tolist()
    head.bias = float(clf.intercept_[0])
    head.suggest_threshold = float(metrics["threshold"])
    head.auto_apply_threshold = float(auto["threshold"]) if auto else None
    head.n_pos = len(pos)
    head.n_neg = len(neg)
    head.ap = float(metrics["ap"])
    head.precision_cv = float(metrics["precision"])
    head.recall = float(metrics["recall"])
    head.trained_at = datetime.now(UTC)
    head.metrics = {"f1": metrics["f1"], "auto_apply": auto}
    return True


# --- Scoring (async, API worker) -----------------------------------------
# Score one image against every current head to produce the rail's suggestions.
# A tiny in-process cache holds the stacked weight matrix keyed on (count,
# max(trained_at)) so a retrain invalidates it without per-request weight loads.
_HEADS_CACHE: dict[str, Any] = {"key": None, "heads": None}


async def _current_heads(session: AsyncSession, embedding_version: str):
    """Stacked (W, b, thresholds, tag_id/name/category) for heads matching the
    current embedding, cached until the next retrain."""
    import numpy as np

    sig = (
        await session.execute(
            select(func.count(), func.max(TagHead.trained_at)).where(
                TagHead.embedding_version == embedding_version
            )
        )
    ).one()
    key = f"{embedding_version}:{sig[0]}:{sig[1].isoformat() if sig[1] else '-'}"
    cached = _HEADS_CACHE.get("heads")
    if cached is not None and _HEADS_CACHE.get("key") == key:
        return cached
    rows = (
        await session.execute(
            select(
                TagHead.tag_id, Tag.name, Tag.kind, Tag.is_system,
                TagHead.weights, TagHead.bias,
                TagHead.suggest_threshold, TagHead.auto_apply_threshold,
            )
            .join(Tag, Tag.id == TagHead.tag_id)
            .where(TagHead.embedding_version == embedding_version)
        )
    ).all()
    if not rows:
        loaded = {"W": None, "rows": []}
    else:
        W = np.vstack([np.asarray(r.weights, dtype=np.float32) for r in rows])
        b = np.asarray([r.bias for r in rows], dtype=np.float32)
        thr = np.asarray([r.suggest_threshold for r in rows], dtype=np.float32)
        meta = [
            {
                "tag_id": r.tag_id,
                "name": r.name,
                # System tags (wip/banner/editor) are kind=general but group under
                # their OWN "system" suggestion category so the operator reviews
                # them apart from content tags (they still train as general heads).
                "category": "system" if r.is_system else _CATEGORY.get(r.kind, "general"),
                "auto_apply_threshold": r.auto_apply_threshold,
                "is_system": bool(r.is_system),
            }
            for r in rows
        ]
        loaded = {"W": W, "b": b, "thr": thr, "meta": meta}
    _HEADS_CACHE["key"] = key
    _HEADS_CACHE["heads"] = loaded
    return loaded


async def _image_bag(
    session: AsyncSession, image_id: int, cur_version: str,
) -> tuple[list, list[dict | None]]:
    """The max-over-bag inputs for one image: the whole-image SigLIP vector (when
    it's in the current model's space) PLUS every concept-region crop embedded in
    that space. Returns (bag, bag_meta) as PARALLEL lists — bag_meta[i] is None for
    the whole-image row, else the region's {bbox, kind, detector} so a surfaced tag
    can point back at the crop that produced it (#1206 grounding).

    Only current-version embeddings enter the bag: mid model-swap (#1190) an image
    still carrying an OLD-version whole-image vector is skipped rather than scored
    by heads trained in a different space; a legacy NULL version is treated as
    current (those predate per-row stamping). Shared by live scoring (score_image)
    and on-demand applied-tag grounding (ground_applied_tag, #1206 Step 4)."""
    import numpy as np

    img = await session.get(ImageRecord, image_id)
    bag: list = []
    bag_meta: list[dict | None] = []
    if img is None:
        return bag, bag_meta
    if img.siglip_embedding is not None and img.siglip_model_version in (
        cur_version, None,
    ):
        bag.append(np.asarray(img.siglip_embedding, dtype=np.float32))
        bag_meta.append(None)
    region_rows = (
        await session.execute(
            select(
                ImageRegion.siglip_embedding,
                ImageRegion.rx, ImageRegion.ry, ImageRegion.rw, ImageRegion.rh,
                ImageRegion.kind, ImageRegion.detector_version,
            )
            .where(ImageRegion.image_record_id == image_id)
            .where(ImageRegion.siglip_embedding.is_not(None))
            .where(ImageRegion.embedding_version == cur_version)
        )
    ).all()
    for vec, rx, ry, rw, rh, kind, detector in region_rows:
        if vec is not None:
            bag.append(np.asarray(vec, dtype=np.float32))
            bag_meta.append(
                {"bbox": [rx, ry, rw, rh], "kind": kind, "detector": detector}
            )
    return bag, bag_meta


async def score_image(
    session: AsyncSession, image_id: int, threshold_override: float | None = None,
) -> list[dict]:
    """Suggestions for one image from the trained heads: [{tag_id, name,
    category, score, above_threshold, grounding}], ranked. A concept is INCLUDED
    when its score clears the head's own suggest_threshold — or, when
    threshold_override is given (the typed-dropdown "show everything" mode), that
    flat floor instead (0 → every head). System-tag heads (wip/banner/editor)
    instead use a flat _SYSTEM_TAG_SUGGEST_FLOOR so their false positives surface
    for rejection (still overridden by threshold_override). Empty if the image has
    no embedding or no heads exist yet.

    ``above_threshold`` is reported SEPARATELY from inclusion: it's always whether
    the score cleared the head's NATURAL cut (suggest_threshold, or the system
    floor), regardless of any override. So the single min=0 fetch returns every
    head, and the caller can split panel (above_threshold) from dropdown (all)
    without a second request.

    MAX-OVER-BAG: the image is scored as a BAG of embeddings — the whole-image
    vector PLUS every concept-region crop the agent embedded (same model
    version) — and each head takes its MAX score across the bag. A small/local
    concept (glasses, a stomach bulge) that the whole-image vector washes out
    can still surface from the crop where it dominates. The whole-image vector is
    always in the bag, so this can never score lower than whole-image alone."""
    import numpy as np

    settings = await _settings_async(session)
    cur_version = settings.embedder_model_version
    heads = await _current_heads(session, cur_version)
    if heads["W"] is None:
        return []
    bag, bag_meta = await _image_bag(session, image_id, cur_version)
    if not bag:
        return []

    X = np.vstack(bag)                                  # (B, D)
    norms = np.linalg.norm(X, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    Xn = X / norms
    Z = Xn @ heads["W"].T + heads["b"]                  # (B, H)
    probs_bag = _sigmoid(Z, np)                         # (B, H)
    probs = probs_bag.max(axis=0)                       # (H,) best over the bag
    # ARGMAX beside the max: WHICH bag row won each head → the region that grounds
    # the tag (bag_meta[win]); None when the whole-image vector won (#1206).
    winners = probs_bag.argmax(axis=0)                  # (H,)
    out = []
    for i, p in enumerate(probs):
        m = heads["meta"][i]
        # The head's NATURAL suggest cut — system tags use the flat floor (see
        # _SYSTEM_TAG_SUGGEST_FLOOR) so their false positives show up for the
        # operator to reject; content heads use their own precision-tuned
        # threshold. This is what "above threshold" means (drives the panel).
        natural = (
            _SYSTEM_TAG_SUGGEST_FLOOR if m["is_system"] else float(heads["thr"][i])
        )
        # INCLUSION is looser under threshold_override (dropdown show-all,
        # override=0): every head comes back so a low-confidence concept can still
        # be typed + picked, each carrying its own above_threshold flag.
        cut = threshold_override if threshold_override is not None else natural
        if p >= cut:
            out.append({
                "tag_id": m["tag_id"],
                "name": m["name"],
                "category": m["category"],
                "score": float(p),
                "above_threshold": bool(p >= natural),
                "grounding": bag_meta[int(winners[i])],
            })
    out.sort(key=lambda d: d["score"], reverse=True)
    return out


async def ground_applied_tag(
    session: AsyncSession, image_id: int, tag_id: int,
) -> tuple[dict | None, bool]:
    """On-demand grounding for an ALREADY-APPLIED tag (#1206 Step 4). Applied tags
    aren't scored live, so recompute the max-over-bag argmax for just this tag's
    head — which crop region best explains the tag on this image — mirroring what
    score_image records for live suggestions. Returns (grounding, has_head):

    - has_head False → the tag has no head in the current embedding space (manual/
      artist/meta tags, or a concept below the head floor). Nothing to localize
      with, so the UI shows no overlay (distinct from "the whole image won").
    - grounding None (has_head True) → the whole-image vector best explains it,
      not any crop; the UI shows the subtle whole-image frame.
    - grounding {bbox, kind, detector} → the winning region.

    Character heads are covered too (character is a head kind); this deliberately
    reuses the SigLIP head bag rather than the CCIP figure path so every applied
    concept grounds through one consistent mechanism."""
    import numpy as np

    cur_version = (await _settings_async(session)).embedder_model_version
    row = (
        await session.execute(
            select(TagHead.weights, TagHead.bias).where(
                TagHead.tag_id == tag_id,
                TagHead.embedding_version == cur_version,
            )
        )
    ).one_or_none()
    if row is None:
        return None, False
    bag, bag_meta = await _image_bag(session, image_id, cur_version)
    if not bag:
        return None, True

    X = np.vstack(bag)
    norms = np.linalg.norm(X, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    Xn = X / norms
    # The sigmoid is monotonic in the logit, so the highest-probability bag row is
    # just argmax of the raw score — no need to exponentiate to pick the winner.
    z = Xn @ np.asarray(row.weights, dtype=np.float32) + float(row.bias)   # (B,)
    return bag_meta[int(z.argmax())], True


async def _settings_async(session: AsyncSession) -> MLSettings:
    return await MLSettings.load(session)


# --- Earned auto-apply (sync, ml worker) ---------------------------------
# A graduated head can apply its tag to images it scores above the head's
# auto_apply_threshold, without a human. Gated by a master switch + a support
# floor so a precise-looking but under-supported head can't spray tags.

_AUTO_APPLY_CHUNK = 5000


class HeadAutoApplyAlreadyRunning(Exception):
    """Raised when an auto-apply sweep is already in flight."""


class HeadAutoApplyDisabled(Exception):
    """Raised when a real (non-dry-run) sweep is requested but the master
    switch (head_auto_apply_enabled) is off."""


def start_head_auto_apply_run(session: Session, params: dict[str, Any]) -> int:
    """Create a HeadAutoApplyRun + dispatch the ml-queue sweep. dry_run previews
    (writes nothing); a real sweep needs the master switch on. One run at a time."""
    dry_run = bool((params or {}).get("dry_run", False))
    existing = session.execute(
        select(HeadAutoApplyRun.id).where(HeadAutoApplyRun.status == "running")
    ).scalar_one_or_none()
    if existing is not None:
        raise HeadAutoApplyAlreadyRunning(existing)
    if not dry_run and not _settings(session).head_auto_apply_enabled:
        raise HeadAutoApplyDisabled()
    run = HeadAutoApplyRun(
        dry_run=dry_run, params={"dry_run": dry_run}, status="running",
        last_progress_at=datetime.now(UTC),
    )
    session.add(run)
    session.flush()
    run_id = run.id
    from ...tasks.ml import apply_head_tags as _task
    _task.delay(run_id)
    return run_id


def _auto_apply_heads(session: Session, embedding_version: str, min_pos: int):
    """Eligible CONTENT heads to fire: graduated (auto_apply_threshold set),
    enough support, current embedding, NON-system. System tags never auto-apply
    via this path — `wip` never auto-applies at all, and banner/editor screenshot
    go through the presentation path at their own flat threshold (#141). Returns
    the row list (tag_id/name/weights/...)."""
    return session.execute(
        select(
            TagHead.tag_id, Tag.name, TagHead.weights, TagHead.bias,
            TagHead.auto_apply_threshold,
        )
        .join(Tag, Tag.id == TagHead.tag_id)
        .where(TagHead.embedding_version == embedding_version)
        .where(TagHead.auto_apply_threshold.is_not(None))
        .where(TagHead.n_pos >= min_pos)
        .where(~Tag.is_system)
    ).all()


def auto_apply_sweep(
    session: Session, run: HeadAutoApplyRun, dry_run: bool
) -> dict[str, Any]:
    """Score every embedded image against the eligible heads and apply (or, for
    dry_run, just count) each head's tag where score >= its auto_apply_threshold
    and the tag isn't already applied or rejected on that image. Streams
    embeddings in chunks; commits per chunk on a real run. Returns
    {n_applied, concepts:[{tag_id,name,applied,scanned,threshold}]}."""
    import numpy as np

    settings = _settings(session)
    rows = _auto_apply_heads(
        session, settings.embedder_model_version,
        settings.head_auto_apply_min_positives,
    )
    if not rows:
        return {"n_applied": 0, "concepts": []}

    W = np.vstack([np.asarray(r.weights, dtype=np.float32) for r in rows])
    b = np.asarray([r.bias for r in rows], dtype=np.float32)
    thr = np.asarray([r.auto_apply_threshold for r in rows], dtype=np.float32)
    tag_ids = [r.tag_id for r in rows]
    names = [r.name for r in rows]

    # Skip images that already carry, or have rejected, each tag.
    skip = _applied_or_rejected(session, tag_ids)

    applied = [0] * len(rows)
    scanned = 0
    all_ids = list(session.execute(
        select(ImageRecord.id).where(ImageRecord.siglip_embedding.is_not(None))
    ).scalars())
    for start in range(0, len(all_ids), _AUTO_APPLY_CHUNK):
        chunk = all_ids[start:start + _AUTO_APPLY_CHUNK]
        emb = _load_embeddings(session, chunk)
        cids = [i for i in chunk if i in emb]
        if not cids:
            continue
        Xn = _l2norm(np.vstack([emb[i] for i in cids]).astype(np.float32), np)
        probs = _sigmoid(Xn @ W.T + b, np)              # (N, H)
        scanned += len(cids)
        for h in range(len(rows)):
            tid = tag_ids[h]
            for idx in np.where(probs[:, h] >= thr[h])[0]:
                iid = cids[int(idx)]
                if iid in skip[tid]:
                    continue
                skip[tid].add(iid)
                applied[h] += 1
                if not dry_run:
                    session.execute(
                        pg_insert(image_tag)
                        .values(image_record_id=iid, tag_id=tid, source="head_auto")
                        .on_conflict_do_nothing()
                    )
        if not dry_run:
            session.commit()
        run.last_progress_at = datetime.now(UTC)
        session.commit()

    concepts = [
        {"tag_id": tag_ids[h], "name": names[h], "applied": applied[h],
         "scanned": scanned, "threshold": float(thr[h])}
        for h in range(len(rows))
    ]
    return {"n_applied": sum(applied), "concepts": concepts}


_PRESENTATION_SOURCE = "presentation_auto"
_PROCESS_SOURCE = "process_auto"

# System-tag auto-apply modes (#1464). Both modes run the identical sweep — apply
# a system tag at a flat threshold with a PROVISIONAL source + a ring-loud review
# guard — and differ ONLY in which tags, which settings knobs, and which
# source/review-mode. 'chrome' (banner) is HIDDEN from the gallery; 'process'
# (wip / editor screenshot) stays VISIBLE (the hide is a gallery-query effect of
# the tag's group membership, not of this sweep).
_SWEEP_MODES = {
    "chrome": {
        "names": CHROME_SYSTEM_TAGS,
        "enabled": "presentation_auto_apply_enabled",
        "threshold": "presentation_auto_apply_threshold",
        "conflict": "presentation_conflict_threshold",
        "source": _PRESENTATION_SOURCE,
    },
    "process": {
        "names": PROCESS_SYSTEM_TAGS,
        "enabled": "process_auto_apply_enabled",
        "threshold": "process_auto_apply_threshold",
        "conflict": "process_conflict_threshold",
        "source": _PROCESS_SOURCE,
    },
}


def _system_tag_heads(session: Session, embedding_version: str, names):
    """Trained heads for a system-tag group (chrome banner / process wip+editor).
    They fire at the group's FLAT threshold regardless of graduation — a head
    exists once the operator has labelled enough (head_min_positives)."""
    return session.execute(
        select(TagHead.tag_id, Tag.name, TagHead.weights, TagHead.bias)
        .join(Tag, Tag.id == TagHead.tag_id)
        .where(TagHead.embedding_version == embedding_version)
        .where(Tag.is_system.is_(True))
        .where(Tag.name.in_(names))
    ).all()


def _conflict_heads(session: Session, embedding_version: str):
    """ALL content (non-system) heads — the "does this ALSO look like real
    content" signal for the presentation conflict guard (#141)."""
    return session.execute(
        select(TagHead.tag_id, TagHead.weights, TagHead.bias)
        .join(Tag, Tag.id == TagHead.tag_id)
        .where(TagHead.embedding_version == embedding_version)
        .where(~Tag.is_system)
    ).all()


def _valued_image_ids(session: Session) -> set[int]:
    """Images the operator has shown they value: carrying a HUMAN or CONFIRMED
    content (non-system) tag. The presentation sweep never auto-hides these
    (guard 1) — you tagged it, so the model doesn't get to bury it (#141)."""
    confirmed = exists().where(
        TagPositiveConfirmation.image_record_id == image_tag.c.image_record_id,
        TagPositiveConfirmation.tag_id == image_tag.c.tag_id,
    )
    rows = session.execute(
        select(image_tag.c.image_record_id)
        .join(Tag, Tag.id == image_tag.c.tag_id)
        .where(~Tag.is_system)
        .where(image_tag.c.source.not_in(_AUTO_SOURCES) | confirmed)
    ).all()
    return {r[0] for r in rows}


def system_tag_auto_apply_sweep(
    session: Session, *, mode: str, dry_run: bool = False
) -> dict:
    """Auto-apply a system-tag group at its FLAT threshold. mode='chrome' (banner,
    #141) hides the image; mode='process' (wip / editor screenshot, #1464) keeps it
    VISIBLE — the ONLY difference is the tag group's gallery membership, not this
    sweep. Two guards keep it safe: (1) never touch an image carrying a
    human/confirmed content tag; (2) if the image ALSO scores >= the conflict
    threshold on a content head, still apply but flag it (PresentationReview,
    mode=<mode>) so the review strip surfaces "also looks like <X>". The source is
    PROVISIONAL so the head never trains on its own output. No-op unless the mode's
    enabled flag is set. numpy-only (no sklearn). Returns {n_applied, n_flagged,
    concepts}."""
    import numpy as np

    cfg = _SWEEP_MODES[mode]
    settings = _settings(session)
    if not dry_run and not getattr(settings, cfg["enabled"]):
        return {"n_applied": 0, "n_flagged": 0, "concepts": []}
    ver = settings.embedder_model_version
    pres = _system_tag_heads(session, ver, cfg["names"])
    if not pres:
        return {"n_applied": 0, "n_flagged": 0, "concepts": []}
    thr = float(getattr(settings, cfg["threshold"]))
    conflict_thr = float(getattr(settings, cfg["conflict"]))
    source = cfg["source"]

    Wp = np.vstack([np.asarray(r.weights, dtype=np.float32) for r in pres])
    bp = np.asarray([r.bias for r in pres], dtype=np.float32)
    pres_tag_ids = [r.tag_id for r in pres]
    pres_names = [r.name for r in pres]

    conf = _conflict_heads(session, ver)
    Wc = bc = conf_tag_ids = None
    if conf:
        Wc = np.vstack([np.asarray(r.weights, dtype=np.float32) for r in conf])
        bc = np.asarray([r.bias for r in conf], dtype=np.float32)
        conf_tag_ids = [r.tag_id for r in conf]

    valued = _valued_image_ids(session)

    # Skip images that already carry, or have rejected, each presentation tag.
    skip = _applied_or_rejected(session, pres_tag_ids)

    applied = [0] * len(pres)
    n_flagged = 0
    scanned = 0
    all_ids = list(session.execute(
        select(ImageRecord.id).where(ImageRecord.siglip_embedding.is_not(None))
    ).scalars())
    for start in range(0, len(all_ids), _AUTO_APPLY_CHUNK):
        chunk = all_ids[start:start + _AUTO_APPLY_CHUNK]
        emb = _load_embeddings(session, chunk)
        cids = [i for i in chunk if i in emb]
        if not cids:
            continue
        Xn = _l2norm(np.vstack([emb[i] for i in cids]).astype(np.float32), np)
        probs = _sigmoid(Xn @ Wp.T + bp, np)                  # (N, P)
        if Wc is not None:
            max_c, arg_c = _conflict_scores(Xn, Wc, bc, np)   # (N,), (N,)
        scanned += len(cids)
        for p in range(len(pres)):
            tid = pres_tag_ids[p]
            for idx in np.where(probs[:, p] >= thr)[0]:
                iid = cids[int(idx)]
                if iid in skip[tid] or iid in valued:
                    continue
                skip[tid].add(iid)
                applied[p] += 1
                if not dry_run:
                    session.execute(
                        pg_insert(image_tag)
                        .values(
                            image_record_id=iid, tag_id=tid,
                            source=source,
                        )
                        .on_conflict_do_nothing()
                    )
                # Guard 2: also looks like real content → still apply, but flag it
                # for the review strip instead of silently marking (chrome hides,
                # process stays visible — either way the operator gets a heads-up).
                if Wc is not None and float(max_c[idx]) >= conflict_thr:
                    n_flagged += 1
                    if not dry_run:
                        _insert_presentation_review(
                            session,
                            image_record_id=iid, tag_id=tid,
                            conflict_tag_id=conf_tag_ids[int(arg_c[idx])],
                            conflict_score=float(max_c[idx]),
                            mode=mode,
                        )
        if not dry_run:
            session.commit()

    concepts = [
        {"tag_id": pres_tag_ids[p], "name": pres_names[p],
         "applied": applied[p], "scanned": scanned, "threshold": thr}
        for p in range(len(pres))
    ]
    return {
        "n_applied": sum(applied), "n_flagged": n_flagged, "concepts": concepts,
    }


def soft_wip_conflict_audit(session: Session, dry_run: bool = False) -> dict:
    """Ring-loud audit for the SOFT WIP-title cohort (#1474). Images auto-tagged
    `wip` from a low-precision sketch/doodle title (source='wip_title_soft') that ALSO
    score >= the process conflict threshold on a content head are probably FINISHED
    art mis-tagged as process — flag them (PresentationReview, mode='process') so the
    review strip surfaces them ("also looks like <X>", Keep tag / Remove tag). Does
    NOT remove the tag; the operator decides. No-op when there are no content heads.
    numpy-only. Returns {n_scanned, n_flagged}."""
    import numpy as np

    from ..wip_title import WIP_TITLE_SOFT_SOURCE, resolve_wip_tag_id

    settings = _settings(session)
    ver = settings.embedder_model_version
    conflict_thr = float(settings.process_conflict_threshold)
    conf = _conflict_heads(session, ver)
    wip_id = resolve_wip_tag_id(session)
    if not conf or wip_id is None:
        return {"n_scanned": 0, "n_flagged": 0}
    Wc = np.vstack([np.asarray(r.weights, dtype=np.float32) for r in conf])
    bc = np.asarray([r.bias for r in conf], dtype=np.float32)
    conf_tag_ids = [r.tag_id for r in conf]

    soft_ids = [iid for (iid,) in session.execute(
        select(image_tag.c.image_record_id)
        .where(image_tag.c.tag_id == wip_id)
        .where(image_tag.c.source == WIP_TITLE_SOFT_SOURCE)
    )]
    # Skip images already flagged for this tag (idempotent re-runs).
    flagged = {iid for (iid,) in session.execute(
        select(PresentationReview.image_record_id)
        .where(PresentationReview.tag_id == wip_id)
    )}
    soft_ids = [i for i in soft_ids if i not in flagged]

    n_flagged = 0
    scanned = 0
    for start in range(0, len(soft_ids), _AUTO_APPLY_CHUNK):
        chunk = soft_ids[start:start + _AUTO_APPLY_CHUNK]
        emb = _load_embeddings(session, chunk)
        cids = [i for i in chunk if i in emb]
        if not cids:
            continue
        scanned += len(cids)
        Xn = _l2norm(np.vstack([emb[i] for i in cids]).astype(np.float32), np)
        max_c, arg_c = _conflict_scores(Xn, Wc, bc, np)
        for k in range(len(cids)):
            if float(max_c[k]) >= conflict_thr:
                n_flagged += 1
                if not dry_run:
                    _insert_presentation_review(
                        session,
                        image_record_id=cids[k], tag_id=wip_id,
                        conflict_tag_id=conf_tag_ids[int(arg_c[k])],
                        conflict_score=float(max_c[k]),
                        mode="process",
                    )
        if not dry_run:
            session.commit()
    return {"n_scanned": scanned, "n_flagged": n_flagged}


def retract_auto_applied_heads(session: Session) -> int:
    """Soft auto-apply (milestone 139): re-score every standing source='head_auto'
    tag against its CURRENT head and REMOVE the ones now BELOW the head's
    auto_apply_threshold — i.e. the head sharpened (or the operator raised the bar)
    and no longer supports them. Skips operator-confirmed tags
    (TagPositiveConfirmation). SILENT: a low score isn't proof the tag was wrong,
    so no hard negative is recorded — that's reserved for an operator removal.
    No-op unless head_auto_apply_enabled. Only re-scores the images that ALREADY
    carry the auto-tag (bounded), never the whole library. Returns n_retracted."""
    import numpy as np

    settings = _settings(session)
    if not settings.head_auto_apply_enabled:
        return 0
    heads = session.execute(
        select(
            TagHead.tag_id, TagHead.weights, TagHead.bias,
            TagHead.auto_apply_threshold,
        )
        .where(TagHead.embedding_version == settings.embedder_model_version)
        .where(TagHead.auto_apply_threshold.is_not(None))
    ).all()
    retracted = 0
    for tag_id, weights, bias, thr in heads:
        auto_ids = [
            iid for (iid,) in session.execute(
                select(image_tag.c.image_record_id)
                .where(image_tag.c.tag_id == tag_id)
                .where(image_tag.c.source == "head_auto")
            )
        ]
        if not auto_ids:
            continue
        confirmed = {
            iid for (iid,) in session.execute(
                select(TagPositiveConfirmation.image_record_id)
                .where(TagPositiveConfirmation.tag_id == tag_id)
                .where(TagPositiveConfirmation.image_record_id.in_(auto_ids))
            )
        }
        candidates = [i for i in auto_ids if i not in confirmed]
        emb = _load_embeddings(session, candidates)
        cids = [i for i in candidates if i in emb]
        if not cids:
            continue
        Xn = _l2norm(np.vstack([emb[i] for i in cids]).astype(np.float32), np)
        w = np.asarray(weights, dtype=np.float32)
        probs = _sigmoid(Xn @ w + float(bias), np)
        below = [cids[k] for k in np.where(probs < float(thr))[0]]
        for iid in below:
            session.execute(
                image_tag.delete()
                .where(image_tag.c.image_record_id == iid)
                .where(image_tag.c.tag_id == tag_id)
                .where(image_tag.c.source == "head_auto")
            )
            retracted += 1
        session.commit()
    return retracted