src/scripts/twolayer_attend.py

"""Job 191 (runs/CSEL/ATTESTED_END_PREREG.md): ATTESTED LINE ENDS at decode, on top of 188's tl24fix decode
(twolayer_full.py decode --stutter-fix, syllable unit, model runs/TLLONG/model_h24, no training).

Why (189 failure analysis, runs/CONFIRM2/fail/): 45 of 55 judge-flagged system วรรค have the error span reaching the line
END - the line-final (rhyme) word does not attach to the word before it (ได้อย่างไรไหม, ที่ซึ้งไหน, ละมุนมาน); ~10 poems are
whole-word stutters (ขวัญขวัญขวัญ, ตามใจตามใจ). In the syllable unit the line-final rhyme cell is committed EARLY (it is the
most constrained cell; committing it first is doing real work, job 173), i.e. before the words to its left exist.

The change (two decode flags; both default OFF; with both off the output is byte-identical to twolayer_full.py decode):
  --attested-end BIGRAMS.json  when a วรรค's last cell is committed (the line is complete), its text is cut with pythainlp
        newmm (= runs/CONFIRM2/fail/bigram.py) and the LAST word-bigram is looked up in the training corpus's bigram set.
        If unattested, the cells of the line's final lexicon word are RE-OPENED and re-filled by the same model and the
        same commit rule (one syllable per step, by confidence), restricted to that region, now with the full left context
        visible; the failed word's last-syllable token is banned at the final cell, so each try draws a different word.
        Rhyme family, tone class and receiver position stay fixed (the re-opened rhyme cell gets its family back in the
        sound layer, exactly as the wak/bot refinement does), so the rhyme still holds. Up to --end-tries K (default 6);
        if no try is attested, the ORIGINAL word is restored (no change; counted as failed). In-model: no external text,
        no LLM.
  --word-stutter WREP.json     the stutter ban (--stutter-fix) covers same-SOUND adjacent syllables only, so a repeated
        MULTI-syllable word (ตามใจตามใจ) is never banned, and an attested pair (ขวัญ|ขวัญ, 2x in training) allows a triple
        (ขวัญขวัญขวัญ). Adds: a multi-syllable word may not sit directly next to itself unless that word repeat is
        attested >= 2x in the training grids; no word three times in a row unless that triple is attested >= 2x.

  bigrams TRAIN OUT.json              newmm word-bigram set of the training poems (line by line), sha-stamped
  wordrep DATA_DIR OUT.json [MIN]     adjacent whole-word repeat pairs / triples in DATA_DIR/poems.jsonl grids
  human DEV PROV OUT.jsonl            dev topics' human controls (confirm2_select.tier_pick on the test split, spaces stripped)
  decode ... (= twolayer_full.py decode) [--attested-end B.json [--end-tries K]] [--word-stutter W.json]
  gates OUT ATT BIGRAMS A-B CTL=DIR E=DIR S=DIR   pre-judge gates per arm + the split proxy readout -> OUT/GATES.json; PASSING ...
  pools OUT A-B HUMAN C2HUMAN CTL=DIR [E=DIR] [S=DIR]   human-balanced pools per seed and topic quarter: system labels ctl/ate/ats,
                                               as many human labels hum/humb/humc (1:1), + anchors (build_pool, 1 per 5)
  score OUT A-B CTL=DIR [E=DIR] [S=DIR]        after judging runs/ATE_{CLEAN,ALIGN}{s}{q}/run1.json: one decision per arm
Job 191 decodes TWO arms, one change each: E = --attested-end only, S = --word-stutter only (both on top of --stutter-fix).
  test                                         CPU tests (toy + real lexicon with a fake model; no weights)
New file; nothing shared is patched (twolayer_full / twolayer_canvas are imported, never edited).
"""
import os, sys, json, collections, argparse, hashlib
sys.path[:0] = [".", "scripts"]

PREFIX = "ATE"
LABELS = {"CTL": "ctl", "E": "ate", "S": "ats"}      # job arm -> pool label: E = --attested-end only, S = --word-stutter only
HUM, HUMB, HUMC = "hum", "humb", "humc"
HUMANS = [HUM, HUMB, HUMC]
NPART = 4                                            # topic parts per seed: 3 system + 3 human labels x 13 + anchors <= 20


def rows(p): return [json.loads(l) for l in open(p, encoding="utf-8") if l.strip()]


def wrows(p, R):
    with open(p + ".tmp", "w", encoding="utf-8") as f:
        for r in R: f.write(json.dumps(r, ensure_ascii=False) + "\n")
    os.replace(p + ".tmp", p)


def sha(p): return hashlib.sha256(open(p, "rb").read()).hexdigest()
def seeds(spec): a, b = map(int, spec.split("-")); return range(a, b + 1)


# ---------------------------------------------------------------- word bigrams (= runs/CONFIRM2/fail/bigram.py)

def toks(s):
    from pythainlp.tokenize import word_tokenize
    return [w for w in word_tokenize(s, engine="newmm") if w.strip()]


def build_bigrams(train, out):
    import pythainlp
    """newmm bigrams of every training line AS WRITTEN (= runs/CONFIRM2/fail/bigram.py, the validated proxy) UNION those of the
    same line with spaces removed: decoded lines never contain spaces, and a space can split a word pair the poet wrote
    (review 2026-10-01: 3 of 588 flagged-line checks differ). The added pairs are counted and printed."""
    bg, bg_raw, n = set(), set(), 0
    for l in open(train, encoding="utf-8"):
        for line in json.loads(l)["poem"].split("\n"):
            t = toks(line); bg_raw.update(zip(t, t[1:])); n += 1
            if " " in line: t = toks(line.replace(" ", "")); bg.update(zip(t, t[1:]))
    bg |= bg_raw
    d = {"source": train, "source_sha256": sha(train), "tokenizer": f"pythainlp {pythainlp.__version__} newmm",
         "n_lines": n, "n_raw_pairs": len(bg_raw), "pairs": sorted(f"{a}\t{b}" for a, b in bg)}
    json.dump(d, open(out, "w", encoding="utf-8"), ensure_ascii=False, indent=0)
    print(f"  bigrams {out}: {len(bg)} distinct newmm word bigrams over {n} lines ({len(bg_raw)} as written = 189's proxy set, "
          f"+{len(bg) - len(bg_raw)} from space-stripped lines); sha256 {sha(out)}")


def load_bigrams(path):
    import pythainlp
    d = json.load(open(path, encoding="utf-8"))
    want = f"pythainlp {pythainlp.__version__} newmm"
    assert d["tokenizer"] == want, f"bigrams built with {d['tokenizer']}, running {want}"
    return frozenset(tuple(x.split("\t")) for x in d["pairs"])


def end_attested(line, BG):
    t = toks(line); return len(t) < 2 or (t[-2], t[-1]) in BG


PROXY_KEYS = ("end_unattested", "adj_repeat", "adj_unattested", "tri_nonadj")


def proxy(poem, BG):
    """189's failure proxy per poem, split (job 191 stand-in: the combined 'adjacent-or-3x' count mixed three things):
      end_unattested  lines whose last newmm bigram never occurs in training
      adj_repeat      lines with an adjacent repeated token (a a), any
      adj_unattested  lines with an adjacent repeat (a a) that training never has (attested reduplication like นานนาน excluded)
      tri_nonadj      lines where a token occurs >= 3 times but never adjacently (topic echo: นานแล้ว..นาน..นาน)"""
    out = dict.fromkeys(PROXY_KEYS, 0)
    for ln in poem.split("\n"):
        t = toks(ln); c = collections.Counter(t)
        adj = [(a, b) for a, b in zip(t, t[1:]) if a == b]
        out["end_unattested"] += int(len(t) > 1 and (t[-2], t[-1]) not in BG)
        out["adj_repeat"] += int(bool(adj)); out["adj_unattested"] += int(any(p not in BG for p in adj))
        out["tri_nonadj"] += int(bool(t) and not adj and max(c.values()) >= 3)
    return out


# ---------------------------------------------------------------- whole-word repeats in the training grids

def line_words(cells):
    """grid line [word, k, n, phon]... -> word occurrences"""
    return [c[0] for c in cells if c[1] == 0]


def word_repeats(grids):
    pairs, triples = collections.Counter(), collections.Counter()
    for line in grids:
        w = line_words(line)
        pairs.update(a for a, b in zip(w, w[1:]) if a == b)
        triples.update(a for a, b, c in zip(w, w[1:], w[2:]) if a == b == c)
    return pairs, triples


def build_wordrep(data_dir, out, min_n=2):
    p, t = word_repeats(line for l in open(f"{data_dir}/poems.jsonl", encoding="utf-8") for line in json.loads(l)["grid"])
    d = {"min": int(min_n), "source": f"{data_dir}/poems.jsonl", "source_sha256": sha(f"{data_dir}/poems.jsonl"),
         "pairs": {w: n for w, n in sorted(p.items()) if n >= int(min_n)}, "triples": {w: n for w, n in sorted(t.items()) if n >= int(min_n)}}
    json.dump(d, open(out, "w", encoding="utf-8"), ensure_ascii=False, indent=0)
    print(f"  word repeats {out}: pairs {len(d['pairs'])} of {len(p)} attested >= {min_n}x, triples {len(d['triples'])} of {len(t)}")


def load_wordrep(path, X):
    d = json.load(open(path, encoding="utf-8"))
    ids = lambda ws: frozenset(X.wid[w] for w in ws if X.wid.get(w, 0) >= 4)
    return ids(d["pairs"]), ids(d["triples"])


# ---------------------------------------------------------------- decoder

def make_decoder_class():
    import torch
    import twolayer_canvas as TC, twolayer_full as TF

    class AttendDecoder(TF.make_decoder_class()):
        """StutterDecoder + whole-word stutter ban (WREP) + attested line ends (BIGRAMS). Class attributes are set by
        decode_cli; None = that part off (then this class behaves exactly as StutterDecoder)."""
        BIGRAMS = None; END_TRIES = 6; WREP = None

        # ---- whole-word stutters
        def _tok_at(self, i, j):
            if not (0 <= j < self.lens[i]): return None
            t = self.neighbour_tok(self.idx(i, j))
            return t if t is not None and t >= 4 else None

        def word_ending(self, i, j):
            """(word id, start) of the word whose LAST syllable sits (committed or forced) at cell j, else None"""
            t = self._tok_at(i, j); X = self.X
            if t is None or int(X.tok_k[t]) != int(X.tok_n[t]) - 1: return None
            return int(X.tok_word[t]), j - int(X.tok_n[t]) + 1

        def word_starting(self, i, j):
            t = self._tok_at(i, j); X = self.X
            if t is None or int(X.tok_k[t]) != 0: return None
            return int(X.tok_word[t]), j + int(X.tok_n[t]) - 1

        def word_stutter_bans(self, i, j):
            """token ids banned at open cell (i, j): first syllable of the word just left of it / last syllable of the word
            just right of it, when that would make an unattested multi-syllable pair or any unattested triple"""
            X = self.X; pair_ok, tri_ok = self.WREP; out = []
            wl = self.word_ending(i, j - 1)
            if wl is not None:
                w, s = wl; n = len(X.sounds[X.words[w]])
                bad = (n >= 2 and w not in pair_ok)
                if not bad and w not in tri_ok:
                    prev = self.word_ending(i, s - 1); nxt = self.word_starting(i, j + n)
                    bad = (prev is not None and prev[0] == w) or (nxt is not None and nxt[0] == w)
                if bad: out.append(X.first[w])
            wr = self.word_starting(i, j + 1)
            if wr is not None:
                w, e = wr; n = len(X.sounds[X.words[w]])
                bad = (n >= 2 and w not in pair_ok)
                if not bad and w not in tri_ok:
                    nxt = self.word_starting(i, e + 1); prev = self.word_ending(i, j - n)
                    bad = (nxt is not None and nxt[0] == w) or (prev is not None and prev[0] == w)
                if bad: out.append(X.first[w] + n - 1)
            return out

        def ban_word_stutters(self, lp):
            for i in range(TC.L):
                for j in range(self.lens[i]):
                    k = self.idx(i, j)
                    if k in self.committed: continue
                    for t in self.word_stutter_bans(i, j):
                        if float(lp[k, t]) > -1e8: self.n_word_banned += 1
                        lp[k, t] = lp[k, t] - 1e9
            return lp

        def lp(self, hold=True):
            lp = super().lp(hold)
            if self.WREP is not None: self.ban_word_stutters(lp)
            for k, ts in getattr(self, "end_ban", {}).items():
                if k not in self.committed: lp[k, torch.tensor(sorted(ts), device=lp.device)] -= 1e9
            return lp

        # ---- attested line ends
        def held(self):
            h = super().held()
            only = getattr(self, "_only", None)
            if only is not None:
                h |= {self.idx(i, j) for i in range(TC.L) if i != only for j in range(self.lens[i])}
            return h

        def line_done(self, i):
            return all(self.idx(i, j) in self.committed for j in range(self.lens[i]))

        def end_ok(self, i):
            self.ae_checks += 1
            return end_attested(self.text()[i], self.BIGRAMS)

        def reopen(self, cells):
            """= Decoder.remask for an explicit cell set: a re-opened rhyme cell shows its family again"""
            for k in cells:
                self.committed.discard(k); self.tok[k] = TC.MASK; self.snd[k] = 0; self.forced.pop(k, None); self.ri_in[k] = 0
                for gi in self.fam:
                    if any(self.idx(*c) == k for c in self.fam_cells(gi)): self.ri_in[k] = self.fam[gi]
            for key in [key for key in self.wscore if self.idx(*key) in cells]: del self.wscore[key]

        def snapshot(self):
            """everything a re-draw can consume or change: canvas, commits, families, scores, topic rows, RNG, counters"""
            import copy
            return {"tok": self.tok.clone(), "snd": self.snd.clone(), "ri_in": self.ri_in.clone(), "committed": set(self.committed),
                    "forced": dict(self.forced), "fam": dict(self.fam), "gpos": dict(self.gpos), "wscore": dict(self.wscore),
                    "topic_rows": [r.clone() for r in self.topic_rows], "gen": self.gen.get_state(), "n_fwd": self.n_fwd,
                    "n_banned": self.n_banned, "n_word_banned": self.n_word_banned,
                    "end_ban": copy.deepcopy(getattr(self, "end_ban", {}))}

        def restore(self, S):
            self.tok, self.snd, self.ri_in = S["tok"].clone(), S["snd"].clone(), S["ri_in"].clone()
            self.committed, self.forced, self.fam, self.gpos = set(S["committed"]), dict(S["forced"]), dict(S["fam"]), dict(S["gpos"])
            self.wscore = dict(S["wscore"]); self.topic_rows = [r.clone() for r in S["topic_rows"]]; self.gen.set_state(S["gen"])
            self.n_fwd, self.n_banned, self.n_word_banned = S["n_fwd"], S["n_banned"], S["n_word_banned"]
            self.end_ban = S["end_ban"]

        def fix_end(self, i, t):
            """re-draw an unattested line-final word (see the module doc). A FAILED attempt restores the whole decoder state,
            RNG included, so the rest of the poem is exactly what the decode without the flag gives; its forward passes are
            logged in ae_fwd (not in forward_passes)."""
            if self.end_ok(i): return
            X = self.X; self.ae_unattested += 1
            last = self.idx(i, self.lens[i] - 1); n = int(X.tok_n[int(self.tok[last])])
            region = [self.idx(i, j) for j in range(self.lens[i] - n, self.lens[i])]
            S = self.snapshot()
            self.end_ban = getattr(self, "end_ban", {})
            for _ in range(self.END_TRIES):
                self.end_ban.setdefault(last, set()).add(int(self.tok[last]))
                self.reopen(region); self._only = i; filled = True
                try:
                    while any(k not in self.committed for k in region):
                        if not self.step_cell_or_word(t): filled = False; break
                finally: self._only = None
                self.ae_tries += 1
                if filled and self.end_ok(i): self.ae_fixed += 1; return
                if not filled: break
            self.ae_fwd += self.n_fwd - S["n_fwd"]
            self.restore(S); self.ae_failed += 1                          # no attested end: everything as before

        def run(self):
            self.n_word_banned = 0; self.ae_checks = self.ae_unattested = self.ae_fixed = self.ae_failed = self.ae_tries = self.ae_fwd = 0
            if self.BIGRAMS is None or self.unit not in ("syllable", "word"): return super().run()
            ok, done = True, set()
            for t in range(200):
                if not any(self.idx(i, j) not in self.committed for i in range(TC.L) for j in range(self.lens[i])): break
                if not self.step_cell_or_word(t) and t >= self.fam_delay: ok = False; break
                for i in range(TC.L):
                    if i not in done and self.line_done(i): done.add(i); self.fix_end(i, t)
            return self.text(), ok

    return AttendDecoder


def decode_cli(argv):
    """= twolayer_full.py decode (same records, same RNG streams) + --attested-end / --word-stutter; extra keys only with them"""
    import torch, random as _r
    import twolayer_canvas as TC, twolayer_full as TF
    from klon8 import prosody as P
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True); ap.add_argument("--prompts", required=True); ap.add_argument("--out", required=True)
    ap.add_argument("--seed", type=int, default=0); ap.add_argument("--temp", type=float, default=0.8); ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--unit", choices=["syllable", "word", "wak", "bot"], default="word")
    ap.add_argument("--refine", type=int, default=3); ap.add_argument("--topic-boost", type=float, default=1.5)
    ap.add_argument("--stutter-fix", action="store_true", help="ban same-sound neighbours unless attested (default off)")
    ap.add_argument("--allow", default="", help="allowlist from twolayer_full.py allow (required with --stutter-fix)")
    ap.add_argument("--attested-end", default="", help="bigram set from `bigrams`: re-draw unattested line-final words (default off)")
    ap.add_argument("--end-tries", type=int, default=6)
    ap.add_argument("--word-stutter", default="", help="word-repeat allowlist from `wordrep`: ban whole-word stutters (default off)")
    a = ap.parse_args(argv)
    model, X = TC.load(a.model)
    new = bool(a.attested_end or a.word_stutter)
    if a.stutter_fix:
        assert a.allow, "--stutter-fix needs --allow"
        SD = make_decoder_class(); SD.ALLOW = TF.load_allow(a.allow, X)
        if a.attested_end:
            assert a.unit in ("syllable", "word"), "--attested-end: syllable or word unit"
            SD.BIGRAMS = load_bigrams(a.attested_end); SD.END_TRIES = a.end_tries
        if a.word_stutter: SD.WREP = load_wordrep(a.word_stutter, X)
        TC.Decoder = SD                                      # decode_one builds TC.Decoder (this process only)
    else: assert not new, "--attested-end / --word-stutter extend --stutter-fix"
    rows_ = [json.loads(l) for l in open(a.prompts, encoding="utf-8")]
    if a.limit: rows_ = rows_[:a.limit]
    rng_len = _r.Random(a.seed)
    with open(a.out + ".tmp", "w", encoding="utf-8") as f:
        for n, r in enumerate(rows_):
            text = r.get("text") or r.get("prompt") or r.get("title") or r.get("topic")
            assert text, f"prompt row without text: {r}"
            want = TC.pick_lens(rng_len)
            lines, D, tries = TC.decode_one(model, X, text, unit=a.unit, temp=a.temp, lens=want, refine=a.refine,
                                            topic_boost=a.topic_boost, gen=torch.Generator().manual_seed(a.seed * 100003 + n))
            poem = "\n".join(lines); sc = P.score(poem, n_bot_expected=1)
            got = [len(P.g2p_line(x)) for x in lines]
            rec = {"id": r.get("id"), "prompt": text, "sample": 0, "poem": poem, "unit": a.unit, "syl_target": want,
                   "syl": got, "retries": tries, "forward_passes": D.n_fwd, "all_pass": sc["all_pass"],
                   "structure": sc["structure"], "ext_rhyme": sc["ext_rhyme"], "tone": sc["tone"], "internal": sc["internal"]}
            if a.stutter_fix:
                rec["stutter_fix"] = True; rec["n_banned"] = D.n_banned; rec["stutters_left"] = len(TF.stutter_pairs(D))
            if a.word_stutter: rec["word_stutter"] = True; rec["n_word_banned"] = D.n_word_banned
            if a.attested_end:
                rec["attested_end"] = True
                for k in ("ae_checks", "ae_unattested", "ae_fixed", "ae_failed", "ae_tries", "ae_fwd"): rec[k] = getattr(D, k)
                rec["end_unattested_left"] = sum(not end_attested(x, SD.BIGRAMS) for x in lines)
            f.write(json.dumps(rec, ensure_ascii=False) + "\n"); f.flush()
            print(f"{n}/{len(rows_)} all_pass={sc['all_pass']} syl={got} want={want} retries={tries}"
                  + (f" ends unattested {D.ae_unattested} fixed {D.ae_fixed} failed {D.ae_failed}" if a.attested_end else ""),
                  file=sys.stderr, flush=True)
    os.replace(a.out + ".tmp", a.out)
    P.save_cache()


# ---------------------------------------------------------------- human controls for the dev topics

def build_human(dev, prov, out):
    """one human poem per dev topic: the TEST-split stanza under the same title key, confirm2_select.tier_pick (opening post
    stanza 0 > opening post's first stanza > the thread's first stanza; ties by lowest id), spaces stripped (189's rule)"""
    import fresh_titles as FT, confirm2_select as CS
    by = collections.defaultdict(list)
    for r in rows(prov): by[FT.key(r["title"])].append(r)
    out_, tiers = [], collections.Counter()
    for d in rows(dev):
        k = FT.key(d["text"]); assert k == d["text"], f"{d['id']}: topic is not its own key"
        assert by.get(k), f"{d['id']}: no human stanza in the test split"
        tier, r = CS.tier_pick(by[k]); tiers[tier] += 1
        poem = r["poem"].replace(" ", "")
        assert len(poem.split("\n")) == 4 and all(poem.split("\n")), f"{d['id']}: not 4 lines"
        out_.append({"id": d["id"], "prompt": d["text"], "poem": poem, "src_id": r["id"], "thread": r["thread"], "post": r["post"],
                     "stanza": r["stanza"], "tier": tier})
    wrows(out, out_)
    print(f"  {out}: {len(out_)} human controls, tiers {dict(sorted(tiers.items()))}; sha256 {sha(out)}")


# ---------------------------------------------------------------- gates (pre-judge)

def arm_specs(specs):
    A = [tuple(x.split("=", 1)) for x in specs]
    assert A and A[0][0] == "CTL" and len({n for n, _ in A}) == len(A) and all(n in LABELS for n, _ in A), \
        f"arms: CTL=DIR then E=DIR and/or S=DIR ({list(LABELS)})"
    return A


def gate_fail(c, f):
    """pre-registered, per arm (ATTESTED_END_PREREG.md, gates): form >= 99% (absolute); UNATTESTED stutters <= control + 0.10
    per poem (188's gate); and three guards against buying the proxy with blander or echoing text (reviews 2026-10-01):
    final-word bits >= control - 0.3, word 3x in a line (non-adjacent; topic/word echo) <= control + 0.05 lines per poem,
    distinct-2 >= control - 0.03"""
    dead = []
    if f["form"] < 99: dead.append("form < 99%")
    if f["stutters_unattested"] > c["stutters_unattested"] + 0.10: dead.append("unattested stutters")
    if f["final_bits"] < c["final_bits"] - 0.3: dead.append("final-word bits")
    if f["tri_nonadj"] > c["tri_nonadj"] + 0.05: dead.append("word echo (3x in a line)")
    if f["distinct2"] < c["distinct2"] - 0.03: dead.append("distinct-2")
    return dead


def topic_echo(prompt, poem):
    """poem tokens that are content tokens of the topic (newmm, >= 2 characters), per poem (reported)"""
    tt = {w for w in toks(prompt) if len(w) >= 2}
    return sum(w in tt for ln in poem.split("\n") for w in toks(ln))


def gates(out, att_json, bigrams, spec, *arms):
    import twolayer_pipe as TP, twolayer_full as TF
    A = arm_specs(arms); BG = load_bigrams(bigrams)
    att = frozenset(json.load(open(att_json, encoding="utf-8"))["kept"]); assert len(att) > 100, "attested set too small"
    S = {}
    for n, d in A:
        R = [r for s in seeds(spec) for r in rows(f"{d}/c{s}.jsonl")]
        px = [proxy(r["poem"], BG) for r in R]
        S[n] = dict(TP.stats(R), stutters_unattested=sum(TF.stutters_unattested(r["poem"], att) for r in R) / len(R),
                    proxy_clean=100 * sum(p["end_unattested"] == 0 and p["adj_unattested"] == 0 for p in px) / len(R),
                    topic_echo=sum(topic_echo(r["prompt"], r["poem"]) for r in R) / len(R),
                    **{k: sum(p[k] for p in px) / len(R) for k in PROXY_KEYS})
        for k in ("ae_unattested", "ae_fixed", "ae_failed", "ae_tries", "n_word_banned"):
            S[n][k] = sum(r.get(k, 0) for r in R) / len(R)
    res = {"stats": S, "dead": {}, "reported_186": {}}
    for n, s in S.items():
        print(f"  {n:4s} n {s['n']:4d}, form {s['form']:.1f}%, stutters raw {s['stutters']:.3f} unattested {s['stutters_unattested']:.3f}, "
              f"repeats {s['repeats']:.2f}, distinct-2 {s['distinct2']:.3f}, final bits {s['final_bits']:.2f}, unseen {s['final_word_unseen']:.3f}")
        print(f"       PROXY lines/poem: unattested end {s['end_unattested']:.3f}; adjacent word repeat {s['adj_repeat']:.3f} "
              f"(unattested {s['adj_unattested']:.3f}, attested reduplication {s['adj_repeat'] - s['adj_unattested']:.3f}); "
              f"word 3x in a line (not adjacent) {s['tri_nonadj']:.3f}; proxy-clean {s['proxy_clean']:.1f}%; topic-word tokens/poem {s['topic_echo']:.2f}")
    for n, _ in A[1:]:
        a = S[n]
        if n == "E": print("  NOTE: E optimises the end-bigram proxy directly, so its proxy readout is circular (and human dev lines end "
                           "on an unattested bigram MORE often than the control: 37% vs 25%) - only the judge decides")
        if n == "E": print(f"  E per poem: lines with an unattested end {a['ae_unattested']:.2f} -> fixed {a['ae_fixed']:.2f}, failed "
                           f"(original kept) {a['ae_failed']:.2f}, re-draws {a['ae_tries']:.2f}")
        if n == "S": print(f"  S per poem: whole-word stutter bans {a['n_word_banned']:.1f} (cell x token, summed over steps)")
        res["dead"][n] = gate_fail(S["CTL"], a)
        res["reported_186"][n] = [x for x in TP.gate_fail(S["CTL"], a) if x != "form < 98%"]
        print(f"  {n}: " + ("KILLED before judging: " + ", ".join(res["dead"][n]) if res["dead"][n] else "pre-judge gates PASS")
              + (f"; 186's gates vs the control (reported only): {', '.join(res['reported_186'][n])}" if res["reported_186"][n] else "; 186's gates: all within"))
    os.makedirs(out, exist_ok=True); json.dump(res, open(f"{out}/GATES.json", "w"), indent=1)
    print("PASSING " + " ".join(n for n, d in res["dead"].items() if not d))


# ---------------------------------------------------------------- pools (human-balanced)

def parts(ids, k=NPART):
    """sorted ids in k contiguous near-equal parts (49 -> 13/12/12/12)"""
    ids = sorted(ids); q, r = divmod(len(ids), k); out, i = [], 0
    for j in range(k):
        n = q + (j < r); out.append(ids[i:i + n]); i += n
    return out


OPENING_TIERS = (1, 2)      # dev human stanzas from the thread's opening post: written to the title (reply-topic rule)


def c2_rows(c2, s, label, n_topics):
    """189's human controls for one seed, rotating: humb takes slots (s*n + j) mod N, humc (s*n + n + j), and the hum
    substitutes (s*n + 2n + j) (n topics, N poems). Within a seed each label is distinct; a poem in two labels sits >= 17
    topic slots apart for n = 49, N = 81 (offsets 49, 98 = 17, and N - those), so never in one part (<= 13 topics).
    pools() asserts it."""
    H = rows(c2); N = len(H); assert N >= n_topics, f"need >= {n_topics} 189 human poems"
    off = s * n_topics + {HUMB: 0, HUMC: n_topics, HUM: 2 * n_topics}[label]
    return [H[(off + j) % N] for j in range(n_topics)]


def pools(out, spec, human, c2human, *arms):
    """per seed s and topic part q (NPART parts), equal-size labels (build_pool asserts the balance), RAW text:
      system labels: ctl + every passing arm (ate = E, ats = S); human labels, as many as system labels:
      hum  = the dev topic's human stanza under the same topic string, ONLY if it is an opening-post stanza (tier 1/2:
             written to the thread title); a REPLY stanza must not be judged under the thread title (reply-topic rule), so
             for those topics hum carries a 189 human control under its own frozen topic instead (id G..._sS)
      humb, humc = 189's human controls (own topics)"""
    A = arm_specs(arms); Hall = {r["id"]: r for r in rows(human)}
    Hd = {i: r for i, r in Hall.items() if r["tier"] in OPENING_TIERS}
    hl = HUMANS[:len(A)]; assert len(hl) == len(A), "not enough human labels for 1:1"
    os.makedirs(out, exist_ok=True); n = 0
    for s in seeds(spec):
        R = {name: {r["id"]: r for r in rows(f"{d}/c{s}.jsonl")} for name, d in A}
        ids = sorted(R["CTL"])
        assert all(set(R[k]) == set(ids) for k in R) and all(len(rows(f"{d}/c{s}.jsonl")) == len(ids) for _, d in A), f"seed {s}: topics differ"
        assert all(R[k][i]["prompt"] == R["CTL"][i]["prompt"] == Hall[i]["prompt"] for k in R for i in ids), f"seed {s}: prompts differ"
        assert all(R[k][i]["unit"] == "syllable" for k in R for i in ids), f"seed {s}: not the syllable unit"
        slot = {i: j for j, i in enumerate(ids)}
        C2 = {lab: c2_rows(c2human, s, lab, len(ids)) for lab in hl}
        for q, part in enumerate(parts(ids)):
            for name, _ in A:
                J = [{"id": f"{i}_s{s}", "prompt": R[name][i]["prompt"], "poem": R[name][i]["poem"]} for i in part]
                assert all(len(r["poem"].split("\n")) == 4 for r in J), f"{name} seed {s}: poems without exactly 4 lines"
                wrows(f"{out}/{LABELS[name]}_s{s}q{q}.jsonl", J); n += len(J)
            used = set()
            for lab in hl:
                J = []
                for i in part:
                    if lab == HUM and i in Hd: J.append({"id": f"{i}_s{s}", "prompt": Hd[i]["prompt"], "poem": Hd[i]["poem"]})
                    else:
                        r = C2[lab][slot[i]]
                        assert r["id"] not in used, f"seed {s} part {q}: 189 human poem {r['id']} twice in one pool"
                        used.add(r["id"]); J.append({"id": f"{r['id']}_s{s}", "prompt": r["prompt"], "poem": r["poem"]})
                wrows(f"{out}/{lab}_s{s}q{q}.jsonl", J)
    print(f"  human-balanced pools: system {[LABELS[x] for x, _ in A]}, human {hl}, seeds {spec}, {NPART} parts per seed, "
          f"{n} system poems + as many human; hum = the dev topic's opening-post stanza for {len(Hd)} of {len(Hall)} topics, "
          f"a 189 human control (own topic) for the {len(Hall) - len(Hd)} whose first stanza is a reply")
    return [LABELS[x] for x, _ in A] + hl


# ---------------------------------------------------------------- score

def decide2(s):   # = fbj_raw.decide2, applied to the clean endpoint
    if s["diff"] >= 5 and s["ci"][0] > 0: return "CONTINUE"
    if s["ci"][1] < 5: return "STOP"
    return "GREY"


def labels(spec, root_pools="runs"):
    """= fixon_raw.labels over NPART parts per seed: {kind: {(src, topic, seed): bool}}; anchor clean share per CLEAN pool"""
    lab = {"clean": {}, "aligned": {}}; anchors = {}
    for s in seeds(spec):
        for q in range(NPART):
            for kind, inst in (("clean", "CLEAN"), ("aligned", "ALIGN")):
                d = f"{root_pools}/{PREFIX}_{inst}{s}{q}"
                km = json.load(open(f"{d}/keymap.json", encoding="utf-8"))
                res = {r["item"]: r for r in json.load(open(f"{d}/run1.json", encoding="utf-8"))}
                assert set(res) == set(km), f"{d}: judge output does not cover the pool"
                hum = []
                for k, v in km.items():
                    r = res[k]
                    val = (not any(e.get("severity") in ("critical", "major") for e in r["errors"])) if kind == "clean" else int(r["align"]) >= 4
                    if v["src"] == "human": hum.append(val); continue
                    t, ss = v["id"].rsplit("_s", 1); assert int(ss) == s, f"{d}: {v} is not seed {s}"
                    assert (v["src"], t, s) not in lab[kind], f"{v} judged twice"
                    lab[kind][(v["src"], t, s)] = val
                if kind == "clean": anchors[d] = sum(hum) / max(len(hum), 1)
    return lab, anchors


def ci_level(a, b, level, seed=0, n_boot=10000):
    """topic-clustered paired bootstrap CI of b - a (pp) at any level (= fixbase_pipe.paired_stats' resampling)"""
    import random as _r
    keys = sorted(set(a) & set(b)); by = collections.defaultdict(list)
    for k in keys: by[k[0]].append(k)
    T = sorted(by); rng = _r.Random(seed); boots = []
    for _ in range(n_boot):
        ks = [k for t in (rng.choice(T) for _ in T) for k in by[t]]
        boots.append(100 * (sum(b[k] for k in ks) - sum(a[k] for k in ks)) / len(ks))
    boots.sort(); lo = (1 - level) / 2
    return boots[int(lo * n_boot)], boots[int((1 - lo) * n_boot) - 1]


def score(out, spec, *arms, root_pools="runs"):
    import fixbase_pipe
    A = arm_specs(arms); form = {}
    for name, d in A:
        for s in seeds(spec):
            for r in rows(f"{d}/c{s}.jsonl"): form[(LABELS[name], r["id"], s)] = bool(r["all_pass"])
    lab, anchors = labels(spec, root_pools)
    void = {d: v for d, v in anchors.items() if v < 0.70}
    syslabs = [LABELS[n] for n, _ in A]
    for kind in ("clean", "aligned"):
        assert {k for k in lab[kind] if k[0] in syslabs} == set(form), f"{kind}: judged poems differ from the candidates"
    sysk = set(form)
    full = {k: {x: form[x] and lab[k][x] for x in sysk} for k in ("clean", "aligned")}
    full["all-three"] = {x: full["clean"][x] and full["aligned"][x] for x in sysk}
    full["form"] = {x: form[x] for x in sysk}; full["judged-clean"] = {x: lab["clean"][x] for x in sysk}
    full["judged-aligned"] = {x: lab["aligned"][x] for x in sysk}
    for x in lab["clean"]:      # human controls: no form step
        if x[0] in HUMANS:
            full["clean"][x] = full["judged-clean"][x] = lab["clean"][x]; full["aligned"][x] = lab["aligned"][x]
            full["all-three"][x] = lab["clean"][x] and lab["aligned"][x]
    R = {"anchor_clean_share": anchors, "void_pools": sorted(void), "contrasts": {}, "rates": {}, "decision": {}}
    sel = lambda L, src: {k[1:]: v for k, v in L.items() if k[0] == src}
    narms = len(A) - 1
    C_ = [(f"{l} vs ctl", "ctl", l, list(full)) for l in syslabs[1:]] + \
         [(f"{l} vs hum", HUM, l, ["clean", "aligned", "all-three"]) for l in syslabs]
    for name, a_, b_, kinds in C_:
        R["contrasts"][name] = {}; print(f"  {name}")
        for kind in kinds:
            x, y = sel(full[kind], a_), sel(full[kind], b_)
            if a_ == HUM:          # only the opening-post dev stanzas pair with a topic; the 189 substitutes do not
                x = {k: v for k, v in x.items() if k in y}; y = {k: v for k, v in y.items() if k in x}
            assert x and set(x) == set(y), f"{name} {kind}: unpaired"
            st = R["contrasts"][name][kind] = fixbase_pipe.paired_stats(x, y)
            extra = ""
            if a_ == "ctl" and kind in ("clean", "all-three") and narms > 1:     # reported, not used for the decision
                st["ci_bonferroni"] = ci_level(x, y, 1 - 0.05 / narms); extra = f"  Bonferroni({narms}) CI [{st['ci_bonferroni'][0]:+.1f}, {st['ci_bonferroni'][1]:+.1f}]"
            print(f"    {kind:14s} {a_} {st['a']:5.1f}%  {b_} {st['b']:5.1f}%  diff {st['diff']:+.1f}pp ({st['up']} up / {st['down']} down, "
                  f"n={st['n']})  95% CI [{st['ci'][0]:+.1f}, {st['ci'][1]:+.1f}]" + extra)
        if a_ == "ctl":
            c = R["contrasts"][name]
            R["decision"][b_] = decide2(c["clean"]) if c["form"]["b"] >= 99 else "STOP (form < 99%)"
            print(f"    DECISION {b_} (pre-registered, clean vs ctl): {R['decision'][b_]}")
    sysT = {k[1:] for k in sysk}
    for src in HUMANS:
        v = sel(full["clean"], src)
        if src == HUM:
            dv = {k: x for k, x in v.items() if k in sysT}; sv = {k: x for k, x in v.items() if k not in sysT}
            if dv: R["rates"]["hum (dev, opening post) clean"] = 100 * sum(dv.values()) / len(dv)
            if sv: R["rates"]["hum (189 substitutes) clean"] = 100 * sum(sv.values()) / len(sv)
        elif v: R["rates"][f"{src} (189 controls) clean"] = 100 * sum(v.values()) / len(v)
    print("  human clean (ONE judge run, wrapper v1, these pools): " + ", ".join(f"{k} {v:.1f}%" for k, v in R["rates"].items())
          + "  - a different instrument from 189's 3-run majority (88.9%); not directly comparable")
    R["decision_rule"] = ("per arm, primary clean (form & judged clean) vs ctl: CONTINUE iff diff >= +5 and lower > 0; STOP iff "
                          "upper < +5; else GREY; form >= 99% required; two arms, no multiplicity correction in the decision "
                          "(Bonferroni CIs reported)")
    os.makedirs(out, exist_ok=True); json.dump(R, open(f"{out}/RESULT.json", "w"), indent=1, default=str)
    print(f"  anchors clean: {anchors}")
    if void:
        R["decision"] = {k: f"provisional ({v}) - void pools, re-judge" for k, v in R["decision"].items()}
        json.dump(R, open(f"{out}/RESULT.json", "w"), indent=1, default=str)
        print(f"  VOID pools (anchors < 70% clean): {void} - decisions PROVISIONAL until re-judged"); sys.exit(3)
    print(f"  DECISIONS (pre-registered, per arm): {R['decision']}")


# ---------------------------------------------------------------- tests

class _FakeModel:
    """random logits that DEPEND ON THE CANVAS: a fixed table L [T, NT] plus 0.5 x L rolled along the vocabulary by a hash
    of the current tokens and rime channel, so every commit changes the next step's distribution; no weights"""
    def __init__(self, NT, T, seed=0):
        import torch
        g = torch.Generator().manual_seed(seed); self.L = torch.randn(T, NT, generator=g) * 3
    def __call__(self, tok, snd, ri):
        import torch
        h = int((tok[0].long() * torch.arange(1, tok.shape[1] + 1, device=tok.device)).sum() + 7 * ri[0].long().sum()) % 9973
        return (self.L + 0.5 * torch.roll(self.L, shifts=h, dims=1))[None], None


def _test():
    import tempfile, torch, contextlib, io
    import twolayer_canvas as TC, twolayer_full as TF
    from klon8 import prosody as P
    t = tempfile.mkdtemp(prefix="attend_", dir=os.environ.get("AE_TMP") or None)
    # 1. bigram set + end check: same newmm cut as bigram.py
    wrows(f"{t}/train.jsonl", [{"poem": "ฉันรักเธอมาก\nเธอรักฉันไหม"}, {"poem": "ฝนตกหนัก\nรอนานนาน"}])
    with contextlib.redirect_stdout(io.StringIO()): build_bigrams(f"{t}/train.jsonl", f"{t}/bg.json")
    BG = load_bigrams(f"{t}/bg.json")
    assert toks("ฉันรักเธอมาก") == ["ฉัน", "รัก", "เธอ", "มาก"] and toks("เธอรักฉันไหม") == ["เธอ", "รัก", "ฉัน", "ไหม"]
    assert {("ฉัน", "รัก"), ("รัก", "เธอ"), ("เธอ", "มาก"), ("ฉัน", "ไหม"), ("นาน", "นาน")} <= BG and ("เธอ", "ไหม") not in BG
    assert end_attested("ฉันรักเธอมาก", BG) and end_attested("เธอรักฉันไหม", BG)
    assert not end_attested("ฉันรักเธอไหม", BG), "(เธอ, ไหม) never occurs in the toy corpus"
    wrows(f"{t}/train_sp.jsonl", [{"poem": "ฝนตก หนัก\nรักเธอ"}])          # a space splits a pair: added from the stripped line
    with contextlib.redirect_stdout(io.StringIO()): build_bigrams(f"{t}/train_sp.jsonl", f"{t}/bg_sp.json")
    sp_ = load_bigrams(f"{t}/bg_sp.json"); st_ = toks("ฝนตกหนัก")
    assert (st_[-2], st_[-1]) in sp_ and json.load(open(f"{t}/bg_sp.json"))["n_raw_pairs"] <= len(sp_)
    assert end_attested("ฝน", BG), "one-word line is never unattested"
    px_ = proxy("ฉันรักเธอมาก\nตามใจตามใจตามใจ\nรักนานนาน\nนานแล้วหนานานว่านานกว่า", BG)
    assert px_["adj_repeat"] == 2, px_                                  # ตามใจตามใจ.. and นานนาน
    assert ("นาน", "นาน") in BG and px_["adj_unattested"] == 1, px_     # นาน|นาน is attested reduplication in the toy text
    assert px_["tri_nonadj"] == 1, px_                                  # นาน x3, never adjacent
    d = json.load(open(f"{t}/bg.json")); assert d["source_sha256"] == sha(f"{t}/train.jsonl") and d["n_lines"] == 4
    # 2. word repeats in grids
    g = [[["ตามใจ", 0, 2, "taam0"], ["ตามใจ", 1, 2, "caj0"], ["ตามใจ", 0, 2, "taam0"], ["ตามใจ", 1, 2, "caj0"], ["ดี", 0, 1, "dii0"]],
         [["ตามใจ", 0, 2, "taam0"], ["ตามใจ", 1, 2, "caj0"], ["ตามใจ", 0, 2, "taam0"], ["ตามใจ", 1, 2, "caj0"]],
         [["ขวัญ", 0, 1, "khwan4"], ["ขวัญ", 0, 1, "khwan4"], ["ขวัญ", 0, 1, "khwan4"], ["นานา", 0, 2, "naa0"], ["นานา", 1, 2, "naa0"]]]
    p_, t_ = word_repeats(g)
    assert p_ == {"ตามใจ": 2, "ขวัญ": 2} and t_ == {"ขวัญ": 1}, (p_, t_)
    os.makedirs(f"{t}/d"); wrows(f"{t}/d/poems.jsonl", [{"grid": g}])
    with contextlib.redirect_stdout(io.StringIO()): build_wordrep(f"{t}/d", f"{t}/wrep.json", 2)
    w_ = json.load(open(f"{t}/wrep.json")); assert w_["pairs"] == {"ตามใจ": 2, "ขวัญ": 2} and w_["triples"] == {}
    # 3. whole-word stutter ban on a toy lexicon (no model)
    lex = {"words": ["ตามใจ", "ฟาดหาง", "ขวัญ", "ดี", "นาน", "เลย"],
           "sounds": {"ตามใจ": ["taam0", "caj0"], "ฟาดหาง": ["faat2", "haaN4"], "ขวัญ": ["khwan4"], "ดี": ["dii0"],
                      "นาน": ["naan0"], "เลย": ["l@@j0"]}, "count": {}}
    X = TC.Lex(lex); AD = make_decoder_class()
    json.dump({"pairs": {"ขวัญ|ขวัญ": 2}}, open(f"{t}/allow.json", "w"))
    AD.ALLOW = TF.load_allow(f"{t}/allow.json", X)
    json.dump({"pairs": {"ตามใจ": 2, "ขวัญ": 2}, "triples": {}}, open(f"{t}/wr.json", "w"), ensure_ascii=False)
    AD.WREP = load_wordrep(f"{t}/wr.json", X)
    D = AD.__new__(AD); D.X = X; D.lens = [8] * TC.L; D.committed = set(); D.forced = {}
    D.tok = torch.full((TC.T,), TC.MASK, dtype=torch.long); D.n_banned = 0; D.n_word_banned = 0
    def put(i, j, w, k=0): kk = D.idx(i, j); D.tok[kk] = X.tok(w, k); D.committed.add(kk)
    put(0, 0, "ฟาดหาง", 0); put(0, 1, "ฟาดหาง", 1)                 # ฟาดหาง at 0-1 (repeat not attested)
    put(1, 0, "ตามใจ", 0); put(1, 1, "ตามใจ", 1)                   # ตามใจ at 0-1 (pair attested)
    put(2, 0, "ขวัญ"); put(2, 1, "ขวัญ")                            # ขวัญขวัญ (attested pair; triple not)
    put(3, 3, "ฟาดหาง", 0); D.forced[D.idx(3, 4)] = X.tok("ฟาดหาง", 1)   # right side: a word starting at (3,3)
    lp = torch.zeros(TC.T, X.NT); D.ban_word_stutters(lp); D.ban_stutters(lp)
    ok = lambda i, j, w, k=0: bool(lp[D.idx(i, j), X.tok(w, k)] > -1e8)
    assert not ok(0, 2, "ฟาดหาง", 0), "unattested whole-word repeat ฟาดหางฟาดหาง must be banned"
    assert ok(0, 2, "ตามใจ", 0) and ok(0, 2, "ดี")
    assert ok(1, 2, "ตามใจ", 0), "attested ตามใจ|ตามใจ allowed"
    assert not ok(2, 2, "ขวัญ"), "third ขวัญ in a row (triple unattested) banned"
    assert ok(2, 2, "ดี")
    assert not ok(3, 2, "ฟาดหาง", 1), "the last syllable of ฟาดหาง right before ฟาดหาง is banned"
    assert ok(3, 2, "ดี") and ok(3, 2, "ตามใจ", 1)
    put(1, 4, "ตามใจ", 0); put(1, 5, "ตามใจ", 1)                   # gap of 2 cells between two ตามใจ: middle triple
    lp = torch.zeros(TC.T, X.NT); D.ban_word_stutters(lp)
    assert not ok(1, 2, "ตามใจ", 0) and not ok(1, 3, "ตามใจ", 1), "ตามใจ between two ตามใจ makes a triple"
    # 4. full decodes on the REAL lexicon with a canvas-dependent fake model (no weights)
    lexf = "runs/TWOLAYER/data/lexicon.json"
    if os.path.exists(lexf):
        XR = TC.Lex(lexf); M = _FakeModel(XR.NT, TC.T)
        SD0 = TF.make_decoder_class(); SD0.ALLOW = ({}, {})
        AD0 = make_decoder_class(); AD0.ALLOW = ({}, {})
        topics = ["ฝนตก", "คิดถึงบ้าน", "รักแรก"]
        base = {}
        for n, top in enumerate(topics):                                   # (a) flags off: AttendDecoder == StutterDecoder
            outs = []
            for cls in (SD0, AD0):
                TC.Decoder = cls
                lines, Dd, tries = TC.decode_one(M, XR, top, unit="syllable", lens=[8, 9, 8, 8], gen=torch.Generator().manual_seed(n))
                outs.append((lines, Dd.n_fwd, tries, Dd.n_banned))
            assert outs[0] == outs[1], f"flags off: AttendDecoder differs from StutterDecoder: {outs}"
            base[top] = outs[0]
        # (b) forced re-draws: a per-LINE test double - each line's FIRST end check says 'unattested', later ones 'attested'
        class AD1(make_decoder_class()):
            ALLOW = ({}, {}); END_TRIES = 3; BIGRAMS = frozenset()
            def end_ok(self, i):
                self.ae_checks += 1; self.seen_line = getattr(self, "seen_line", collections.Counter()); self.seen_line[i] += 1
                return self.seen_line[i] > 1
        for n, top in enumerate(topics):
            TC.Decoder = AD1
            lines, D1, tries = TC.decode_one(M, XR, top, unit="syllable", lens=[8, 9, 8, 8], gen=torch.Generator().manual_seed(n))
            assert D1.ae_unattested == 4 and D1.ae_fixed == 4 and D1.ae_failed == 0 and D1.ae_tries == 4, \
                (D1.ae_unattested, D1.ae_fixed, D1.ae_failed, D1.ae_tries)
            assert dict(D1.seen_line) == {0: 2, 1: 2, 2: 2, 3: 2}, D1.seen_line
            for i in range(TC.L):
                last = D1.idx(i, D1.lens[i] - 1); tl = int(D1.tok[last])
                assert len(D1.end_ban[last]) == 1 and tl not in D1.end_ban[last], "the banned (original) end word came back"
                assert bool(XR.tok_last[tl]), "line-final cell must hold a word's last syllable"
                if i in TC.TONE_AT: assert int(XR.tok_to[tl]) in TC.TONE_AT[i], "tone class at the line end"
            for gi in range(len(TC.GROUPS)):                                   # every rhyme cell carries its group's family
                for c in D1.gcells(gi):
                    assert int(XR.tok_ri[int(D1.tok[D1.idx(*c)])]) == D1.fam[gi], f"group {gi} cell {c}: rhyme lost"
            want = [[XR.phons[int(XR.tok_sound[int(D1.tok[D1.idx(i, j)])])] for j in range(D1.lens[i])] for i in range(TC.L)]
            got = [[s_.phon for s_ in P.g2p_line(l)] for l in lines]
            assert got == want, f"prosody readback differs from the cells after re-draws ({top}, {tries} retries)"
            sc = P.score("\n".join(lines), n_bot_expected=1)
            assert sc["ext_rhyme"] == 1.0, (lines, sc)                        # the checker agrees: the rhyme still holds
            assert lines != base[top][0], "re-drawn ends must change the poem"
            print(f"  forced re-draw, topic {top}: 4 lines re-drawn, readback == cells, ext_rhyme 1.0")
        # (c) never attested -> K tries per line, every one restored: byte-identical to the flags-off decode (RNG included)
        AD2 = make_decoder_class(); AD2.ALLOW = ({}, {}); AD2.END_TRIES = 2; AD2.BIGRAMS = frozenset()
        for n, top in enumerate(topics):
            TC.Decoder = AD2
            lines, D2, tries = TC.decode_one(M, XR, top, unit="syllable", lens=[8, 9, 8, 8], gen=torch.Generator().manual_seed(n))
            assert (lines, D2.n_fwd, tries, D2.n_banned) == base[top], f"all-fail E differs from the control decode ({top})"
            assert D2.ae_failed == 4 and D2.ae_fixed == 0 and D2.ae_tries >= 4 and D2.ae_fwd > 0, (D2.ae_failed, D2.ae_tries, D2.ae_fwd)
        print(f"  real-lexicon fake-model decodes ({len(topics)} topics): flags-off identity, forced re-draws keep rhyme/tone, "
              f"all-fail == control byte for byte")
    else: print("  (real lexicon missing: decode tests skipped)")
    # 4b. the REAL model, flags off through AttendDecoder, == the control's c0 (env AE_TEST_MODEL=RUN AE_TEST_REF=c0.jsonl)
    if os.environ.get("AE_TEST_MODEL"):
        ref = os.environ["AE_TEST_REF"]; dev = os.environ.get("AE_TEST_PROMPTS", "data/prompts/dev49_pilot.jsonl")
        with contextlib.redirect_stderr(io.StringIO()):
            decode_cli(["--model", os.environ["AE_TEST_MODEL"], "--prompts", dev, "--unit", "syllable", "--seed", "0", "--limit", "3",
                        "--stutter-fix", "--allow", os.environ["AE_TEST_ALLOW"], "--out", f"{t}/reg3.jsonl"])
        assert TC.Decoder.__name__ == "AttendDecoder", "the regression must run through AttendDecoder"
        assert open(f"{t}/reg3.jsonl", "rb").read() == b"".join(open(ref, "rb").readlines()[:3]), "flags-off AttendDecoder != control c0"
        print("  real model: flags-off AttendDecoder decode == the control's c0 on 3 topics")
    # 5. gates / decision rules
    c = {"form": 99.6, "stutters_unattested": 0.05, "final_bits": 5.0, "tri_nonadj": 0.10, "distinct2": 0.90}
    ok_ = {"form": 99.0, "stutters_unattested": 0.15, "final_bits": 4.7, "tri_nonadj": 0.15, "distinct2": 0.87}
    assert gate_fail(c, ok_) == []                                       # every guard exactly at its limit passes
    assert gate_fail(c, dict(ok_, form=98.9, stutters_unattested=0.16)) == ["form < 99%", "unattested stutters"]
    assert gate_fail(c, dict(ok_, final_bits=4.69)) == ["final-word bits"]
    assert gate_fail(c, dict(ok_, tri_nonadj=0.151)) == ["word echo (3x in a line)"]
    assert gate_fail(c, dict(ok_, distinct2=0.869)) == ["distinct-2"]
    assert topic_echo("รักแรก", "รักแรกพบ\nรักเธอ\nฝน\nตก") == sum(w in set(toks("รักแรก")) for w in toks("รักแรกพบ") + toks("รักเธอ"))
    D_ = lambda d, lo, hi: {"diff": d, "ci": (lo, hi)}
    assert decide2(D_(6, 1, 11)) == "CONTINUE" and decide2(D_(6, -1, 12)) == "GREY" and decide2(D_(0, -5, 4.9)) == "STOP"
    assert decide2(D_(5, 0, 10)) == "GREY"
    assert [len(x) for x in parts([f"T{i:02d}" for i in range(49)])] == [13, 12, 12, 12]
    try: arm_specs(["E=a", "CTL=b"]); raise AssertionError("control not first accepted")
    except AssertionError as e: assert "CTL" in str(e), e
    # 6. pools + score on fake labels (real build_pool; 9 topics x 2 seeds -> parts 3/2/2/2), both arms, 1:1 human
    import build_pool, random as _r
    NTp = 9
    for dd, badf in (("ctl", 3), ("E", 1), ("S", 5)):
        os.makedirs(f"{t}/{dd}")
        for s in (0, 1):
            wrows(f"{t}/{dd}/c{s}.jsonl", [{"id": f"T{i:02d}", "prompt": f"p{i}", "unit": "syllable", "all_pass": i != badf,
                                          "poem": f"ก{dd}{i}{s}\nข\nค\nง"} for i in range(NTp)])
    REPLY = {2, 5, 7}                                    # dev topics whose human stanza is a reply (tier 3)
    wrows(f"{t}/hum.jsonl", [{"id": f"T{i:02d}", "prompt": f"p{i}", "poem": f"ม{i}\nข\nค\nง", "tier": 3 if i in REPLY else 1 + i % 2}
                             for i in range(NTp)])
    wrows(f"{t}/c2.jsonl", [{"id": f"G{i:03d}", "prompt": f"g{i}", "poem": f"ช{i}\nข\nค\nง"} for i in range(40)])
    A6 = [f"CTL={t}/ctl", f"E={t}/E", f"S={t}/S"]
    with contextlib.redirect_stdout(io.StringIO()): labs = pools(f"{t}/p", "0-1", f"{t}/hum.jsonl", f"{t}/c2.jsonl", *A6)
    assert labs == ["ctl", "ate", "ats", HUM, HUMB, HUMC], labs
    for s in (0, 1):
        for lab in (HUMB, HUMC):
            ids = [x["id"] for q in range(NPART) for x in rows(f"{t}/p/{lab}_s{s}q{q}.jsonl")]
            assert len(ids) == len(set(ids)) == NTp, ids
        hum_ = [x for q in range(NPART) for x in rows(f"{t}/p/{HUM}_s{s}q{q}.jsonl")]
        assert [x["id"].startswith("G") for x in hum_] == [i in REPLY for i in range(NTp)], "reply stanzas must be substituted"
        assert all(x["prompt"] == f"p{i}" for i, x in enumerate(hum_) if i not in REPLY), "dev stanza under its own topic"
        assert all(x["prompt"].startswith("g") for i, x in enumerate(hum_) if i in REPLY), "substitute keeps its 189 topic"
        for q, nsys in enumerate((3, 2, 2, 2)):
            g_ = [x["id"] for lab in (HUM, HUMB, HUMC) for x in rows(f"{t}/p/{lab}_s{s}q{q}.jsonl") if x["id"].startswith("G")]
            assert len(g_) == len(set(g_)), "a 189 human poem twice in one pool"
            for inst, al in (("CLEAN", False), ("ALIGN", True)):
                with contextlib.redirect_stdout(io.StringIO()):
                    km = build_pool.build(f"{t}/runs/{PREFIX}_{inst}{s}{q}", [(l, f"{t}/p/{l}_s{s}q{q}.jsonl") for l in labs],
                                          nlines=4, align=al, no_drop=True)
                comp = collections.Counter(v["src"] for v in km.values())
                assert comp == {**{l: nsys for l in labs}, "human": -(-6 * nsys // 5)}, comp
    # one arm killed -> 2 system + 2 human labels
    with contextlib.redirect_stdout(io.StringIO()): labs2 = pools(f"{t}/p2", "0-0", f"{t}/hum.jsonl", f"{t}/c2.jsonl", f"CTL={t}/ctl", f"S={t}/S")
    assert labs2 == ["ctl", "ats", HUM, HUMB], labs2
    rng = _r.Random(0)
    for s in (0, 1):
        for q in range(NPART):
            for inst in ("CLEAN", "ALIGN"):
                dd = f"{t}/runs/{PREFIX}_{inst}{s}{q}"; km = json.load(open(f"{dd}/keymap.json")); res = []
                for k, v in km.items():
                    bad = (inst == "CLEAN" and v["src"] == "ctl" and v["id"].startswith("T00")) or \
                          (inst == "CLEAN" and v["src"] == HUM and v["id"].startswith("T04"))
                    res.append({"item": k, "errors": [{"severity": "major"}] if bad else [], "align": rng.choice([4, 5])})
                json.dump(res, open(f"{dd}/run1.json", "w"))
    with contextlib.redirect_stdout(io.StringIO()): score(f"{t}/o", "0-1", *A6, root_pools=f"{t}/runs")
    Rr = json.load(open(f"{t}/o/RESULT.json")); cE = Rr["contrasts"]["ate vs ctl"]["clean"]; cS = Rr["contrasts"]["ats vs ctl"]["clean"]
    # ctl: T00 judged bad, T03 form-fail -> 7/9; E: T01 form-fail -> 8/9; S: T05 form-fail -> 8/9 per seed
    assert cE["n"] == 18 and abs(cE["a"] - 700 / 9) < 1e-9 and abs(cE["b"] - 800 / 9) < 1e-9 and cE["up"] == 4 and cE["down"] == 2, cE
    assert cS["up"] == 4 and cS["down"] == 2 and "ci_bonferroni" in cS and cS["ci_bonferroni"][0] <= cS["ci"][0]
    assert set(Rr["decision"]) == {"ate", "ats"}
    # vs hum: only the 6 opening-post topics pair (12 items); there ctl = 4/6 (T00 bad, T03 form), hum = 5/6 (T04 bad)
    ch = Rr["contrasts"]["ctl vs hum"]["clean"]
    assert ch["n"] == 12 and abs(ch["a"] - 500 / 6) < 1e-9 and abs(ch["b"] - 400 / 6) < 1e-9, ch
    assert abs(Rr["rates"]["hum (dev, opening post) clean"] - 500 / 6) < 1e-9 and Rr["rates"]["hum (189 substitutes) clean"] == 100
    assert Rr["rates"]["humb (189 controls) clean"] == Rr["rates"]["humc (189 controls) clean"] == 100
    assert not any(str(v).startswith("provisional") for v in Rr["decision"].values())
    # a void pool (anchors dirty) -> exit 3 and every decision marked provisional
    dd = f"{t}/runs/{PREFIX}_CLEAN00"; km = json.load(open(f"{dd}/keymap.json")); res = json.load(open(f"{dd}/run1.json"))
    for r in res:
        if km[r["item"]]["src"] == "human": r["errors"] = [{"severity": "major"}]
    json.dump(res, open(f"{dd}/run1.json", "w"))
    try:
        with contextlib.redirect_stdout(io.StringIO()): score(f"{t}/o2", "0-1", *A6, root_pools=f"{t}/runs")
        raise AssertionError("void pool not caught")
    except SystemExit as e: assert e.code == 3
    assert all(v.startswith("provisional") for v in json.load(open(f"{t}/o2/RESULT.json"))["decision"].values())
    print(f"twolayer_attend: tests pass (temp {t})")


if __name__ == "__main__":
    cmd, args = sys.argv[1], sys.argv[2:]
    if cmd == "decode": decode_cli(args)
    elif cmd == "bigrams": build_bigrams(*args)
    elif cmd == "wordrep": build_wordrep(args[0], args[1], *(args[2:3] or [2]))
    elif cmd == "human": build_human(*args)
    elif cmd == "gates": gates(*args)
    elif cmd == "pools": print("LABELS " + " ".join(pools(*args)))
    elif cmd == "score": score(*args)
    elif cmd == "test": _test()
    else: sys.exit(f"unknown command {cmd}")