"""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.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 PRESENTATION_SYSTEM_TAGS, image_tag
from .training_data import (
    _AUTO_SOURCES,
    _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

# 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


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 session.execute(
        select(MLSettings).where(MLSettings.id == 1)
    ).scalar_one()


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)), 0.5), 0.999)
    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}], ranked. A concept surfaces 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.

    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 = 1.0 / (1.0 + np.exp(-Z))                # (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):
        if threshold_override is not None:
            cut = threshold_override
        elif heads["meta"][i]["is_system"]:
            # System tags surface at the flat floor (see _SYSTEM_TAG_SUGGEST_FLOOR)
            # so their false positives show up for the operator to reject.
            cut = _SYSTEM_TAG_SUGGEST_FLOOR
        else:
            cut = heads["thr"][i]
        if p >= cut:
            m = heads["meta"][i]
            out.append({
                "tag_id": m["tag_id"],
                "name": m["name"],
                "category": m["category"],
                "score": float(p),
                "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 session.execute(select(MLSettings).where(MLSettings.id == 1))
    ).scalar_one()


# --- 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
    from sqlalchemy.dialects.postgresql import insert as pg_insert

    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 = {tid: set() for tid in tag_ids}
    for tid in tag_ids:
        for (iid,) in session.execute(
            select(image_tag.c.image_record_id).where(image_tag.c.tag_id == tid)
        ):
            skip[tid].add(iid)
        for (iid,) in session.execute(
            select(TagSuggestionRejection.image_record_id).where(
                TagSuggestionRejection.tag_id == tid
            )
        ):
            skip[tid].add(iid)

    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 = 1.0 / (1.0 + np.exp(-(Xn @ W.T + b)))   # (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"


def _presentation_heads(session: Session, embedding_version: str):
    """Trained heads for the presentation chrome tags (banner / editor screenshot).
    They fire at the FLAT presentation threshold regardless of graduation — a head
    exists once the operator has labelled enough chrome (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_(PRESENTATION_SYSTEM_TAGS))
    ).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 presentation_auto_apply_sweep(session: Session, dry_run: bool = False) -> dict:
    """Auto-hide presentation chrome (banner / editor screenshot) at the FLAT
    presentation threshold (#141) — NOT the per-head graduated threshold. Two
    guards keep it safe: (1) never hide an image carrying a human/confirmed content
    tag; (2) if an image about to be hidden ALSO scores >= the conflict threshold
    on a content head, still hide it but flag it (PresentationReview) so the Hidden
    view surfaces "also looks like <X>" for review. No-op unless
    presentation_auto_apply_enabled. numpy-only (no sklearn). Returns
    {n_applied, n_flagged, concepts}."""
    import numpy as np
    from sqlalchemy.dialects.postgresql import insert as pg_insert

    settings = _settings(session)
    if not dry_run and not settings.presentation_auto_apply_enabled:
        return {"n_applied": 0, "n_flagged": 0, "concepts": []}
    ver = settings.embedder_model_version
    pres = _presentation_heads(session, ver)
    if not pres:
        return {"n_applied": 0, "n_flagged": 0, "concepts": []}
    thr = float(settings.presentation_auto_apply_threshold)
    conflict_thr = float(settings.presentation_conflict_threshold)

    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 = {tid: set() for tid in pres_tag_ids}
    for tid in pres_tag_ids:
        for (iid,) in session.execute(
            select(image_tag.c.image_record_id).where(image_tag.c.tag_id == tid)
        ):
            skip[tid].add(iid)
        for (iid,) in session.execute(
            select(TagSuggestionRejection.image_record_id).where(
                TagSuggestionRejection.tag_id == tid
            )
        ):
            skip[tid].add(iid)

    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 = 1.0 / (1.0 + np.exp(-(Xn @ Wp.T + bp)))       # (N, P)
        if Wc is not None:
            cprobs = 1.0 / (1.0 + np.exp(-(Xn @ Wc.T + bc)))  # (N, C)
            max_c = cprobs.max(axis=1)
            arg_c = cprobs.argmax(axis=1)
        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=_PRESENTATION_SOURCE,
                        )
                        .on_conflict_do_nothing()
                    )
                # Guard 2: also looks like content → hide but flag for review.
                if Wc is not None and float(max_c[idx]) >= conflict_thr:
                    n_flagged += 1
                    if not dry_run:
                        session.execute(
                            pg_insert(PresentationReview)
                            .values(
                                image_record_id=iid, tag_id=tid,
                                conflict_tag_id=conf_tag_ids[int(arg_c[idx])],
                                conflict_score=float(max_c[idx]),
                            )
                            .on_conflict_do_nothing()
                        )
        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 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 = 1.0 / (1.0 + np.exp(-(Xn @ w + float(bias))))
        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