"""Label taxonomy and the coarse-grained collapses.

Kept out of `detector.py` so it imports without torch — same reason
`forensics.py` is its own module. `test_collapse.py` exercises it directly.

DroidDetect predicts four classes. Two coarser views are more useful to a human
moderator, and the Droid paper reports accuracy at each level (Large: ~97%
weighted F1 binary, ~88.8% ternary — the full 4-way split is the hardest,
because "adversarial" and "generated" are the same text with different styling).
"""

LABELS = [
    "HUMAN_GENERATED",
    "MACHINE_GENERATED",
    "MACHINE_REFINED",
    "MACHINE_GENERATED_ADVERSARIAL",
]

# "Who actually wrote this": the participant, a model, or both.
# MACHINE_GENERATED_ADVERSARIAL folds into MACHINE — model output styled to look
# hand-written is still model output.
TERNARY_OF = {
    "HUMAN_GENERATED": "HUMAN",
    "MACHINE_GENERATED": "MACHINE",
    "MACHINE_GENERATED_ADVERSARIAL": "MACHINE",
    "MACHINE_REFINED": "REFINED",
}
TERNARY_LABELS = ["HUMAN", "MACHINE", "REFINED"]

# The yes/no the old CodeT5+ Space returned.
BINARY_OF = {
    "HUMAN_GENERATED": "HUMAN",
    "MACHINE_GENERATED": "MACHINE",
    "MACHINE_GENERATED_ADVERSARIAL": "MACHINE",
    "MACHINE_REFINED": "MACHINE",
}
BINARY_LABELS = ["HUMAN", "MACHINE"]

TERNARY_MEANING = {
    "HUMAN": "Written by the participant",
    "MACHINE": "Model-written (including output styled to look hand-typed)",
    "REFINED": "Participant started it, a model completed or rewrote it",
}


def collapse(probs, mapping, order):
    """Sum a 4-class distribution into a coarser one.

    Returns (coarse_probs, top_label, top_confidence).

    Summing is the correct collapse because the four classes are mutually
    exclusive outcomes of one softmax. It also fixes a real failure of the raw
    argmax: when a model's output is genuinely ambiguous between "generated"
    and "adversarial", each class can sit below a human plurality while their
    sum is decisively machine. e.g. 0.40 human vs 0.31 + 0.27 machine — the
    4-way argmax says HUMAN, the ternary correctly says MACHINE at 0.58.
    """
    out = {k: 0.0 for k in order}
    for label, p in probs.items():
        out[mapping[label]] += p
    out = {k: round(v, 6) for k, v in out.items()}
    top = max(out, key=out.get)
    return out, top, out[top]