app/fp/corpus2.py

"""Contract v2.1 synthetic corpus where *layout* is a locked experimental axis.

The pilot (v1) trained on three layouts and one per harder split, so a held-out
layout scored zero and the result could not be attributed. This corpus makes
layout families first-class, named and disjoint across splits: arm A trains on
the narrow set, arm B on the broad set, and both are scored on identical rows
whose layouts neither arm ever saw. That is the difference between measuring
memorization and measuring generalization over formats.

Three things make the gold here trustworthy rather than convenient:

*  Every record is derived by `fp.normalize.record`, the same contract code that
   scores predictions. The generator never writes an output value by hand, so a
   case the contract cannot represent fails here instead of becoming a silent
   disagreement between gold and prediction.
*  Season values come from span *text* through `normalize.season_value`. There is
   no value head and no annotator number, so "the model learned the numbers"
   cannot be confused with "the code parsed the numbers".
*  Every row is strict: each atom carries a kind, including DISCARD for junk, so
   unexplained debris is impossible by construction and `check_debris` proves it.

Layout parts are explicit so a typo cannot become silent text: a `str` is a
literal that must contain no word characters, `("a", key)` is an atom (a span),
and `("j", text)` is junk claimed as DISCARD.
"""

import argparse
import hashlib
import json
import random
import re
from collections import Counter
from pathlib import Path

from fp.normalize import record as contract_record

ROOT = Path("data/parser-v2")
SEED = 20260915
MAX_CHARS = 120  # the character encoders cap at 128 positions; nothing truncates
# Tolerated exactly as `fp.gold` tolerates them: a real outcome to be reproduced,
# not a generator bug. Anything else means the row cannot be represented.
TOLERATED = {"season-unspecified", "dimensions-no-class"}

# The span kinds this experiment's head may emit, and therefore the kinds the
# corpus must demonstrate. Kept here, not in the trainer, so a kind that exists
# in one place only is a test failure rather than an unmeasured blind spot.
KINDS = ("TITLE", "EPISODE_TITLE", "GROUP", "SEASON", "EPISODE", "RANGE", "LIST", "YEAR",
         "RESOLUTION", "DIMENSIONS", "DISCARD")

LANGS = ("en", "pt", "es", "ru", "ja", "zh")
WORDS = {
    "en": ["Quiet", "Harbor", "Copper", "Moon", "Garden", "River", "Paper", "Forest",
           "Station", "Lantern", "Winter", "Signal"],
    "pt": ["Porto", "Lua", "Jardim", "Rio", "Papel", "Bosque", "Pedra", "Silencio",
           "Estacao", "Farol", "Inverno", "Sinal"],
    "es": ["Puerto", "Luna", "Jardin", "Rio", "Papel", "Bosque", "Piedra", "Silencio",
           "Estacion", "Faro", "Invierno", "Senal"],
    "ru": ["Тихий", "Порт", "Луна", "Сад", "Река", "Бумага", "Лес", "Камень",
           "Вокзал", "Маяк", "Зима", "Сигнал"],
    "ja": ["静かな", "港", "月", "庭", "川", "紙", "森", "石",
           "駅", "灯台", "冬", "信号"],
    "zh": ["安静", "港湾", "月亮", "花园", "河流", "纸张", "森林", "石头",
           "车站", "灯塔", "冬天", "信号"],
}
# Authored second-script forms, not real transliterations: a controlled bilingual
# vocabulary for testing alias handling. The spans are the point, not the
# linguistics, and nobody should read these rows as language data.
ALIAS = {
    "ru": ["Tihiy", "Port", "Luna", "Sad", "Reka", "Bumaga", "Les", "Kamen",
           "Vokzal", "Mayak", "Zima", "Signal"],
    "ja": ["Shizuka", "Minato", "Tsuki", "Niwa", "Kawa", "Kami", "Mori", "Ishi",
           "Eki", "Toudai", "Fuyu", "Shingou"],
    "zh": ["Anjing", "Gangwan", "Yueliang", "Huayuan", "Heliu", "Zhiwang", "Senlin",
           "Shitou", "Chezhan", "Dengta", "Dongtian", "Xinhao"],
}
SUBTITLE = {
    "en": ["Rain Memory", "The Last Promise", "A Quiet Harbor", "Paper Lanterns",
           "North Wind", "Signals at Dusk"],
    "pt": ["Memoria da Chuva", "A Ultima Promessa", "Noite Serena", "Vento Norte",
           "Farol Apagado", "Cidade Velha"],
    "es": ["Recuerdo de Lluvia", "La Ultima Promesa", "Noche Serena", "Viento Norte",
           "Farol Apagado", "Ciudad Vieja"],
    "ru": ["Память дождя", "Последнее обещание", "Тихая ночь", "Северный ветер",
           "Погасший маяк", "Старый город"],
    "ja": ["雨の記憶", "最後の約束", "静かな夜", "北の風", "消えた灯台", "古い街"],
    "zh": ["雨的记忆", "最后的约定", "安静的夜", "北风", "消失的灯塔", "旧城"],
}
SEASON_MARKS = {
    "en": ["Season {n}", "{ord} Season", "Season {n}", "s{n}"],
    "pt": ["{ord} Temporada", "Temporada {n}", "temporada {n}"],
    "es": ["{ord} Temporada", "Temporada {n}", "temp {n}"],
    "ru": ["{ord} Сезон", "Сезон {n}", "сезон {n}"],
    "ja": ["第{n}期", "シーズン{n}", "{ord}シーズン"],
    "zh": ["第{ord}季", "第{n}季"],
}
ORDINALS = {
    "en": ["First", "Second", "Third", "Fourth"],
    "pt": ["Primeira", "Segunda", "Terceira", "Quarta"],
    "es": ["Primera", "Segunda", "Tercera", "Cuarta"],
    "ru": ["Первый", "Второй", "Третий", "Четвёртый"],
    "ja": ["一", "二", "三", "四"],
    "zh": ["一", "二", "三", "四"],
}
FINAL_MARK = {"en": "Final Season", "pt": "Temporada Final", "es": "Temporada Final",
              "ru": "Последний Сезон", "ja": "最終シーズン", "zh": "最终季"}
GROUPS = ["Nebula-raws", "Lantern Works", "Aoi-Subs", "Meridian", "Kite-Team",
          "Harbor-Encode", "Sable-raws"]
JUNK = ["REMASTERED", "BluRay", "x264", "MULTI-AUDIO", "DVDRIP", "Hi10P", "WEBRip",
        "AC3", "PROPER", "字幕组", "双语", "Batch", "Complete", "CD1", "Vol02"]
EXTENSIONS = [".mkv", ".mp4", ".avi", ".webm"]
RESOLUTIONS = ["480p", "720p", "1080p", "2160p", "720P", "1080P"]
DIMENSIONS = ["1280x720", "1920x1080", "640x480", "720x576", "3840x2160"]


def digest(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def layout_inventory():
    """(name, tag, parts). `narrow` mirrors the pilot's three shapes; `broad` is
    the format-coverage hypothesis; `heldout-*` trains nothing in this experiment.
    """
    return [
        ("n1", "narrow", ["[", ("a", "g"), "] ", ("a", "t"), " - ", ("a", "e"),
                          " [", ("a", "R"), "]"]),
        ("n2", "narrow", [("a", "t"), " [", ("a", "s"), "] ", ("a", "e"), " ",
                          ("a", "R"), "-", ("a", "d1")]),
        ("n3", "narrow", [("a", "t"), "_", ("a", "s"), "_", ("a", "e"), "_", ("a", "R"),
                          "_[", ("a", "d1"), "]"]),
        ("b1", "broad", [("a", "g"), " ", ("a", "t"), " - ", ("a", "e"), " [", ("a", "R"),
                         "][", ("a", "d1"), "]"]),
        ("b2", "broad", [("a", "t"), " (", ("a", "y"), ") - ", ("a", "e"), " [", ("a", "R"),
                         "] [", ("a", "d1"), "]"]),
        ("b3", "broad", [("a", "t"), ".", ("a", "y"), ".", ("a", "R"), ".", ("a", "d1"),
                         "-", ("a", "d2"), "-", ("a", "g")]),
        ("b4", "broad", [("a", "t"), " ", ("a", "x"), " [", ("a", "R"), "] [", ("a", "g"), "]"]),
        ("b5", "broad", [("a", "g"), " - ", ("a", "t"), " - ", ("a", "e"), " - ", ("a", "x")]),
        ("b6", "broad", [("a", "t"), " ", ("a", "a"), " [", ("a", "s"), "] [", ("a", "R"), "]"]),
        ("b7", "broad", [("a", "t"), " ", ("a", "a"), " ", ("a", "s"), " ", ("a", "e")]),
        ("b8", "broad", [("a", "t"), " ", ("a", "sc"), ("a", "ec"), " [", ("a", "R"), "]"]),
        ("b9", "broad", ["[", ("a", "d1"), "] ", ("a", "t"), " [", ("a", "x"), "] ",
                         ("a", "e"), " (", ("a", "R"), ")"]),
        ("b10", "broad", [("a", "t"), " - ", ("a", "e"), " ", ("a", "d1"), " ", ("a", "R"),
                          " [", ("a", "g"), "]"]),
        ("b11", "broad", ["[", ("a", "g"), "] ", ("a", "t"), "(", ("a", "y"), ") ",
                          ("a", "s"), " ", ("a", "l")]),
        ("b12", "broad", [("a", "t"), " ", ("a", "s"), " ", ("a", "r"), " ", ("a", "R"),
                          " ", ("a", "d1")]),
        ("b13", "broad", [("a", "t"), " [", ("a", "D"), "] ", ("a", "e"), " [", ("a", "d1"),
                          "] [", ("a", "g"), "]"]),
        ("b14", "broad", [("a", "t"), " ", ("a", "fm"), " ", ("a", "e"), " [", ("a", "R"), "]"]),
        ("b15", "broad", [("a", "g"), "/", ("a", "t"), "/", ("a", "td"), " ", ("a", "e"),
                          " [", ("a", "R"), "]"]),
        ("b16", "broad", [("a", "t"), " ", ("a", "s"), " ", ("a", "ej"), " [", ("a", "R"), "]"]),
        # held-out format families
        ("h1", "heldout-layout-a", [("j", "TV"), "/", ("a", "t"), "/", ("a", "s"), "/", ("a", "td"), " ",
                                    ("a", "d1"), " ", ("a", "e")]),
        ("h2", "heldout-layout-a", [("a", "t"), "(", ("a", "y"), ")", ("a", "sc"),
                                     "(", ("a", "R"), ")(", ("a", "d1"), ")"]),
        ("h3", "heldout-layout-a", [("a", "t"), " - ", ("a", "e3"), ("a", "vt"), " - ",
                                    ("a", "x"), " [", ("a", "R"), "]"]),
        ("h4", "heldout-layout-a", ["[", ("a", "g"), "] ", ("a", "t"), " ", ("a", "ej"),
                                    " [", ("a", "D"), "]"]),
        ("h5", "heldout-layout-b", ["【", ("a", "g"), "】", ("a", "t"), "【", ("a", "ej"),
                                    "】【", ("a", "R"), "】"]),
        ("h6", "heldout-layout-b", [("a", "t"), " | ", ("a", "a"), " | ", ("a", "e"), " | ",
                                    ("a", "d1")]),
        ("h7", "heldout-layout-b", ["(", ("a", "d1"), ") ", ("a", "t"), " ", ("a", "ej"),
                                    " ", ("a", "x"), " ", ("a", "R")]),
        ("h8", "heldout-layout-b", [("a", "t"), "_", ("a", "s"), "_", ("a", "e"), "-",
                                    ("a", "x"), "_", ("a", "d1"), "_[", ("a", "g"), "]"]),
        ("h9", "heldout-layout-b", [("j", "Series"), "/", ("a", "t"), "/", ("a", "epp"), "-", ("a", "td"),
                                    " ", ("a", "d1")]),
        ("h10", "heldout-layout-b", [("a", "t"), " ", ("a", "a"), " ", ("a", "tc"), " ",
                                      ("a", "l"), " [", ("a", "R"), "]"]),
        # held-out lexical markers in familiar shapes
        ("k1", "heldout-lexical", [("a", "g"), " ", ("a", "t"), " ", ("a", "epw"), " [",
                                   ("a", "R"), "]"]),
        ("k2", "heldout-lexical", [("a", "t"), " ", ("a", "chp"), ("j", " of 24"), " [", ("a", "R"), "]"]),
        ("k3", "heldout-lexical", [("a", "t"), " ", ("a", "sw"), " ", ("a", "epw")]),
        # held-out values: 4-digit episodes, wide ranges, big season numbers
        ("v1", "heldout-value", [("a", "t"), " - ", ("a", "e4"), " [", ("a", "R"), "] [",
                                 ("a", "d1"), "]"]),
        ("v2", "heldout-value", [("a", "t"), " ", ("a", "sbig"), " ", ("a", "rwide"), " [",
                                 ("a", "R"), "]"]),
        ("v3", "heldout-value", [("a", "t"), " ", ("a", "sc"), ("a", "e4"), " [", ("a", "R"),
                                 "] [", ("a", "g"), "]"]),
        # no work title: the row must abstain
        ("z1", "narrow-neg", [("a", "d1"), ".", ("a", "d2"), ".", ("a", "R")]),
        ("z2", "broad-neg", ["[", ("a", "d1"), "] ", ("a", "tk"), " ", ("a", "d2")]),
        ("z3", "broad-neg", [("a", "tk"), "/", ("a", "s"), "/", ("a", "d1")]),
        ("z4", "heldout-neg", [("a", "d1"), "-", ("a", "d2"), "-", ("a", "D"), " [",
                               ("a", "g"), "]"]),
    ]


LAYOUTS = {name: (tag, parts) for name, tag, parts in layout_inventory()}
# Which layouts an arm may see, and which corpus splits those layouts produced.
ARM_TAGS = {"A-narrow": ["narrow", "narrow-neg"],
            "B-broad": ["narrow", "broad", "narrow-neg", "broad-neg"]}
ARM_SPLITS = {"A-narrow": ["train-narrow"], "B-broad": ["train-narrow", "train-extra"]}
NEEDS_ALIAS = {"b6", "b7", "h6", "h10"}
LITERAL_OK = re.compile(r"^[\W_]*$", re.UNICODE)


def make_atoms(rng, lang, title, alias):
    """The row's typed fragments. Probes reuse this dict and change exactly one key."""
    season_number = rng.randrange(1, 5)
    episode = rng.randrange(1, 60)
    return {
        "t": ("TITLE", title),
        "a": ("TITLE", alias or title),
        "td": ("DISCARD", title),  # a path component repeating the title: no new evidence
        "x": ("EPISODE_TITLE", rng.choice(SUBTITLE[lang])),
        "s": ("SEASON", rng.choice(SEASON_MARKS[lang]).format(
            n=season_number, p=rng.randrange(1, 3), ord=rng.choice(ORDINALS[lang]))),
        "fm": ("SEASON", FINAL_MARK[lang]),
        "sc": ("SEASON", f"S{season_number:02d}"),
        "tc": ("SEASON", f"T{season_number:02d}"),
        "sbig": ("SEASON", f"Season {rng.randrange(10, 24)}"),
        "sw": ("SEASON", rng.choice(["Season number {n}", "Season {n}"]).format(n=season_number)),
        "n2": ("EPISODE", f"{episode:02d}"),
        "n3": ("EPISODE", f"{episode}"),
        "epp": ("EPISODE", f"EP{episode:02d}"),
        "e": ("EPISODE", f"{episode:02d}"),
        "ec": ("EPISODE", f"E{episode:02d}"),
        "e3": ("EPISODE", f"{episode:03d}"),
        "e4": ("EPISODE", f"{episode * 100:04d}"),
        "ej": ("EPISODE", f"第{episode:02d}话"),
        "epw": ("EPISODE", f"Episode {episode}"),
        "chp": ("EPISODE", f"Chapter {episode:02d}"),
        "y": ("YEAR", str(rng.randrange(1985, 2026))),
        "l": ("LIST", ",".join(f"{x:02d}" for x in sorted(rng.sample(range(1, 24), 3)))),
        "r": ("RANGE", f"{rng.randrange(1, 12):02d}-{rng.randrange(13, 24):02d}"),
        "rwide": ("RANGE", f"{rng.randrange(1, 5):02d} ~ {rng.randrange(20, 26):02d}"),
        "R": ("RESOLUTION", rng.choice(RESOLUTIONS)),
        "D": ("DIMENSIONS", rng.choice(DIMENSIONS)),
        "g": ("GROUP", rng.choice(GROUPS)),
        "d1": ("DISCARD", rng.choice(JUNK)),
        "d2": ("DISCARD", rng.choice(JUNK)),
        "tk": ("DISCARD", f"{rng.randrange(16 ** 7, 16 ** 8):x}"),
        "vt": ("DISCARD", f"v{rng.randrange(2, 4)}"),
    }


def render(atoms, parts, extension):
    query, spans = "", []
    for part in parts:
        if isinstance(part, str):
            if not LITERAL_OK.match(part):
                raise AssertionError(f"literal carries word characters: {part!r}")
            query += part
            continue
        kind, text = atoms[part[1]] if part[0] == "a" else ("DISCARD", part[1])
        spans.append({"kind": kind, "text": text, "start": len(query),
                      "end": len(query) + len(text)})
        query += text
    return query + extension, spans


def finalize(query, spans, layout, lang, family, used_queries, quarantine):
    """Contract-validate one rendered row, or quarantine it with the reason."""
    if len(query) > MAX_CHARS or query in used_queries:
        quarantine["collision-or-too-long"] += 1
        return None
    derived, issues = contract_record(query, spans, check_debris=True)
    blocking = [i for i in issues if i not in TOLERATED]
    if blocking or (derived is None and not layout.startswith("z")):
        quarantine[",".join(blocking) or "no-record"] += 1
        return None
    used_queries.add(query)
    return {"query": query, "spans": spans, "record": derived,
            "tolerated": sorted(set(issues)), "layout": layout, "language": lang,
            "family": family, "script": "cjk" if lang in ("ja", "zh") else "latin-cyrillic"}


JOINERS = (" ", " - ", "_", " | ", "-", "(", ")", " ")
PATH_ROOTS = ("Archive", "Collection", "Library", "Shelf", "Videos")
# A gap the boundary view may rewrite: separators only. Brackets are excluded
# because a span's `[` is structure, not spacing - removing it manufactures text
# like `Title1080p]`, which no real filename contains and which the contract has no
# way to reject, since unclaimed punctuation is debris-free by definition.
SEPARATOR_GAP = re.compile(r"[\s._-]+")
TIGHT_JOINERS = ("", ".", "|", "-", "_")
# Round 5's marker-role vocabulary. Deliberately disjoint from every marker word
# that appears in the held-out lexical split: `Episode`, `Chapter` and
# `Season number` must stay unseen, or the lexical axis would measure memory of
# the probe's own words instead of transfer. `test_marker_vocabulary_is_disjoint
# _from_the_heldout_lexical_split` enforces that, because a comment cannot.
EPISODE_MARKERS = ("Ep", "Eps", "Episodio", "Эпизод", "Bölüm")
EPISODE_FORMS = ("{m} {n}", "{m}{n}", "{n} {m}")


def _single_number(text):
    """The span's number, only if the text contains exactly one digit group.

    The contract's EPISODE rule requires one number, and a two-number fragment
    ("05-07") is a RANGE the augmentation must not rewrite into a fake episode.
    """
    numbers = re.findall(r"\d+", text)
    return numbers[0] if len(numbers) == 1 else None


def augment_markers(row, rng, attempts=4):
    """Attach a word marker to numeric evidence, keeping the derived record fixed.

    Training layouts express episodes as bare numbers (`05`, `E05`, `第05话`) and
    seasons through `SEASON_MARKS`, so nothing in the corpus ever tells the model
    that a *word* belongs to a number's span. That missing supervision is the
    round-4 explanation for the lexical collapse: faced with `Season 4 Episode 21`
    the model has no reason to prefer EPISODE over EPISODE_TITLE.

    Like the value and path views, gold is re-derived by the contract and the view
    is refused unless the record is byte-identical: the marker must be extra text
    inside an already-typed span, never a re-labelled value.
    """
    if row["record"] is None:
        return None
    ordered = sorted(row["spans"], key=lambda span: span["start"])
    eligible = []
    for index, span in enumerate(ordered):
        if span["kind"] not in ("EPISODE", "SEASON"):
            continue
        text = row["query"][span["start"]:span["end"]]
        number = _single_number(text)
        if number is None or len(number) > 4:
            continue
        eligible.append((index, span, text, number))
    if not eligible:
        return None
    for _ in range(attempts):
        chosen = [item for item in eligible if rng.random() < 0.7] or [rng.choice(eligible)]
        replacements = {}
        for index, span, text, number in chosen:
            if span["kind"] == "SEASON":
                # Re-render with another mark from the same language's table: the
                # number and its value are unchanged, only which word carries it.
                marks = [mark for mark in SEASON_MARKS[row["language"]]
                         if mark.format(n=number, ord=number, p=number) != text]
                if not marks:
                    continue
                value = int(number)
                replacements[index] = rng.choice(marks).format(
                    n=value, p=min(value, 3), ord=ORDINALS[row["language"]][min(value, 4) - 1])
                continue
            marker = rng.choice(EPISODE_MARKERS)
            replacements[index] = rng.choice(EPISODE_FORMS).format(m=marker, n=number)
        if not replacements:
            continue
        query, spans = _replace_span_text(row, replacements)
        if len(query) > MAX_CHARS:
            continue
        derived, issues = contract_record(query, spans, check_debris=True)
        if derived != row["record"] or [i for i in issues if i not in TOLERATED]:
            continue
        return {**row, "query": query, "spans": spans, "record": derived,
                "tolerated": sorted(set(issues)), "augmented": True,
                "augmented_axes": sorted(set(row.get("augmented_axes", [])) | {"marker"})}
    return None


def augment_boundary(row, rng, attempts=6):
    """Tighten the separator between two adjacent atoms, brackets untouched.

    Every layout in the corpus separates atoms with visible whitespace, so the
    model can treat "there is always a gap" as a free signal; the held-out layouts
    use `.`, `|`, `_` and bare adjacency instead. Round 3's error analysis put the
    residual failure on boundary cuts rather than vocabulary, so this view makes the
    cut itself the task while keeping the atom types and the record fixed.

    Only gaps made purely of separator characters are eligible, and the replacement
    is another separator, never the empty-and-bracketless deletion of structure:
    deleting `" ["` would yield `Title1080p]` - text no real filename contains. The
    contract cannot catch that (unclaimed punctuation is debris-free by design), so
    the surface form is constrained here instead, and gold equality is still checked.
    """
    if row["record"] is None:
        return None
    ordered = sorted(row["spans"], key=lambda span: span["start"])
    gaps = []
    for index in range(len(ordered) - 1):
        gap = row["query"][ordered[index]["end"]:ordered[index + 1]["start"]]
        if SEPARATOR_GAP.fullmatch(gap):
            gaps.append((index, gap))
    if not gaps:
        return None
    for _ in range(attempts):
        index, gap = rng.choice(gaps)
        left, right = ordered[index], ordered[index + 1]
        if left["kind"] == right["kind"]:
            continue
        replacement = rng.choice(TIGHT_JOINERS)
        if replacement == gap:
            continue
        start = left["end"]
        query = row["query"][:start] + replacement + row["query"][start + len(gap):]
        shift = len(replacement) - len(gap)
        spans = [dict(span) if position <= index else
                 {**span, "start": span["start"] + shift, "end": span["end"] + shift}
                 for position, span in enumerate(ordered)]
        if len(query) > MAX_CHARS:
            continue
        derived, issues = contract_record(query, spans, check_debris=True)
        if derived != row["record"] or [i for i in issues if i not in TOLERATED]:
            continue
        return {**row, "query": query, "spans": spans, "record": derived,
                "tolerated": sorted(set(issues)), "augmented": True,
                "augmented_axes": sorted(set(row.get("augmented_axes", [])) | {"boundary"})}
    return None


def augment(row, rng, attempts=4):
    """The same atoms in a different arrangement, gold re-derived by the contract.

    Round 1 measured the opposite of the obvious hypothesis: 19 hand-written
    layout families fit their own distribution to 83.6 % and scored 0.4 % on
    unseen arrangements, *worse* than 3 layouts. Adding formats teaches the model
    to look up a format. Breaking position inside every training row instead makes
    the lookup impossible, so the only signal left is local: digits, brackets,
    known group names, resolution tokens.

    Identity fragments keep their relative order. `record` joins title fragments
    in file order, so a blind shuffle would change the gold title and the
    augmentation would no longer be the same row in a different costume.

    Returns None when a shuffle lands on text the contract cannot represent -
    two adjacent numeric spans, say - rather than shipping a row whose gold is a
    guess.
    """
    original = [(span["kind"], span["text"]) for span in
                sorted(row["spans"], key=lambda s: s["start"])]
    extension = row["query"][row["query"].rindex("."):]
    for _ in range(attempts):
        parts = list(original)
        rng.shuffle(parts)
        slots = [i for i, part in enumerate(parts) if part[0] in ("TITLE", "EPISODE_TITLE")]
        identity = [part for part in original if part[0] in ("TITLE", "EPISODE_TITLE")]
        for slot, part in zip(slots, identity):
            parts[slot] = part
        query, spans = "", []
        for index, (kind, text) in enumerate(parts):
            if index:
                query += rng.choice(JOINERS)
            spans.append({"kind": kind, "text": text, "start": len(query),
                          "end": len(query) + len(text)})
            query += text
        query += extension
        if len(query) > MAX_CHARS:
            continue
        derived, issues = contract_record(query, spans, check_debris=True)
        if [i for i in issues if i not in TOLERATED] or derived != row["record"]:
            continue
        return {**row, "query": query, "spans": spans, "record": derived,
                "tolerated": sorted(set(issues)), "augmented": True}
    return None


def _replace_span_text(row, replacements):
    """Rebuild a row after replacing span text, retaining every literal byte.

    `replacements` is keyed by the span's index in start order. Rebuilding rather
    than editing offsets in place keeps span alignment auditable and preserves the
    parent layout's punctuation exactly.
    """
    old = sorted(row["spans"], key=lambda span: span["start"])
    query, spans, cursor = "", [], 0
    for index, span in enumerate(old):
        query += row["query"][cursor:span["start"]]
        text = replacements.get(index, row["query"][span["start"]:span["end"]])
        start = len(query)
        query += text
        spans.append({"kind": span["kind"], "text": text,
                      "start": start, "end": len(query)})
        cursor = span["end"]
    query += row["query"][cursor:]
    return query, spans


def augment_values(row, rng, attempts=4):
    """Broaden numeric evidence without changing any non-numeric field.

    Positional augmentation removes the layout cue that distinguished a four-digit
    episode from YEAR/RESOLUTION, but round 2 never supplied four-digit EPISODE
    evidence. This view changes digits *inside already typed spans*, re-derives the
    record through the contract, and refuses any collateral field change.
    """
    if row["record"] is None:
        return None
    ordered = sorted(row["spans"], key=lambda span: span["start"])
    eligible = [(index, span) for index, span in enumerate(ordered)
                if span["kind"] in ("EPISODE", "SEASON")
                and re.search(r"\d+", row["query"][span["start"]:span["end"]])]
    if not eligible:
        return None
    for _ in range(attempts):
        replacements = {}
        chosen = [pair for pair in eligible if rng.random() < 0.7]
        if not chosen:
            chosen = [rng.choice(eligible)]
        for index, span in chosen:
            text = row["query"][span["start"]:span["end"]]
            if span["kind"] == "EPISODE":
                # Current held-out 4-digit episodes are multiples of 100. Train on
                # the width, never their exact values.
                number = rng.randrange(1001, 10000)
                while number % 100 == 0:
                    number = rng.randrange(1001, 10000)
            else:
                # heldout-value's big seasons are 10-23.
                number = rng.randrange(24, 31)
            replacements[index] = re.sub(r"\d+", str(number), text, count=1)
        query, spans = _replace_span_text(row, replacements)
        if len(query) > MAX_CHARS:
            continue
        derived, issues = contract_record(query, spans, check_debris=True)
        if derived is None or [issue for issue in issues if issue not in TOLERATED]:
            continue
        changed = {key for key in set(row["record"]) | set(derived)
                   if row["record"].get(key) != derived.get(key)}
        if not changed or not changed <= {"season", "episodes"}:
            continue
        return {**row, "query": query, "spans": spans, "record": derived,
                "tolerated": sorted(set(issues)),
                "augmented": True,
                "augmented_axes": sorted(set(row.get("augmented_axes", [])) | {"value"})}
    return None


def augment_path(row, rng, attempts=4):
    """Place the row in randomized path components and add duplicate evidence.

    Roots are generic development vocabulary, deliberately not the held-out `TV`
    or `Series` literals. A duplicate title is DISCARD, as contract §2 requires;
    successful views must derive the parent record exactly.
    """
    if row["record"] is None:
        return None
    extension = row["query"][row["query"].rindex("."):]
    original = [(span["kind"], row["query"][span["start"]:span["end"]])
                for span in sorted(row["spans"], key=lambda span: span["start"])]
    titles = [text for kind, text in original if kind == "TITLE"]
    if not titles:
        return None
    for _ in range(attempts):
        parts = [("DISCARD", rng.choice(PATH_ROOTS)), *original]
        # A duplicate fragment is not a second title (§2). Put it in a varying
        # component so the view teaches identity comparison rather than one layout.
        parts.insert(rng.randrange(2, len(parts) + 1), ("DISCARD", rng.choice(titles)))
        boundaries = len(parts) - 1
        path_boundaries = set(rng.sample(range(boundaries), k=min(3, boundaries)))
        query, spans = "", []
        for index, (kind, text) in enumerate(parts):
            if index:
                query += rng.choice(("/", "\\")) if index - 1 in path_boundaries else rng.choice(JOINERS)
            start = len(query)
            query += text
            spans.append({"kind": kind, "text": text, "start": start, "end": len(query)})
        query += extension
        if len(query) > MAX_CHARS:
            continue
        derived, issues = contract_record(query, spans, check_debris=True)
        if derived != row["record"] or [issue for issue in issues if issue not in TOLERATED]:
            continue
        return {**row, "query": query, "spans": spans, "record": derived,
                "tolerated": sorted(set(issues)),
                "augmented": True,
                "augmented_axes": sorted(set(row.get("augmented_axes", [])) | {"path"})}
    return None


def augment_training(row, mode, position_rng, value_rng, path_rng, marker_rng=None,
                     boundary_rng=None):
    """Compose the supported axes while keeping each axis's random stream separate.

    Every mode starts from the positional shuffle, so the added axes are always
    measured *on top of* the layout cue being broken. The extra RNGs exist for the
    same reason the value and path ones do: sharing a stream with the position
    shuffle would make the round-2 paired controls depend on which later axis the
    arm happened to use.
    """
    view = augment(row, position_rng)
    if view is None:
        return None
    axes = set(view.get("augmented_axes", [])) | {"position"}
    view["augmented_axes"] = sorted(axes)
    if mode in ("shuffle-values", "shuffle-rich"):
        view = augment_values(view, value_rng)
        if view is None:
            return None
    if mode in ("shuffle-paths", "shuffle-rich"):
        view = augment_path(view, path_rng)
        if view is None:
            return None
    if mode == "shuffle-markers":
        if marker_rng is None:
            raise ValueError("shuffle-markers needs its own rng")
        view = augment_markers(view, marker_rng)
        if view is None:
            return None
    if mode == "shuffle-boundary":
        if boundary_rng is None:
            raise ValueError("shuffle-boundary needs its own rng")
        view = augment_boundary(view, boundary_rng)
        if view is None:
            return None
    return view


def build_split(rng, names, count, families, used_queries, quarantine):
    picked = []
    for index in range(count):
        lang, title, alias = families[index]
        choices = [n for n in names if n not in NEEDS_ALIAS or alias]
        if not choices:
            quarantine["no-layout-for-family"] += 1
            continue
        layout = rng.choice(choices)
        atoms = make_atoms(rng, lang, title, alias)
        query, spans = render(atoms, LAYOUTS[layout][1], rng.choice(EXTENSIONS))
        row = finalize(query, spans, layout, lang, title, used_queries, quarantine)
        if row:
            picked.append(row)
    return picked


# Paired probes: one controlled change per row, everything else identical to the
# matching `probe-base` row. The pilot showed a bare held-out number cannot be
# attributed, so each probe answers exactly one question about which factor broke
# the model. `probe-*` rows are diagnostics, never claim metrics.
PROBE_BASE = "n1"
PROBE_LETTERS = ("ru", "ja", "zh")  # the script probe needs a second script to exist


def build_probes(rng, families, used_queries, quarantine):
    base, variants = [], {f"probe-{name}": [] for name in
                          ("layout", "script", "value", "episode-title")}
    heldout = [n for n, (tag, parts) in LAYOUTS.items()
               if tag.startswith("heldout-layout") and ("a", "a") not in parts]
    base_parts = LAYOUTS[PROBE_BASE][1]
    with_title = base_parts[:5] + [("a", "x"), "-"] + base_parts[5:]
    for lang, title, alias in families:
        if lang not in PROBE_LETTERS or not alias:
            continue
        atoms = make_atoms(rng, lang, title, alias)
        query, spans = render(atoms, base_parts, ".mkv")
        parent = finalize(query, spans, PROBE_BASE, lang, title, used_queries, quarantine)
        if not parent:
            continue
        base.append(parent)
        number = len(base)
        cases = {"layout": (LAYOUTS[rng.choice(heldout)][1], None),
                 "script": (base_parts, {"t": ("TITLE", alias)}),
                 "value": (base_parts, {"e": ("EPISODE", f"{number * 100:04d}")}),
                 "episode-title": (with_title, None)}
        for name, (parts, changed) in cases.items():
            probe_atoms = dict(atoms)
            probe_atoms.update(changed or {})
            query, spans = render(probe_atoms, parts, ".mkv")
            row = finalize(query, spans, f"probe-{name}", lang, title, used_queries, quarantine)
            if row:
                row["probe_parent"] = parent["query"]
                variants[f"probe-{name}"].append(row)
    return {"probe-base": base, **variants}


def generate():
    if ROOT.exists():
        raise SystemExit(f"refusing to overwrite the locked corpus at {ROOT}")
    ROOT.mkdir(parents=True)
    rng = random.Random(SEED)
    families = []
    seen_titles = set()
    for lang in LANGS:
        words = WORDS[lang]
        for a, b, c in [(x, y, z) for x in words for y in words for z in words
                        if x != y and y != z]:
            title = f"{a} {b} {c}"
            if title in seen_titles:
                continue  # an identical string in two languages would leak families
            seen_titles.add(title)
            table = dict(zip(words, ALIAS.get(lang, [])))
            families.append((lang, title, " ".join(table.get(w, w) for w in title.split(" "))
                             if lang in ALIAS else None))
    rng.shuffle(families)
    spec = [("train-narrow", ["narrow", "narrow-neg"], 1600),
            ("train-extra", ["broad", "broad-neg"], 3200),
            ("validation", ["narrow", "broad", "narrow-neg", "broad-neg"], 240),
            ("easy", ["narrow", "broad", "narrow-neg", "broad-neg"], 240),
            ("heldout-layout", ["heldout-layout-a", "heldout-layout-b"], 240),
            ("heldout-lexical", ["heldout-lexical"], 160),
            ("heldout-value", ["heldout-value"], 160),
            ("heldout-negative", ["heldout-neg"], 120)]
    total = sum(count for _, _, count in spec) + 600  # 600 more families for probes
    if len(families) < total:
        raise SystemExit(f"only {len(families)} disjoint title families for {total} rows")
    manifest = {"seed": SEED, "max_chars": MAX_CHARS, "schema": "contract v2.1",
                "labels_from": "authored layout atoms, derived by fp.normalize.record",
                "claim": "format generalization on controlled data, not real-world accuracy",
                "arms": {k: {"layout_tags": v, "train_splits": ARM_SPLITS[k]}
                             for k, v in ARM_TAGS.items()},
                "layouts": {name: {"tag": tag, "atoms": sorted({p[1] for p in parts
                                                                if not isinstance(p, str)})}
                            for name, (tag, parts) in LAYOUTS.items()},
                "splits": {}, "quarantine": {}}
    cursor, used, quarantine = 0, set(), Counter()

    def emit(split, rows):
        path = ROOT / f"{split}.jsonl"
        path.write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in rows))
        manifest["splits"][split] = {
            "n": len(rows), "sha256": digest(path),
            "layouts": sorted({r["layout"] for r in rows}),
            "languages": dict(Counter(r["language"] for r in rows)),
            "fields": dict(Counter(k for r in rows for k in (r["record"] or {}))),
            "abstentions": sum(1 for r in rows if r["record"] is None),
            "max_chars": max(len(r["query"]) for r in rows)}
        for name in manifest["splits"][split]["layouts"]:
            if split.startswith("train") and LAYOUTS[name][0].startswith("heldout"):
                raise AssertionError(f"held-out layout {name} leaked into {split}")

    for split, tags, count in spec:
        names = [n for n, (tag, _) in LAYOUTS.items() if tag in tags]
        subset = families[cursor:cursor + count]
        cursor += count
        rows = build_split(rng, names, count, subset, used, quarantine)
        assert rows, f"{split} produced nothing"
        emit(split, rows)
    # Probes get their own families, so a probe row never shares a title with a
    # scored split; the pairing is inside the probe set, not across splits.
    probe_subset = families[cursor:cursor + 600]
    for split, rows in build_probes(rng, probe_subset, used, quarantine).items():
        emit(split, rows)
    manifest["quarantine"] = dict(quarantine)
    (ROOT / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n")
    print(json.dumps({k: v["n"] for k, v in manifest["splits"].items()} |
                     {"quarantine": dict(quarantine)}, ensure_ascii=False, indent=2))


def rows(split):
    manifest = json.loads((ROOT / "manifest.json").read_text())
    path = ROOT / f"{split}.jsonl"
    data = [json.loads(line) for line in path.read_text().splitlines()]
    assert digest(path) == manifest["splits"][split]["sha256"], f"{split}: hash mismatch"
    assert len(data) == manifest["splits"][split]["n"], f"{split}: row count changed"
    return data


def verify():
    manifest = json.loads((ROOT / "manifest.json").read_text())
    problems, queries, families = [], set(), {}
    for split in manifest["splits"]:
        group = "probes" if split.startswith("probe-") else split
        for row in rows(split):
            queries.add(row["query"])
            owner = families.setdefault(row["family"], group)
            if owner != group:
                problems.append(f"{row['family']}: family in {group} and {owner}")
            derived, issues = contract_record(row["query"], row["spans"], check_debris=True)
            if derived != row["record"] or sorted(set(issues)) != row["tolerated"]:
                problems.append(f"{row['query'][:40]}: gold not reproducible {issues}")
            if len(row["query"]) > manifest["max_chars"]:
                problems.append(f"{row['query'][:40]}: exceeds locked length")
    print(json.dumps({"rows": len(queries), "families": len(families),
                      "problem_count": len(problems), "problems": problems[:20]},
                     ensure_ascii=False, indent=2))
    return 1 if problems else 0


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--verify", action="store_true")
    args = parser.parse_args(argv)
    return verify() if args.verify else (generate(), 0)[1]


if __name__ == "__main__":
    raise SystemExit(main())