FabledCurator/tests/test_ml_suggestions.py

"""Suggestion read-path (tagging-v2): suggestions come from trained HEADS, not
Camie predictions or centroids. Heads are inserted directly (training needs
scikit-learn, ml image only); scoring is numpy-only (available via pgvector)."""
import pytest
from sqlalchemy import select

from backend.app.models import ImageRecord, ImageRegion, MLSettings, TagHead, TagKind
from backend.app.models.tag import image_tag
from backend.app.services.ml.allowlist import AllowlistService
from backend.app.services.ml.heads import ground_applied_tag
from backend.app.services.ml.suggestions import SuggestionService
from backend.app.services.tag_service import TagService

pytestmark = pytest.mark.integration


def _emb(slot: int, val: float = 3.0) -> list[float]:
    """An embedding pointing along axis `slot` (so its L2-normalized form is the
    unit vector e_slot — a head with weights e_slot scores it sigmoid(1)≈0.73)."""
    v = [0.0] * 1152
    v[slot] = val
    return v


async def _img(db, sha: str, emb=None) -> ImageRecord:
    img = ImageRecord(
        path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg",
        width=1, height=1, origin="imported_filesystem",
        integrity_status="unknown", siglip_embedding=emb,
    )
    db.add(img)
    await db.flush()
    return img


async def _embver(db) -> str:
    s = (await db.execute(select(MLSettings).where(MLSettings.id == 1))).scalar_one()
    return s.embedder_model_version


async def _head(db, tag_id: int, slot: int, suggest_threshold: float = 0.5):
    weights = [0.0] * 1152
    weights[slot] = 1.0
    db.add(TagHead(
        tag_id=tag_id, embedding_version=await _embver(db),
        weights=weights, bias=0.0, suggest_threshold=suggest_threshold,
        auto_apply_threshold=None, n_pos=10, n_neg=30,
        ap=0.8, precision_cv=0.9, recall=0.6,
    ))


@pytest.mark.asyncio
async def test_head_suggestion_surfaces_for_matching_image(db):
    tag = await TagService(db).find_or_create("glasses", TagKind.general)
    img = await _img(db, "a" * 64, _emb(0))
    await _head(db, tag.id, slot=0)
    await db.commit()

    sl = await SuggestionService(db).for_image(img.id)
    general = sl.by_category["general"]
    assert len(general) == 1
    s = general[0]
    assert s.canonical_tag_id == tag.id
    assert s.source == "head"
    assert s.creates_new_tag is False
    assert s.via_alias is False and s.raw_name is None
    assert s.score > 0.5


@pytest.mark.asyncio
async def test_no_embedding_means_no_suggestions(db):
    img = await _img(db, "b" * 64, None)
    tag = await TagService(db).find_or_create("cat", TagKind.general)
    await _head(db, tag.id, slot=0)
    await db.commit()
    assert (await SuggestionService(db).for_image(img.id)).by_category == {}


@pytest.mark.asyncio
async def test_no_heads_means_no_suggestions(db):
    img = await _img(db, "c" * 64, _emb(0))
    await db.commit()  # no heads trained yet
    assert (await SuggestionService(db).for_image(img.id)).by_category == {}


@pytest.mark.asyncio
async def test_applied_tag_not_suggested(db):
    tag = await TagService(db).find_or_create("dog", TagKind.general)
    img = await _img(db, "d" * 64, _emb(0))
    await _head(db, tag.id, slot=0)
    await db.execute(
        image_tag.insert().values(
            image_record_id=img.id, tag_id=tag.id, source="manual"
        )
    )
    await db.commit()
    sl = await SuggestionService(db).for_image(img.id)
    assert "general" not in sl.by_category or not sl.by_category["general"]


@pytest.mark.asyncio
async def test_threshold_override_surfaces_below_cut(db):
    # A head with a high suggest_threshold won't surface on a so-so score, but
    # the dropdown's override=0 floor surfaces every head regardless.
    tag = await TagService(db).find_or_create("horse", TagKind.general)
    img = await _img(db, "e" * 64, _emb(1))   # orthogonal to the head → score 0.5
    await _head(db, tag.id, slot=0, suggest_threshold=0.6)
    await db.commit()
    svc = SuggestionService(db)
    assert svc and not (await svc.for_image(img.id)).by_category.get("general")
    flooded = await svc.for_image(img.id, threshold_override=0.0)
    assert any(s.canonical_tag_id == tag.id for s in flooded.by_category["general"])


@pytest.mark.asyncio
async def test_system_tag_surfaces_at_flat_floor_despite_high_threshold(db):
    # A system tag's head uses the flat _SYSTEM_TAG_SUGGEST_FLOOR (0.65), not its
    # precision-tuned suggest_threshold — so a matching image (~0.73) surfaces it
    # for rejection even though its 0.9 auto threshold would hide it.
    from backend.app.services.ml.heads import _SYSTEM_TAG_SUGGEST_FLOOR
    assert _SYSTEM_TAG_SUGGEST_FLOOR == 0.65
    tag = await TagService(db).find_or_create("wip sketchy", TagKind.general)
    tag.is_system = True
    img = await _img(db, "1" * 64, _emb(0))            # matches head → ~0.73
    await _head(db, tag.id, slot=0, suggest_threshold=0.9)
    await db.commit()
    sl = await SuggestionService(db).for_image(img.id)
    # System tags surface under their OWN "system" category, not "general".
    assert any(
        s.canonical_tag_id == tag.id for s in sl.by_category.get("system", [])
    )
    assert not sl.by_category.get("general")


@pytest.mark.asyncio
async def test_system_tag_below_floor_stays_hidden(db):
    # Below the 0.65 floor (~0.5 on an orthogonal image) the system tag is not
    # surfaced — the floor still gates, it isn't "always show".
    tag = await TagService(db).find_or_create("wip faint", TagKind.general)
    tag.is_system = True
    img = await _img(db, "2" * 64, _emb(1))            # orthogonal → 0.5
    await _head(db, tag.id, slot=0, suggest_threshold=0.9)
    await db.commit()
    sl = await SuggestionService(db).for_image(img.id)
    assert not sl.by_category.get("system")


@pytest.mark.asyncio
async def test_content_tag_ignores_system_floor(db):
    # A NON-system head with the same high threshold does NOT get the floor: the
    # ~0.73 match stays hidden below its 0.9 auto threshold (floor is system-only).
    tag = await TagService(db).find_or_create("glasses hi", TagKind.general)
    img = await _img(db, "3" * 64, _emb(0))            # matches head → ~0.73
    await _head(db, tag.id, slot=0, suggest_threshold=0.9)
    await db.commit()
    sl = await SuggestionService(db).for_image(img.id)
    assert not sl.by_category.get("general")


@pytest.mark.asyncio
async def test_concept_region_surfaces_via_max_over_bag(db):
    # Max-over-bag: the whole-image vector is orthogonal to the head (scores the
    # 0.5 midpoint, under a 0.7 cut → nothing), but a concept CROP that aligns
    # with the head lifts the max over the bag above the cut. A small/local
    # concept surfaces ONLY because of the crop.
    tag = await TagService(db).find_or_create("glasses", TagKind.general)
    img = await _img(db, "b1" * 32, _emb(5))      # whole-image ⟂ head
    await _head(db, tag.id, slot=0, suggest_threshold=0.7)
    await db.commit()
    # Whole-image alone: sigmoid(0)=0.5 < 0.7 → no suggestion.
    assert not (await SuggestionService(db).for_image(img.id)).by_category.get("general")

    # A concept crop aligned with the head, but stamped with a STALE model
    # version → filtered out of the bag, so still nothing.
    db.add(ImageRegion(
        image_record_id=img.id, kind="concept",
        rx=0.1, ry=0.1, rw=0.3, rh=0.3,
        siglip_embedding=_emb(0), embedding_version="stale-embedder-v0",
    ))
    await db.commit()
    assert not (await SuggestionService(db).for_image(img.id)).by_category.get("general")

    # A matching-version concept crop → max-over-bag lifts it over the cut.
    db.add(ImageRegion(
        image_record_id=img.id, kind="concept",
        rx=0.4, ry=0.4, rw=0.3, rh=0.3,
        siglip_embedding=_emb(0), embedding_version=await _embver(db),
    ))
    await db.commit()
    general = (await SuggestionService(db).for_image(img.id)).by_category["general"]
    s = next(x for x in general if x.canonical_tag_id == tag.id)
    assert s.score > 0.7
    # #1206: the winning crop grounds the tag — hover highlights THIS region
    # (the matching-version crop at 0.4,0.4,0.3,0.3), not the whole image.
    assert s.grounding is not None
    assert s.grounding["bbox"] == pytest.approx([0.4, 0.4, 0.3, 0.3])
    assert s.grounding["kind"] == "concept"


@pytest.mark.asyncio
async def test_grounding_none_when_whole_image_wins(db):
    # #1206: when the whole-image vector (not a crop) produces the winning score,
    # grounding is None — the tag came from the global vector, not a region.
    tag = await TagService(db).find_or_create("sky", TagKind.general)
    img = await _img(db, "d1" * 32, _emb(0))          # whole-image ALIGNED w/ head
    await _head(db, tag.id, slot=0, suggest_threshold=0.5)
    db.add(ImageRegion(                                # an orthogonal crop (0.5)
        image_record_id=img.id, kind="concept",
        rx=0.1, ry=0.1, rw=0.2, rh=0.2,
        siglip_embedding=_emb(5), embedding_version=await _embver(db),
    ))
    await db.commit()
    general = (await SuggestionService(db).for_image(img.id)).by_category["general"]
    s = next(x for x in general if x.canonical_tag_id == tag.id)
    assert s.grounding is None


@pytest.mark.asyncio
async def test_stale_embedding_version_excluded_from_scoring(db):
    # Mid model-swap (#1190): an image still carrying an OLD-version whole-image
    # embedding must NOT be scored by heads trained in the new model's space —
    # even though the vector aligns with the head, it's the wrong coordinate
    # system, so nothing surfaces until it's re-embedded.
    tag = await TagService(db).find_or_create("glasses", TagKind.general)
    img = await _img(db, "c1" * 32, _emb(0))
    img.siglip_model_version = "some-old-model-v0"      # != current embedder
    await _head(db, tag.id, slot=0, suggest_threshold=0.5)
    await db.commit()
    assert not (await SuggestionService(db).for_image(img.id)).by_category.get("general")


@pytest.mark.asyncio
async def test_rejected_tag_surfaced_flagged_then_reversible(db):
    # A dismissed suggestion is NOT dropped: it stays flagged rejected so the
    # rail can show it + offer one-click un-reject (operator-asked 2026-06-27).
    tag = await TagService(db).find_or_create("goblin", TagKind.general)
    img = await _img(db, "f" * 64, _emb(0))
    await _head(db, tag.id, slot=0)
    await db.commit()
    await AllowlistService(db).dismiss(img.id, tag.id)
    await db.commit()

    sl = await SuggestionService(db).for_image(img.id)
    s = next(x for x in sl.by_category["general"] if x.canonical_tag_id == tag.id)
    assert s.rejected is True

    await AllowlistService(db).undismiss(img.id, tag.id)
    await db.commit()
    sl2 = await SuggestionService(db).for_image(img.id)
    s2 = next(x for x in sl2.by_category["general"] if x.canonical_tag_id == tag.id)
    assert s2.rejected is False


async def _figure(db, image_id, slot):
    v = [0.0] * 768
    v[slot] = 1.0
    db.add(ImageRegion(
        image_record_id=image_id, kind="figure",
        rx=0.0, ry=0.0, rw=1.0, rh=1.0,
        ccip_embedding=v, embedding_version="ccip-test",
    ))


@pytest.mark.asyncio
async def test_ccip_character_surfaces_in_rail(db):
    # A character with a CCIP reference (a tagged figure) is suggested on a new
    # image whose figure matches — overlaid into the rail alongside the heads.
    raven = await TagService(db).find_or_create("Raven", TagKind.character)
    ref = await _img(db, "0" * 64, None)         # the operator's tagged example
    await _figure(db, ref.id, slot=0)
    await db.execute(image_tag.insert().values(
        image_record_id=ref.id, tag_id=raven.id, source="manual",
    ))
    query = await _img(db, "1" * 64, None)        # untagged, matching figure
    await _figure(db, query.id, slot=0)
    await db.commit()

    sl = await SuggestionService(db).for_image(query.id)
    m = next(
        c for c in sl.by_category.get("character", [])
        if c.canonical_tag_id == raven.id
    )
    assert m.source == "ccip"


# --- #1206 Step 4: on-demand grounding for ALREADY-APPLIED tag chips ---------
# Applied tags aren't scored live, so ground_applied_tag recomputes the winning
# bag row for one tag's head on demand — the data behind hovering an applied chip.


@pytest.mark.asyncio
async def test_ground_applied_tag_returns_winning_region(db):
    # Whole-image vector is orthogonal to the head; a concept crop aligns with it,
    # so the crop is the max-over-bag winner → grounding is THAT region's box.
    tag = await TagService(db).find_or_create("glasses", TagKind.general)
    img = await _img(db, "g1" * 32, _emb(5))            # whole-image ⟂ head
    await _head(db, tag.id, slot=0)
    db.add(ImageRegion(
        image_record_id=img.id, kind="concept",
        rx=0.4, ry=0.4, rw=0.3, rh=0.3,
        siglip_embedding=_emb(0), embedding_version=await _embver(db),
    ))
    await db.commit()

    grounding, has_head = await ground_applied_tag(db, img.id, tag.id)
    assert has_head is True
    assert grounding is not None
    assert grounding["bbox"] == pytest.approx([0.4, 0.4, 0.3, 0.3])
    assert grounding["kind"] == "concept"


@pytest.mark.asyncio
async def test_ground_applied_tag_none_when_whole_image_wins(db):
    # Whole-image vector aligns with the head and out-scores an orthogonal crop →
    # grounding None (has_head True): the tag is best explained by the global
    # vector, so the chip hover shows the subtle whole-image frame, not a box.
    tag = await TagService(db).find_or_create("sky", TagKind.general)
    img = await _img(db, "g2" * 32, _emb(0))            # whole-image ALIGNED
    await _head(db, tag.id, slot=0)
    db.add(ImageRegion(
        image_record_id=img.id, kind="concept",
        rx=0.1, ry=0.1, rw=0.2, rh=0.2,
        siglip_embedding=_emb(5), embedding_version=await _embver(db),
    ))
    await db.commit()

    grounding, has_head = await ground_applied_tag(db, img.id, tag.id)
    assert has_head is True
    assert grounding is None


@pytest.mark.asyncio
async def test_ground_applied_tag_no_head(db):
    # A tag with no head in the current space (manual/artist/meta, or below the
    # head floor) can't be localized → (None, has_head False); the UI draws no
    # overlay at all, distinct from "the whole image won".
    tag = await TagService(db).find_or_create("artist:someone", TagKind.general)
    img = await _img(db, "g3" * 32, _emb(0))
    await db.commit()

    grounding, has_head = await ground_applied_tag(db, img.id, tag.id)
    assert has_head is False
    assert grounding is None