code/tests/test_plan_dataset.py

# test_plan_dataset.py: The three think-block tiers in the dataset.

import json
import os
import subprocess
import sys
import zlib
from pathlib import Path

import pytest

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

torch = pytest.importorskip("torch")


class FakeTokenizer:
    UNK_ID = 0
    UNK_TOKEN = "<|endoftext|>"
    TEXT_BASE = 1000

    def __init__(self, tokens, base_vocab: int = 50277):
        self._token_to_id = {self.UNK_TOKEN: self.UNK_ID}
        self._id_to_token = {self.UNK_ID: self.UNK_TOKEN}
        for offset, token in enumerate(tokens):
            token_id = base_vocab + offset
            self._token_to_id[token] = token_id
            self._id_to_token[token_id] = token

    def __len__(self):
        return max(self._id_to_token) + 1

    def convert_tokens_to_ids(self, tokens):
        if isinstance(tokens, str):
            return self._token_to_id.get(tokens, self.UNK_ID)
        return [self._token_to_id.get(t, self.UNK_ID) for t in tokens]

    def convert_ids_to_tokens(self, ids):
        if isinstance(ids, int):
            return self._id_to_token.get(ids)
        return [self._id_to_token.get(i) for i in ids]

    def encode(self, text, add_special_tokens=False):
        return [self.TEXT_BASE + (zlib.crc32(w.encode()) % 4096)
                for w in text.split()]


def fake_registry():
    import tokenizer as tok
    return tok.TokenRegistry(FakeTokenizer(tok.get_all_new_tokens()))

PROMPT = "alpha beta gamma delta"
REASONING = "Say it warmly and slow down at the end."

PLAN_WORDS = [
    {"word": "alpha", "pitch_bin": 3, "dur_bin": 1, "energy_bin": 0, "dur_frames": 2},
    {"word": "beta", "pitch_bin": None, "dur_bin": 4, "energy_bin": 7, "dur_frames": 4},
    {"word": "gamma", "pitch_bin": 15, "dur_bin": 9, "energy_bin": 3, "dur_frames": 6},
]
SUM_D = 12


def _row(n_frames, *, plan=None, reasoning=REASONING, tags=True, prompt=PROMPT):
    from delay_dataset import NUM_LEVELS
    row = {
        "prompt": prompt,
        "reasoning": reasoning,
        "speaker_id": 3,
        "speech_tokens": list(range(n_frames)),
        "residual_codes": [[(k + 1) % 900 for _ in range(n_frames)]
                           for k in range(NUM_LEVELS - 1)],
    }
    if tags:
        row.update({"emotion_label": "joyful", "pace": "fast", "pitch": "high"})
    if plan is not None:
        row["plan"] = plan
    return row


def write_corpus(path, rows):
    Path(path).write_text("\n".join(json.dumps(r) for r in rows), encoding="utf-8")
    return str(path)


def make_dataset(path, registry, cot_mode, **kw):
    from delay_dataset import DelayMimiDataset
    import tokenizer as tok
    kw.setdefault("max_seq_len", 2048)
    kw.setdefault("aligned", True)
    kw.setdefault("stage", "auto")
    kw.setdefault("lazy", False)
    return DelayMimiDataset(data_path=path, tokenizer=registry.tokenizer,
                            registry=registry, cot_mode=cot_mode, **kw)


def expected_legacy_ids(reg, tok_, think_text, n_frames, aligned=True, n_levels=None):
    from delay_dataset import NUM_LEVELS
    n_levels = NUM_LEVELS if n_levels is None else n_levels
    stagger = 0 if aligned else n_levels - 1
    span = n_frames + stagger
    think = ([reg.think_start_id] + tok_.encode(think_text) + [reg.think_end_id]
             if think_text else [])
    return ([reg.bos_id, reg.user_prompt_id] + tok_.encode(PROMPT) + think
            + [reg.audio_start_id] + [reg.audio_start_id] * span
            + [reg.audio_end_id, reg.eos_id])


def probe(tmpdir) -> dict:
    import config
    from delay_dataset import DelayMimiDataset, NO_WORD

    tmpdir = Path(tmpdir)
    reg = fake_registry()
    tok_ = reg.tokenizer

    rows = [
        _row(SUM_D, plan=list(PLAN_WORDS)),
        _row(SUM_D),
        _row(2 * SUM_D, plan=list(PLAN_WORDS)),
    ]
    path = write_corpus(tmpdir / "plan.jsonl", rows)

    out = {
        "plan_start_id": reg.plan_start_id,
        "plan_end_id": reg.plan_end_id,
        "plan_word_sep_id": reg.plan_word_sep_id,
        "think_start_id": reg.think_start_id,
        "think_end_id": reg.think_end_id,
        "audio_start_id": reg.audio_start_id,
        "audio_end_id": reg.audio_end_id,
        "bos_id": reg.bos_id,
        "eos_id": reg.eos_id,
        "pitch_ids": [reg.pitch_bin_to_id(b) for b in range(config.PITCH_BINS)],
        "duration_ids": [reg.duration_bin_to_id(b) for b in range(config.DURATION_BINS)],
        "energy_ids": [reg.energy_bin_to_id(b) for b in range(config.ENERGY_BINS)],
        "prompt_ids": tok_.encode(PROMPT),
        "reasoning_ids": tok_.encode(REASONING),
        "tag_ids": tok_.encode("EMO=joyful PACE=fast PITCH=high"),
        "tag_ids_after_prose": tok_.encode(" EMO=joyful PACE=fast PITCH=high"),
        "no_word": NO_WORD,
    }

    def dump(ds, i):
        item = ds[i]
        d = {k: (v.tolist() if torch.is_tensor(v) else v) for k, v in item.items()}
        return d

    for mode in ("prose", "tags", "none", "plan", "tags+plan", "full"):
        ds = make_dataset(path, reg, mode)
        out[f"items_{mode}"] = [dump(ds, i) for i in range(len(ds))]
        out[f"counters_{mode}"] = {
            "seen": ds.n_rows_seen, "missing": ds.n_rows_missing_plan,
            "rejected": ds.n_rows_plan_rejected,
            "truncated": ds.n_rows_plan_truncated,
            "prose_dropped": ds.n_rows_reasoning_dropped,
        }
        out[f"audit_{mode}"] = ds.plan_audit

    ds_pack = make_dataset(path, reg, "plan", plan_schedule_mode="pack")
    out["items_plan_pack"] = [dump(ds_pack, i) for i in range(len(ds_pack))]

    ds_la0 = make_dataset(path, reg, "plan", plan_lookahead=0)
    out["item_plan_lookahead0"] = dump(ds_la0, 0)

    ds_w = make_dataset(path, reg, "full", reasoning_weight=0.3, tag_weight=1.0,
                        plan_weight=2.0)
    out["item_full_weighted"] = dump(ds_w, 0)

    from delay_dataset import DelayCollator
    ds_plan = make_dataset(path, reg, "plan")
    batch = DelayCollator(pad_token_id=reg.pad_id)([ds_plan[0], ds_plan[1], ds_plan[2]])
    out["batch_keys"] = sorted(batch.keys())
    out["batch"] = {k: batch[k].tolist() for k in
                    ("plan_schedule", "plan_slots", "plan_durations", "plan_lengths")}
    ds_none = make_dataset(path, reg, "none")
    out["batch_keys_plainarm"] = sorted(
        DelayCollator(pad_token_id=reg.pad_id)([ds_none[0], ds_none[1]]).keys())

    tight = make_dataset(path, reg, "full", max_seq_len=40)
    out["item_tight"] = dump(tight, 0)
    out["counters_tight"] = {"prose_dropped": tight.n_rows_reasoning_dropped,
                             "truncated": tight.n_rows_plan_truncated}
    tighter = make_dataset(path, reg, "full", max_seq_len=20)
    out["item_tighter"] = dump(tighter, 0)
    out["counters_tighter"] = {"truncated": tighter.n_rows_plan_truncated}

    rej = write_corpus(tmpdir / "rej.jsonl", [
        _row(SUM_D, plan={"words": list(PLAN_WORDS), "ok": False}),
        _row(SUM_D, plan=list(PLAN_WORDS)),
    ])
    ds_rej = make_dataset(rej, reg, "plan")
    out["rejected_audit"] = ds_rej.plan_audit
    out["rejected_item0_schedule"] = dump(ds_rej, 0)["plan_schedule"]
    out["rejected_counters"] = {"rejected": ds_rej.n_rows_plan_rejected,
                                "missing": ds_rej.n_rows_missing_plan}

    noplan = write_corpus(tmpdir / "noplan.jsonl", [_row(SUM_D), _row(SUM_D)])
    out["err_zero_coverage"] = _err(make_dataset, noplan, reg, "plan")
    out["err_stage1_plan_arm"] = _err(lambda: make_dataset(path, reg, "plan", stage=1))
    out["err_min_coverage"] = _err(
        lambda: make_dataset(path, reg, "plan", min_plan_coverage=0.95))

    bad_bin = write_corpus(tmpdir / "badbin.jsonl", [
        _row(SUM_D, plan=[{"pitch_bin": 3, "dur_bin": 99, "energy_bin": 0,
                           "dur_frames": 2}])])
    out["err_bad_dur_bin"] = _err(lambda: make_dataset(bad_bin, reg, "plan")[0])
    no_frames = write_corpus(tmpdir / "nodur.jsonl", [
        _row(SUM_D, plan=[{"pitch_bin": 3, "dur_bin": 1, "energy_bin": 0}])])
    out["err_no_dur_frames"] = _err(lambda: make_dataset(no_frames, reg, "plan")[0])
    out["err_no_dur_frames_msg"] = _msg(lambda: make_dataset(no_frames, reg, "plan")[0])
    dur_s = write_corpus(tmpdir / "durs.jsonl", [
        _row(SUM_D, plan=[{"pitch_bin": 3, "dur_bin": 1, "energy_bin": 0,
                           "dur_s": 0.5}])])
    out["dur_s_frames"] = dump(make_dataset(dur_s, reg, "plan"), 0)["plan_durations"]

    out["legacy_expected_prose"] = expected_legacy_ids(
        reg, tok_, REASONING, SUM_D)
    out["legacy_expected_tags"] = expected_legacy_ids(
        reg, tok_, "EMO=joyful PACE=fast PITCH=high", SUM_D)
    out["legacy_expected_none"] = expected_legacy_ids(reg, tok_, "", SUM_D)
    return out


def _err(fn, *args):
    try:
        fn(*args)
    except Exception as exc:
        return type(exc).__name__
    return None


def _msg(fn, *args):
    try:
        fn(*args)
    except Exception as exc:
        return str(exc)
    return ""

_CACHE: dict = {}


def _run_probe(**env_extra) -> dict:
    key = tuple(sorted(env_extra.items()))
    if key in _CACHE:
        return _CACHE[key]
    env = dict(os.environ)
    env.update(env_extra)
    env["PYTHONPATH"] = str(REPO_ROOT) + os.pathsep + env.get("PYTHONPATH", "")
    proc = subprocess.run(
        [sys.executable, str(Path(__file__).resolve())],
        capture_output=True, text=True, cwd=str(REPO_ROOT), env=env)
    assert proc.returncode == 0, (
        f"probe subprocess failed (env={env_extra}):\n"
        f"{proc.stdout[-2000:]}\n{proc.stderr[-4000:]}")
    result = json.loads(proc.stdout.strip().splitlines()[-1])
    _CACHE[key] = result
    return result


@pytest.fixture(scope="module")
def on() -> dict:
    return _run_probe(MVC_ENABLE_PLAN_TOKENS="1", MVC_NUM_SPEECH_TOKENS="64")


@pytest.fixture(scope="module")
def reg():
    import config
    if config.ENABLE_PLAN_TOKENS:
        pytest.skip("suite is running with MVC_ENABLE_PLAN_TOKENS=1; gate-off N/A")
    return fake_registry()


@pytest.fixture
def corpus(tmp_path):
    return write_corpus(tmp_path / "legacy.jsonl",
                        [_row(SUM_D), _row(SUM_D, reasoning="")])


class TestLegacyUnchanged:
    @pytest.mark.parametrize("mode,think_text", [
        ("prose", REASONING),
        ("tags", "EMO=joyful PACE=fast PITCH=high"),
        ("none", ""),
    ])
    def test_sequence_is_the_hand_written_answer(self, corpus, reg, mode, think_text):
        ds = make_dataset(corpus, reg, mode)
        got = ds[0]["input_ids"].tolist()
        assert got == expected_legacy_ids(reg, reg.tokenizer, think_text, SUM_D)

    def test_none_emits_no_think_block_at_all(self, corpus, reg):
        ids = make_dataset(corpus, reg, "none")[0]["input_ids"].tolist()
        assert reg.think_start_id not in ids
        assert reg.think_end_id not in ids

    def test_stage1_row_in_auto_mode_still_gets_no_think(self, corpus, reg):
        ids = make_dataset(corpus, reg, "prose")[1]["input_ids"].tolist()
        assert reg.think_start_id not in ids

    def test_legacy_loss_weights_unchanged(self, corpus, reg):
        item = make_dataset(corpus, reg, "prose", reasoning_weight=0.3)[0]
        w = item["loss_weights"].tolist()
        n_prompt = len(reg.tokenizer.encode(PROMPT))
        lo = 2 + n_prompt
        hi = lo + 1 + len(reg.tokenizer.encode(REASONING)) + 1
        assert w[lo:hi] == pytest.approx([0.3] * (hi - lo))
        assert set(w[:lo]) == {1.0} and set(w[hi:]) == {1.0}

    def test_no_plan_keys_leak_into_a_legacy_batch(self, corpus, reg):
        from delay_dataset import DelayCollator
        ds = make_dataset(corpus, reg, "prose")
        batch = DelayCollator(pad_token_id=reg.pad_id)([ds[0], ds[1]])
        assert not [k for k in batch if k.startswith("plan_")]
        assert set(batch) == {"input_ids", "labels", "codes", "attention_mask",
                              "loss_weights", "speaker_ids"}

    def test_item_exposes_plan_fields_as_none(self, corpus, reg):
        item = make_dataset(corpus, reg, "prose")[0]
        assert item["plan_schedule"] is None
        assert item["plan_slots"] is None
        assert item["plan_durations"] is None


class TestGateOffRefusesPlanArms:
    @pytest.mark.parametrize("mode", ["plan", "tags+plan", "full"])
    def test_plan_mode_without_the_vocabulary_is_fatal(self, corpus, reg, mode):
        with pytest.raises(RuntimeError, match="MVC_ENABLE_PLAN_TOKENS"):
            make_dataset(corpus, reg, mode)

    def test_unknown_cot_mode_still_rejected(self, corpus, reg):
        with pytest.raises(AssertionError):
            make_dataset(corpus, reg, "plann")


class TestSchedule:
    def test_stretch_tiles_the_span_exactly(self):
        from delay_dataset import build_plan_schedule
        assert build_plan_schedule([2, 4, 6], 12) == (
            [0] * 2 + [1] * 4 + [2] * 6)

    def test_stretch_dilates_proportionally(self):
        from delay_dataset import build_plan_schedule
        assert build_plan_schedule([2, 4, 6], 24) == (
            [0] * 4 + [1] * 8 + [2] * 12)

    def test_stretch_covers_every_frame_and_every_word(self):
        from delay_dataset import build_plan_schedule
        sched = build_plan_schedule([3, 1, 7, 2], 37)
        assert -1 not in sched and len(sched) == 37
        assert sorted(set(sched)) == [0, 1, 2, 3]
        assert sched == sorted(sched)

    def test_pack_is_reproducible_from_the_plan_alone(self):
        from delay_dataset import build_plan_schedule, NO_WORD
        assert build_plan_schedule([2, 4, 6], 24, "pack") == (
            [0] * 2 + [1] * 4 + [2] * 6 + [NO_WORD] * 12)

    def test_pack_clips_at_the_span_without_wrapping(self):
        from delay_dataset import build_plan_schedule
        assert build_plan_schedule([2, 4, 6], 5, "pack") == [0, 0, 1, 1, 1]

    def test_more_words_than_frames_does_not_crash(self):
        from delay_dataset import build_plan_schedule
        sched = build_plan_schedule([1, 1, 1, 1], 2)
        assert len(sched) == 2 and sched == sorted(sched)

    def test_empty_and_zero_cases(self):
        from delay_dataset import build_plan_schedule, NO_WORD
        assert build_plan_schedule([], 3) == [NO_WORD] * 3
        assert build_plan_schedule([1, 2], 0) == []

    def test_degenerate_durations_raise(self):
        from delay_dataset import build_plan_schedule
        with pytest.raises(ValueError):
            build_plan_schedule([0, 0], 10)
        with pytest.raises(ValueError):
            build_plan_schedule([1, -2], 10)

    def test_unknown_mode_raises(self):
        from delay_dataset import build_plan_schedule
        with pytest.raises(ValueError, match="schedule mode"):
            build_plan_schedule([1], 1, "squash")


class TestWordNormalisation:
    def _norm(self, words):
        from delay_dataset import normalise_plan_words
        return normalise_plan_words(words, n_pitch_bins=16, n_duration_bins=16,
                                    n_energy_bins=8)

    def test_unvoiced_sentinel_becomes_none(self):
        out = self._norm([{"pitch_bin": 16, "dur_bin": 0, "energy_bin": 0,
                           "dur_frames": 3}])
        assert out[0]["pitch_bin"] is None

    def test_out_of_range_bins_raise_with_coordinates(self):
        from delay_dataset import normalise_plan_words
        with pytest.raises(ValueError, match="word 1"):
            normalise_plan_words(
                [{"pitch_bin": 0, "dur_bin": 0, "energy_bin": 0, "dur_frames": 1},
                 {"pitch_bin": 0, "dur_bin": 0, "energy_bin": 9, "dur_frames": 1}],
                n_pitch_bins=16, n_duration_bins=16, n_energy_bins=8)

    def test_pitch_above_the_sentinel_raises(self):
        with pytest.raises(ValueError, match="pitch_bin"):
            self._norm([{"pitch_bin": 17, "dur_bin": 0, "energy_bin": 0,
                         "dur_frames": 1}])

    def test_seconds_are_converted_at_the_declared_frame_rate(self):
        from delay_dataset import normalise_plan_words
        out = normalise_plan_words(
            [{"pitch_bin": 0, "dur_bin": 0, "energy_bin": 0, "dur_s": 0.5}],
            n_pitch_bins=16, n_duration_bins=16, n_energy_bins=8, frame_rate=20.0)
        assert out[0]["dur_frames"] == 10


def _plan_span(on, item):
    ids = item["input_ids"]
    return ids.index(on["plan_start_id"]), ids.index(on["plan_end_id"]) + 1


class TestPlanAssembly:
    def test_legacy_arms_are_unchanged_with_the_gate_on(self, on):
        assert on["items_prose"][0]["input_ids"] == on["legacy_expected_prose"]
        assert on["items_tags"][0]["input_ids"] == on["legacy_expected_tags"]
        assert on["items_none"][0]["input_ids"] == on["legacy_expected_none"]

    def test_plan_block_is_exactly_the_program(self, on):
        item = on["items_plan"][0]
        lo, hi = _plan_span(on, item)
        expected = (
            [on["plan_start_id"]]
            + [on["plan_word_sep_id"], on["pitch_ids"][3],
               on["duration_ids"][1], on["energy_ids"][0]]
            + [on["plan_word_sep_id"], on["duration_ids"][4], on["energy_ids"][7]]
            + [on["plan_word_sep_id"], on["pitch_ids"][15],
               on["duration_ids"][9], on["energy_ids"][3]]
            + [on["plan_end_id"]])
        assert item["input_ids"][lo:hi] == expected

    def test_plan_only_arm_carries_no_prose_and_no_tags(self, on):
        item = on["items_plan"][0]
        ids = item["input_ids"]
        i = ids.index(on["think_start_id"])
        assert ids[i + 1] == on["plan_start_id"]
        assert ids[ids.index(on["plan_end_id"]) + 1] == on["think_end_id"]
        for tid in on["reasoning_ids"] + on["tag_ids"]:
            assert tid not in ids

    def test_tier_order_is_prose_then_tags_then_plan(self, on):
        ids = on["items_full"][0]["input_ids"]
        i_think = ids.index(on["think_start_id"])
        i_prose = ids.index(on["reasoning_ids"][0])
        i_tags = ids.index(on["tag_ids_after_prose"][0])
        i_plan = ids.index(on["plan_start_id"])
        i_end = ids.index(on["think_end_id"])
        assert i_think < i_prose < i_tags < i_plan < i_end

    def test_tags_plus_plan_has_no_prose(self, on):
        ids = on["items_tags+plan"][0]["input_ids"]
        assert ids[ids.index(on["think_start_id"]) + 1] == on["tag_ids"][0]
        for tid in on["reasoning_ids"]:
            assert tid not in ids

    def test_audio_content_is_identical_across_arms(self, on):
        base = on["items_none"][0]["codes"]
        for mode in ("prose", "tags", "plan", "tags+plan", "full"):
            item = on[f"items_{mode}"][0]
            codes = [c for c in item["codes"] if c[0] != -1]
            assert codes == [c for c in base if c[0] != -1], mode


class TestPerTierLossWeights:
    def test_each_tier_lands_on_exactly_its_own_positions(self, on):
        item = on["item_full_weighted"]
        w = item["loss_weights"]
        ids = item["input_ids"]
        lo = ids.index(on["think_start_id"])
        n_prose = len(on["reasoning_ids"])
        n_tags = len(on["tag_ids_after_prose"])
        p_lo, p_hi = _plan_span(on, item)
        assert w[lo] == pytest.approx(0.3)
        assert w[lo + 1: lo + 1 + n_prose] == pytest.approx([0.3] * n_prose)
        assert w[lo + 1 + n_prose: p_lo] == [1.0] * n_tags
        assert w[p_lo:p_hi] == [2.0] * (p_hi - p_lo)
        assert w[p_hi] == pytest.approx(0.3)
        assert set(w[:lo]) == {1.0} and set(w[p_hi + 1:]) == {1.0}

    def test_tag_weight_defaults_to_the_reasoning_weight(self, on):
        item = on["items_tags"][0]
        w = item["loss_weights"]
        lo = item["input_ids"].index(on["think_start_id"])
        n = 1 + len(on["tag_ids"]) + 1
        assert w[lo:lo + n] == pytest.approx([0.3] * n)

    def test_plan_tokens_are_supervised_targets(self, on):
        item = on["items_plan"][0]
        lo, hi = _plan_span(on, item)
        assert item["labels"][lo:hi] == item["input_ids"][lo:hi]


class TestScheduleAlignment:
    def test_position_to_word_map_is_exact(self, on):
        item = on["items_plan"][0]
        sched = item["plan_schedule"]
        a0 = item["audio_start"]
        expected_frames = [0] * 2 + [1] * 4 + [2] * 6
        for t, word in enumerate(expected_frames):
            assert sched[a0 + t - 1] == word, f"frame {t}"
        assert set(sched[: a0 - 1]) == {on["no_word"]}
        assert set(sched[a0 + len(expected_frames) - 1:]) == {on["no_word"]}

    def test_frame_zero_is_scheduled_on_the_audio_token(self, on):
        item = on["items_plan"][0]
        a0 = item["audio_start"]
        assert item["input_ids"][a0 - 1] == on["audio_start_id"]
        assert item["plan_schedule"][a0 - 1] == 0

    def test_lookahead_zero_shifts_by_exactly_one(self, on):
        shifted = on["item_plan_lookahead0"]["plan_schedule"]
        base = on["items_plan"][0]["plan_schedule"]
        assert shifted[1:] == base[:-1]

    def test_stretch_dilates_when_the_clip_is_longer_than_the_program(self, on):
        item = on["items_plan"][2]
        a0 = item["audio_start"]
        expected = [0] * 4 + [1] * 8 + [2] * 12
        assert item["plan_schedule"][a0 - 1: a0 - 1 + 24] == expected

    def test_pack_leaves_the_tail_unscheduled(self, on):
        item = on["items_plan_pack"][2]
        a0 = item["audio_start"]
        expected = [0] * 2 + [1] * 4 + [2] * 6 + [on["no_word"]] * 12
        assert item["plan_schedule"][a0 - 1: a0 - 1 + 24] == expected

    def test_slots_point_at_the_words_own_plan_tokens(self, on):
        item = on["items_plan"][0]
        ids = item["input_ids"]
        slots = item["plan_slots"]
        assert [ids[s[0]] for s in slots if s[0] != -1] == [
            on["pitch_ids"][3], on["pitch_ids"][15]]
        assert slots[1][0] == -1
        assert [ids[s[1]] for s in slots] == [on["duration_ids"][b]
                                              for b in (1, 4, 9)]
        assert [ids[s[2]] for s in slots] == [on["energy_ids"][b]
                                              for b in (0, 7, 3)]

    def test_durations_match_the_sidecar(self, on):
        assert on["items_plan"][0]["plan_durations"] == [2, 4, 6]

    def test_seconds_fall_back_at_the_declared_frame_rate(self, on):
        assert on["dur_s_frames"] == [11]


class TestMissingPlansAreVisible:
    def test_a_plan_less_row_degrades_to_the_stage1_layout(self, on):
        item = on["items_plan"][1]
        assert on["think_start_id"] not in item["input_ids"]
        assert item["plan_schedule"] is None
        assert item["plan_slots"] is None

    def test_the_degradation_is_counted(self, on):
        assert on["counters_plan"]["missing"] == 1
        assert on["counters_plan"]["seen"] == 3

    def test_construction_audits_coverage(self, on):
        audit = on["audit_plan"]
        assert audit["scanned"] == 3 and audit["with_plan"] == 2
        assert abs(audit["coverage"] - 2 / 3) < 1e-9
        assert audit["mean_words"] == 3.0
        assert audit["sampled"] is False

    def test_legacy_arms_run_no_audit(self, on):
        for mode in ("prose", "tags", "none"):
            assert on[f"audit_{mode}"] is None

    def test_a_corpus_with_no_plans_at_all_is_fatal(self, on):
        assert on["err_zero_coverage"] == "RuntimeError"

    def test_min_coverage_is_enforceable(self, on):
        assert on["err_min_coverage"] == "RuntimeError"

    def test_a_plan_arm_pinned_to_stage_1_is_fatal(self, on):
        assert on["err_stage1_plan_arm"] == "ValueError"

    def test_a_rejected_sidecar_is_dropped_not_used(self, on):
        assert on["rejected_item0_schedule"] is None
        assert on["rejected_counters"]["rejected"] == 1
        assert on["rejected_audit"]["rejected"] == 1
        assert on["rejected_audit"]["with_plan"] == 1

    def test_malformed_plans_raise_loudly(self, on):
        assert on["err_bad_dur_bin"] == "ValueError"
        assert on["err_no_dur_frames"] == "ValueError"
        assert "dur_frames" in on["err_no_dur_frames_msg"]
        assert "plan_extract" in on["err_no_dur_frames_msg"]


class TestTruncationPolicy:
    def test_reasoning_goes_first(self, on):
        item = on["item_tight"]
        ids = item["input_ids"]
        assert on["plan_start_id"] in ids, "the plan was dropped before the prose"
        for tid in on["reasoning_ids"]:
            assert tid not in ids
        assert on["counters_tight"]["prose_dropped"] == 1

    def test_the_window_is_respected(self, on):
        assert len(on["item_tight"]["input_ids"]) <= 40
        assert len(on["item_tighter"]["input_ids"]) <= 20

    def test_a_words_program_is_never_split(self, on):
        item = on["item_tighter"]
        ids = item["input_ids"]
        lo, hi = _plan_span(on, item)
        body = ids[lo + 1: hi - 1]
        assert body, "the whole program vanished"
        assert body[0] == on["plan_word_sep_id"]
        groups, cur = [], []
        for tid in body:
            if tid == on["plan_word_sep_id"]:
                if cur:
                    groups.append(cur)
                cur = []
            else:
                cur.append(tid)
        groups.append(cur)
        for g in groups:
            assert len(g) in (2, 3), f"partial word {g}"
            assert g[-2] in on["duration_ids"] and g[-1] in on["energy_ids"]
            if len(g) == 3:
                assert g[0] in on["pitch_ids"]

    def test_truncation_is_counted(self, on):
        assert on["counters_tighter"]["truncated"] >= 1

    def test_the_schedule_still_covers_only_surviving_words(self, on):
        item = on["item_tighter"]
        n_words = len(item["plan_slots"])
        assert max(item["plan_schedule"]) < n_words
        assert len(item["plan_durations"]) == n_words


class TestCollator:
    def test_plan_keys_appear_only_when_a_plan_is_present(self, on):
        assert not [k for k in on["batch_keys_plainarm"] if k.startswith("plan_")]
        for key in ("plan_schedule", "plan_slots", "plan_durations", "plan_lengths"):
            assert key in on["batch_keys"]

    def test_padding_never_looks_like_word_zero(self, on):
        sched = on["batch"]["plan_schedule"]
        assert len(set(len(r) for r in sched)) == 1, "ragged batch"
        assert set(sched[1]) == {-1}, "the plan-less row must be all NO_WORD"
        assert sched[0][-1] == -1

    def test_slots_and_durations_pad_correctly(self, on):
        slots = on["batch"]["plan_slots"]
        durations = on["batch"]["plan_durations"]
        lengths = on["batch"]["plan_lengths"]
        assert lengths == [3, 0, 3]
        assert slots[1] == [[-1, -1, -1]] * 3
        assert durations[1] == [0, 0, 0]
        assert durations[0] == [2, 4, 6]
        assert sum(durations[1]) == 0

    def test_schedule_rows_align_with_the_padded_input(self, on):
        sched = on["batch"]["plan_schedule"]
        item0 = on["items_plan"][0]
        a0 = item0["audio_start"]
        assert sched[0][a0 - 1] == 0
        assert sched[0][: a0 - 1] == [-1] * (a0 - 1)

if __name__ == "__main__":
    import tempfile
    with tempfile.TemporaryDirectory() as td:
        print(json.dumps(probe(td)))