"""Self-check for the ternary/binary collapse. No torch, no model, instant.

    python test_collapse.py

This is the layer moderators actually read, so a bug here is a bug in every
verdict the Space reports.
"""

from labels import (
    BINARY_LABELS,
    BINARY_OF,
    LABELS,
    TERNARY_LABELS,
    TERNARY_MEANING,
    TERNARY_OF,
    collapse,
)


def check(probs, want_ternary, want_binary, note):
    assert abs(sum(probs.values()) - 1) < 1e-6, f"{note}: fixture isn't a distribution"
    t, tl, tc = collapse(probs, TERNARY_OF, TERNARY_LABELS)
    b, bl, bc = collapse(probs, BINARY_OF, BINARY_LABELS)
    assert abs(sum(t.values()) - 1) < 1e-6, f"{note}: ternary not normalised: {t}"
    assert abs(sum(b.values()) - 1) < 1e-6, f"{note}: binary not normalised: {b}"
    assert tl == want_ternary, f"{note}: ternary {tl} != {want_ternary}"
    assert bl == want_binary, f"{note}: binary {bl} != {want_binary}"
    print(f"  ok  {note:<46} ternary={tl:<8}({tc:.2f})  binary={bl}({bc:.2f})")


def main():
    # Every 4-way class must map into both coarser taxonomies, or a real
    # prediction would KeyError at request time.
    assert set(TERNARY_OF) == set(LABELS), "a class is missing a ternary mapping"
    assert set(BINARY_OF) == set(LABELS), "a class is missing a binary mapping"
    assert set(TERNARY_OF.values()) <= set(TERNARY_LABELS)
    assert set(BINARY_OF.values()) <= set(BINARY_LABELS)
    assert set(TERNARY_MEANING) == set(TERNARY_LABELS), "a ternary label has no blurb"

    check(
        {"HUMAN_GENERATED": 0.80, "MACHINE_GENERATED": 0.10,
         "MACHINE_REFINED": 0.06, "MACHINE_GENERATED_ADVERSARIAL": 0.04},
        "HUMAN", "HUMAN", "clearly human",
    )
    check(
        {"HUMAN_GENERATED": 0.18, "MACHINE_GENERATED": 0.09,
         "MACHINE_REFINED": 0.71, "MACHINE_GENERATED_ADVERSARIAL": 0.02},
        "REFINED", "MACHINE", "refined keeps its own ternary class",
    )
    # The case that justifies collapsing at all: machine mass split across two
    # classes loses the 4-way argmax to a human plurality, but summed it wins.
    check(
        {"HUMAN_GENERATED": 0.40, "MACHINE_GENERATED": 0.31,
         "MACHINE_REFINED": 0.02, "MACHINE_GENERATED_ADVERSARIAL": 0.27},
        "MACHINE", "MACHINE", "split machine mass beats a human plurality",
    )
    check(
        {"HUMAN_GENERATED": 0.05, "MACHINE_GENERATED": 0.10,
         "MACHINE_REFINED": 0.05, "MACHINE_GENERATED_ADVERSARIAL": 0.80},
        "MACHINE", "MACHINE", "adversarial folds into MACHINE",
    )
    # Refined must NOT absorb machine mass at the ternary level — the whole
    # point is telling "co-written" apart from "model-written".
    check(
        {"HUMAN_GENERATED": 0.10, "MACHINE_GENERATED": 0.55,
         "MACHINE_REFINED": 0.30, "MACHINE_GENERATED_ADVERSARIAL": 0.05},
        "MACHINE", "MACHINE", "refined stays separate from generated",
    )

    print("\ncollapse self-check OK")


if __name__ == "__main__":
    main()