"""Shared ML helpers extracted in the DRY pass (milestone #161). These pin the
single sources the auto-apply sweeps now trust, so a future edit can't silently
drift them: `_applied_or_rejected` is the skip-set used by auto_apply_sweep,
system_tag_auto_apply_sweep (heads.py) and scheduled_ccip_auto_apply (tasks/ml.py);
`_sigmoid` is the head score→prob transform used at every scoring site."""
import pytest

from backend.app.models import ImageRecord, Tag, TagKind, TagSuggestionRejection
from backend.app.models.tag import image_tag
from backend.app.services.ml.training_data import _applied_or_rejected


def test_sigmoid_matches_naive_form():
    import numpy as np

    from backend.app.services.ml.heads import _sigmoid

    z = np.array([-3.0, -0.5, 0.0, 1.5, 12.0], dtype=np.float32)
    assert np.allclose(_sigmoid(z, np), 1.0 / (1.0 + np.exp(-z)))
    assert float(_sigmoid(np.array([0.0]), np)[0]) == pytest.approx(0.5)


@pytest.mark.integration
def test_applied_or_rejected_unions_applied_any_source_and_rejected(db_sync):
    a = Tag(name="dry-helper-a", kind=TagKind.general)
    b = Tag(name="dry-helper-b", kind=TagKind.general)
    db_sync.add_all([a, b])
    db_sync.flush()

    imgs = []
    for i in range(5):
        img = ImageRecord(
            path=f"/images/dryhelp{i}.jpg", sha256=f"{i:064d}", size_bytes=1,
            mime="image/jpeg", width=1, height=1, origin="imported_filesystem",
            integrity_status="unknown", siglip_embedding=[0.0] * 1152,
        )
        db_sync.add(img)
        imgs.append(img)
    db_sync.flush()

    # tag a: applied manually (img0), applied by an AUTO source (img1), rejected (img2).
    db_sync.execute(image_tag.insert().values(
        image_record_id=imgs[0].id, tag_id=a.id, source="manual"))
    db_sync.execute(image_tag.insert().values(
        image_record_id=imgs[1].id, tag_id=a.id, source="head_auto"))
    db_sync.add(TagSuggestionRejection(image_record_id=imgs[2].id, tag_id=a.id))
    # tag b: applied to img3 only.
    db_sync.execute(image_tag.insert().values(
        image_record_id=imgs[3].id, tag_id=b.id, source="manual"))
    db_sync.flush()

    skip = _applied_or_rejected(db_sync, [a.id, b.id])

    # Applied-under-ANY-source (manual + head_auto) ∪ rejected, kept per-tag; the
    # untouched image (img4) appears under neither tag.
    assert skip[a.id] == {imgs[0].id, imgs[1].id, imgs[2].id}
    assert skip[b.id] == {imgs[3].id}
    assert imgs[4].id not in skip[a.id]
    assert imgs[4].id not in skip[b.id]


# --- the CCIP auto-apply sweep's scorer ---------------------------------------
#
# `scheduled_ccip_auto_apply` scored one image per matmul, over every image in
# the library, on every daily run — and on 2026-09-23 it hit its 1800s soft
# limit on the operator's instance. `char_maxima` does the same arithmetic in
# blocks. These pin THAT: same answer, whatever the blocking.


def _score_fixture(np):
    """Four images with 1-3 figures each, three characters with 2/5/1
    prototypes. Deliberately ragged — equal group sizes would let a wrong
    `reduceat` offset pass."""
    from backend.app.services.ml.training_data import _l2norm

    rng = np.random.default_rng(7)
    dim = 16
    q_by_image = [
        _l2norm(rng.standard_normal((n, dim)).astype(np.float32), np)
        for n in (1, 3, 2, 1)
    ]
    mats = [
        _l2norm(rng.standard_normal((k, dim)).astype(np.float32), np)
        for k in (2, 5, 1)
    ]
    allref = np.vstack(mats)
    seg = np.cumsum([0] + [len(m) for m in mats])[:-1]
    return q_by_image, allref, seg


def _naive(q_by_image, allref, seg, np):
    """The loop as it was written before batching, kept longhand. The point of
    comparing against this rather than against a stored array is that it is
    the OLD CODE — if the batched form ever diverges, this says so in the
    terms the change was justified in."""
    return np.vstack([
        np.maximum.reduceat((q @ allref.T).max(axis=0), seg) for q in q_by_image
    ])


def test_char_maxima_matches_the_per_image_loop():
    import numpy as np

    from backend.app.services.ml.ccip import char_maxima

    q_by_image, allref, seg = _score_fixture(np)
    got = char_maxima(q_by_image, allref, seg, np)

    assert got.shape == (len(q_by_image), len(seg))
    np.testing.assert_allclose(
        got, _naive(q_by_image, allref, seg, np), rtol=1e-6, atol=1e-6,
    )


def test_the_answer_does_not_depend_on_where_the_blocks_fall():
    """The one thing batching could get wrong. Rows are reduced over the
    PROTOTYPE axis inside a block and over the FIGURE axis afterwards, so a
    block boundary may fall in the middle of an image's figures — which is
    safe only because max does not care how it is grouped. `max_elems=1`
    forces a boundary between every single row."""
    import numpy as np

    from backend.app.services.ml.ccip import char_maxima

    q_by_image, allref, seg = _score_fixture(np)
    whole = char_maxima(q_by_image, allref, seg, np, max_elems=10_000_000)
    split = char_maxima(q_by_image, allref, seg, np, max_elems=1)

    np.testing.assert_allclose(whole, split, rtol=1e-6, atol=1e-6)


def test_one_character_and_one_figure_still_reduces():
    """The degenerate shape `reduceat` is easiest to get wrong: a single
    segment starting at 0, and a single row."""
    import numpy as np

    from backend.app.services.ml.ccip import char_maxima

    q = np.array([[1.0, 0.0]], dtype=np.float32)
    allref = np.array([[1.0, 0.0], [0.0, 1.0]], dtype=np.float32)
    got = char_maxima([q], allref, np.array([0]), np)

    assert got.shape == (1, 1)
    assert got[0][0] == pytest.approx(1.0)