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

The eval (#1130, tag_eval.py) proved the spine; this is its 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. Reuses tag_eval's proven data
  loaders + metric helpers so production heads match the eval's 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, func, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Session

from ...models import (
    HeadAutoApplyRun,
    HeadTrainingRun,
    ImageRecord,
    ImageRegion,
    MLSettings,
    Tag,
    TagHead,
    TagKind,
    TagSuggestionRejection,
)
from ...models.tag import image_tag
from .tag_eval import (
    _auto_apply_point,
    _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"}


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 labelled images — the
    set that gets a head. Counts all sources; source-aware filtering (#1133) is
    a separate, optional refinement."""
    rows = session.execute(
        select(Tag.id)
        .join(image_tag, image_tag.c.tag_id == Tag.id)
        .where(Tag.kind.in_(_HEAD_KINDS))
        .group_by(Tag.id)
        .having(func.count(image_tag.c.image_record_id) >= min_pos)
    ).all()
    return [r[0] for r in rows]


def train_all_heads(
    session: Session, params: dict[str, Any], run: HeadTrainingRun | None = None
) -> dict[str, int]:
    """(Re)train a head for every eligible concept; prune heads whose tag is no
    longer eligible. Commits per head so a SIGKILL leaves trained heads durable
    (training is idempotent). Returns {n_trained, n_skipped}."""
    import numpy as np

    cfg = _normalize_params(session, params)
    embedding_version = _embedder_version(session)
    eligible = _eligible_tag_ids(session, cfg["min_positives"])
    eligible_set = set(eligible)
    trained = 0
    skipped = 0
    for i, tag_id in enumerate(eligible):
        try:
            ok = train_head(session, tag_id, embedding_version, cfg, np)
        except Exception:
            log.exception("train_head failed for tag %d", tag_id)
            ok = False
        session.commit()
        trained += int(ok)
        skipped += 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()
    return {"n_trained": trained, "n_skipped": skipped}


def train_head(
    session: Session, tag_id: int, embedding_version: str, cfg: dict, np
) -> 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

    pos_ids = _ids_with_tag(session, tag_id)
    if len(pos_ids) < cfg["min_positives"]:
        session.execute(delete(TagHead).where(TagHead.tag_id == tag_id))
        return False

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

    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,
                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,
                "category": _CATEGORY.get(r.kind, "general"),
                "auto_apply_threshold": r.auto_apply_threshold,
            }
            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 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). 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

    img = await session.get(ImageRecord, image_id)
    if img is None:
        return []
    settings = await _settings_async(session)
    cur_version = settings.embedder_model_version
    heads = await _current_heads(session, cur_version)
    if heads["W"] is None:
        return []

    # Only embeddings in the CURRENT model's space enter the bag. Mid model-swap
    # (#1190), an image still carrying the 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).
    bag = []
    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))
    region_vecs = (
        await session.execute(
            select(ImageRegion.siglip_embedding)
            .where(ImageRegion.image_record_id == image_id)
            .where(ImageRegion.siglip_embedding.is_not(None))
            .where(ImageRegion.embedding_version == cur_version)
        )
    ).all()
    for (vec,) in region_vecs:
        if vec is not None:
            bag.append(np.asarray(vec, dtype=np.float32))
    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 = (1.0 / (1.0 + np.exp(-Z))).max(axis=0)      # (H,) best over the bag
    out = []
    for i, p in enumerate(probs):
        cut = threshold_override if threshold_override is not None else 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),
            })
    out.sort(key=lambda d: d["score"], reverse=True)
    return out


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 heads to fire: graduated (auto_apply_threshold set), enough
    support, current embedding. 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)
    ).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}