training_code/latent_threads/tasks.py

"""Five delayed-selector tasks with token-exact, dot-replaceable digit CoTs.

Every task renders to [X: scenario prompt][Z: M single-digit body positions][Y: query + boxed
answer]. The query is revealed only in Y (after Z), so under the bottleneck mask (Y cannot attend
X) the dot block must carry the FULL ensemble of thread states. Body positions are single digit
tokens, 1:1 swappable for the filler dot (id 659) by the curriculum. See HYPOTHESES.md.
"""
from __future__ import annotations

import random
from dataclasses import dataclass

from filler_cot.filler import FILLER_ID


def digit_ids(tok) -> list[int]:
    """Single-token ids for digits 0..9, cached on the tokenizer."""
    if not hasattr(tok, "_lt_digit_ids"):
        ids = []
        for d in range(10):
            t = tok(str(d), add_special_tokens=False)["input_ids"]
            assert len(t) == 1, f"digit {d} not single-token: {t}"
            ids.append(t[0])
        tok._lt_digit_ids = ids
    return tok._lt_digit_ids


def dot_mask(n: int, n_dots: int, direction: str) -> list[bool]:
    assert 0 <= n_dots <= n, (n_dots, n)
    if direction == "front":
        return [i < n_dots for i in range(n)]
    if direction == "back":
        return [i >= n - n_dots for i in range(n)]
    raise ValueError(direction)


def body_ids(tok, digits: list[int], n_dots: int, direction: str) -> list[int]:
    did = digit_ids(tok)
    mask = dot_mask(len(digits), n_dots, direction)
    return [FILLER_ID if mask[i] else did[digits[i]] for i in range(len(digits))]


def _table_str(table: list[int]) -> str:
    return "  ".join(f"{i}:{table[i]}" for i in range(10))


def _sample_table(rng: random.Random) -> list[int]:
    return [rng.randrange(10) for _ in range(10)]


def _chase(table: list[int], start: int, m: int) -> list[int]:
    states, cur = [], start
    for _ in range(m):
        cur = table[cur]
        states.append(cur)
    return states


# ---------------------------------------------------------------- 1. parallel_select
@dataclass
class PSInst:
    table: list[int]
    regs: list[int]      # starting values of registers a..d
    results: list[int]   # f(reg) for each
    q: int               # queried register index


class ParallelSelect:
    """Four one-step threads; query names one register after the think block."""
    name = "parallel_select"
    REG = ["a", "b", "c", "d"]
    chance = 0.1

    def __init__(self, n_regs: int = 4):
        self.B = n_regs
        self.M = n_regs

    def sample(self, rng: random.Random) -> PSInst:
        table = _sample_table(rng)
        regs = [rng.randrange(10) for _ in range(self.B)]
        return PSInst(table, regs, [table[v] for v in regs], rng.randrange(self.B))

    def prompt(self, p: PSInst) -> str:
        regs = ", ".join(f"{self.REG[i]}={p.regs[i]}" for i in range(self.B))
        return (
            f'You are given a function f on the digits 0-9, written as "input:output" pairs:\n'
            f"{_table_str(p.table)}\n\n"
            f"Four registers hold starting values: {regs}. Compute f of each register's value and "
            f"remember all four results. Only AFTER your thinking, exactly one register will be "
            f"named -- answer with ONLY that register's result as \\boxed{{d}} (a single digit).\n\n"
            f"Reason inside <think> </think> -- write the four results in register order "
            f"{', '.join(self.REG[: self.B])}."
        )

    def body_digits(self, p: PSInst) -> list[int]:
        return list(p.results)

    def query(self, p: PSInst) -> str:
        return f"Result for register {self.REG[p.q]}: "

    def answer(self, p: PSInst) -> int:
        return p.results[p.q]

    def all_queries(self, p: PSInst) -> list[tuple[str, int]]:
        return [(f"Result for register {self.REG[i]}: ", p.results[i]) for i in range(self.B)]

    def probe_targets(self, p: PSInst) -> dict[str, int]:
        t = {f"res_{self.REG[i]}": p.results[i] for i in range(self.B)}
        t |= {f"start_{self.REG[i]}": p.regs[i] for i in range(self.B)}
        return t


# ---------------------------------------------------------------- 2. chase_select
@dataclass
class CSInst:
    table: list[int]
    starts: list[int]
    states: list[list[int]]  # states[b][t], t in 0..m-1
    q: int


class ChaseSelect:
    """Three interleaved pointer-chase chains (shared table); query names one chain late."""
    name = "chase_select"
    CHAIN = ["a", "b", "c", "d", "e", "f"]  # supports width-generalization up to 6 chains
    chance = 0.1

    def __init__(self, n_chains: int = 3, m: int = 4):
        self.B, self.m = n_chains, m
        self.M = n_chains * m

    def sample(self, rng: random.Random) -> CSInst:
        table = _sample_table(rng)
        starts = rng.sample(range(10), self.B)
        states = [_chase(table, s, self.m) for s in starts]
        return CSInst(table, starts, states, rng.randrange(self.B))

    def prompt(self, p: CSInst) -> str:
        starts = ", ".join(f"{self.CHAIN[b]}={p.starts[b]}" for b in range(self.B))
        return (
            f'You are given a function f on the digits 0-9, written as "input:output" pairs:\n'
            f"{_table_str(p.table)}\n\n"
            f"Three chains start at {starts}. Each chain applies f for {self.m} steps (each step: "
            f"replace the chain's current value v with f(v)). Track ALL three chains in parallel. "
            f"Only AFTER your thinking, exactly one chain will be named -- answer with ONLY that "
            f"chain's final value as \\boxed{{d}} (a single digit).\n\n"
            f"Reason inside <think> </think> -- after each step write the new values of chains "
            f"{', '.join(self.CHAIN[: self.B])} in order."
        )

    def body_digits(self, p: CSInst) -> list[int]:
        # step-major interleave: a1 b1 c1 a2 b2 c2 ...
        return [p.states[b][t] for t in range(self.m) for b in range(self.B)]

    def query(self, p: CSInst) -> str:
        return f"Final value of chain {self.CHAIN[p.q]}: "

    def answer(self, p: CSInst) -> int:
        return p.states[p.q][-1]

    def all_queries(self, p: CSInst) -> list[tuple[str, int]]:
        return [(f"Final value of chain {self.CHAIN[b]}: ", p.states[b][-1]) for b in range(self.B)]

    def probe_targets(self, p: CSInst) -> dict[str, int]:
        return {f"s_{self.CHAIN[b]}{t + 1}": p.states[b][t] for b in range(self.B) for t in range(self.m)}

    def step_states(self, p: CSInst) -> list[list[int]]:
        """Ground-truth parallel state per latent step: step t -> all chains' values after t."""
        return [[p.states[b][t] for b in range(self.B)] for t in range(self.m)]


# ---------------------------------------------------------------- 3. step_select
@dataclass
class SSInst:
    table: list[int]
    start: int
    states: list[int]
    q: int  # queried step, 1-based


class StepSelect:
    """One 8-step chain; the query asks for the running value after step t (t revealed late)."""
    name = "step_select"
    chance = 0.1

    def __init__(self, m: int = 8):
        self.m = m
        self.M = m

    def sample(self, rng: random.Random) -> SSInst:
        table = _sample_table(rng)
        start = rng.randrange(10)
        return SSInst(table, start, _chase(table, start, self.m), 1 + rng.randrange(self.m))

    def prompt(self, p: SSInst) -> str:
        return (
            f'You are given a function f on the digits 0-9, written as "input:output" pairs:\n'
            f"{_table_str(p.table)}\n\n"
            f"Start with the value {p.start}. Apply f repeatedly for {self.m} steps (each step: "
            f"replace the current value v with f(v)). Remember the running value after EVERY step. "
            f"Only AFTER your thinking, you will be asked for the value after one specific step -- "
            f"answer it as \\boxed{{d}} (a single digit).\n\n"
            f"Reason inside <think> </think> -- write the running value after each step."
        )

    def body_digits(self, p: SSInst) -> list[int]:
        return list(p.states)

    def query(self, p: SSInst) -> str:
        return f"Value after step {p.q}: "

    def answer(self, p: SSInst) -> int:
        return p.states[p.q - 1]

    def all_queries(self, p: SSInst) -> list[tuple[str, int]]:
        return [(f"Value after step {t + 1}: ", p.states[t]) for t in range(self.m)]

    def probe_targets(self, p: SSInst) -> dict[str, int]:
        return {f"s{t + 1}": p.states[t] for t in range(self.m)}

    def step_states(self, p: SSInst) -> list[list[int]]:
        return [[p.states[t]] for t in range(self.m)]


# ---------------------------------------------------------------- 4. coin_track
@dataclass
class CTInst:
    starts: list[int]
    events: list[tuple[int, int, int]]  # (entity, delta, new_count)
    finals: list[int]
    q: int


class CoinTrack:
    """NL world-state tracking: 3 players, 8 win/lose events; query names one player late."""
    name = "coin_track"
    NAMES = ["Anna", "Ben", "Cara"]
    chance = 0.1

    def __init__(self, n_events: int = 8):
        self.E = n_events
        self.M = n_events

    def sample(self, rng: random.Random) -> CTInst:
        counts = [rng.randint(2, 7) for _ in range(3)]
        starts = list(counts)
        events = []
        for _ in range(self.E):
            while True:
                e = rng.randrange(3)
                delta = rng.choice([-3, -2, -1, 1, 2, 3])
                if 0 <= counts[e] + delta <= 9:
                    break
            counts[e] += delta
            events.append((e, delta, counts[e]))
        return CTInst(starts, events, list(counts), rng.randrange(3))

    def prompt(self, p: CTInst) -> str:
        ev = " ".join(
            f"{self.NAMES[e]} {'wins' if d > 0 else 'loses'} {abs(d)} coin{'s' if abs(d) > 1 else ''}."
            for e, d, _ in p.events
        )
        return (
            f"Anna, Ben, and Cara play a coin game. Anna starts with {p.starts[0]} coins, Ben with "
            f"{p.starts[1]}, and Cara with {p.starts[2]}. Then, in order: {ev}\n\n"
            f"Track every player's coin count as you read. Only AFTER your thinking, exactly one "
            f"player will be named -- answer with ONLY that player's final coin count as "
            f"\\boxed{{d}} (a single digit).\n\n"
            f"Reason inside <think> </think> -- after each event write the affected player's new "
            f"coin count."
        )

    def body_digits(self, p: CTInst) -> list[int]:
        return [new for _, _, new in p.events]

    def query(self, p: CTInst) -> str:
        return f"Final coin count of {self.NAMES[p.q]}: "

    def answer(self, p: CTInst) -> int:
        return p.finals[p.q]

    def all_queries(self, p: CTInst) -> list[tuple[str, int]]:
        return [(f"Final coin count of {self.NAMES[i]}: ", p.finals[i]) for i in range(3)]

    def probe_targets(self, p: CTInst) -> dict[str, int]:
        t = {f"e{i + 1}": new for i, (_, _, new) in enumerate(p.events)}
        t |= {f"who{i + 1}": e for i, (e, _, _) in enumerate(p.events)}
        t |= {f"final_{self.NAMES[i]}": p.finals[i] for i in range(3)}
        return t


# ---------------------------------------------------------------- 5. agg_select
@dataclass
class AGInst:
    table: list[int]
    starts: list[int]
    states: list[list[int]]
    qtype: int  # 0 value, 1 argmax, 2 argmin, 3 sum mod 10
    qarg: int   # chain index for qtype 0


class AggSelect:
    """Three 2-step chains (distinct finals); late query = value / argmax / argmin / sum mod 10."""
    name = "agg_select"

    def __init__(self, n_chains: int = 3, m: int = 2, qtypes: tuple[int, ...] = (0, 1, 2, 3),
                 sum_in_body: bool = False):
        self.B, self.m = n_chains, m
        self.sum_in_body = sum_in_body  # surface CoT ends with the running sum mod 10 (then dotted)
        self.M = n_chains * m + (1 if sum_in_body else 0)
        self.qtypes = tuple(qtypes)
        self.chance = sum(0.1 if t in (0, 3) else 1 / 3 for t in self.qtypes) / len(self.qtypes)

    def sample(self, rng: random.Random) -> AGInst:
        table = _sample_table(rng)
        for _ in range(200):
            starts = rng.sample(range(10), self.B)
            states = [_chase(table, s, self.m) for s in starts]
            if len({st[-1] for st in states}) == self.B:
                break
            table = _sample_table(rng)
        else:
            raise RuntimeError("could not sample distinct finals")
        qtype = rng.choice(self.qtypes)
        return AGInst(table, starts, states, qtype, rng.randrange(self.B))

    def prompt(self, p: AGInst) -> str:
        starts = ", ".join(f"chain {b + 1} at {p.starts[b]}" for b in range(self.B))
        return (
            f'You are given a function f on the digits 0-9, written as "input:output" pairs:\n'
            f"{_table_str(p.table)}\n\n"
            f"Three chains start: {starts}. Each chain applies f for {self.m} steps (each step: "
            f"replace the chain's current value v with f(v)). Track ALL three chains. Only AFTER "
            f"your thinking, you will be asked ONE of: a chain's final value; the number of the "
            f"chain with the largest final value; the number of the chain with the smallest final "
            f"value; or the last digit of the sum of the three final values. Answer with ONLY a "
            f"single digit as \\boxed{{d}}.\n\n"
            f"Reason inside <think> </think> -- after each step write the new values of chains 1, "
            f"2, 3 in order." + (" Then write the last digit of the sum of the three final values."
                                 if self.sum_in_body else "")
        )

    def body_digits(self, p: AGInst) -> list[int]:
        base = [p.states[b][t] for t in range(self.m) for b in range(self.B)]
        return base + ([sum(self._finals(p)) % 10] if self.sum_in_body else [])

    def _finals(self, p: AGInst) -> list[int]:
        return [st[-1] for st in p.states]

    def _qa(self, p: AGInst, qtype: int, qarg: int) -> tuple[str, int]:
        v = self._finals(p)
        if qtype == 0:
            return f"Final value of chain {qarg + 1}: ", v[qarg]
        if qtype == 1:
            return "Chain number with the largest final value: ", 1 + max(range(self.B), key=v.__getitem__)
        if qtype == 2:
            return "Chain number with the smallest final value: ", 1 + min(range(self.B), key=v.__getitem__)
        return "Last digit of the sum of the three final values: ", sum(v) % 10

    def query(self, p: AGInst) -> str:
        return self._qa(p, p.qtype, p.qarg)[0]

    def answer(self, p: AGInst) -> int:
        return self._qa(p, p.qtype, p.qarg)[1]

    def all_queries(self, p: AGInst) -> list[tuple[str, int]]:
        qs = [self._qa(p, 0, b) for b in range(self.B)]
        return qs + [self._qa(p, 1, 0), self._qa(p, 2, 0), self._qa(p, 3, 0)]

    def probe_targets(self, p: AGInst) -> dict[str, int]:
        v = self._finals(p)
        t = {f"v{b + 1}": v[b] for b in range(self.B)}
        t |= {f"s{b + 1}_1": p.states[b][0] for b in range(self.B)}
        t |= {"argmax": 1 + max(range(self.B), key=v.__getitem__),
              "argmin": 1 + min(range(self.B), key=v.__getitem__),
              "sum10": sum(v) % 10}
        return t


# ---------------------------------------------------------------- 6. hyp_track (semantic)
# Abductive surviving-hypothesis tracking: each THREAD is a candidate suspect whose "train of
# thought" is the accumulated chain of semantic constraint-checks across the clues (recall the
# suspect's established attribute, check it against each clue — an entailment, NOT arithmetic).
# The latent z_t holds the running viability VECTOR (which suspects remain) after clue t — a
# vocabulary-space superposition of parallel extended hypotheses. The unique-survivor "culprit"
# query is correct only if ALL threads are tracked, so every train is load-bearing for it.
_HYP_NAMES = ["Anna", "Ben", "Cara", "Dan", "Eve", "Finn"]
# each attribute: (val0 clause, val0 clue, val1 clause, val1 clue)
_HYP_ATTRS = [
    ("is left-handed", "The culprit is left-handed.", "is right-handed", "The culprit is right-handed."),
    ("has dark hair", "The culprit has dark hair.", "has fair hair", "The culprit has fair hair."),
    ("is tall", "The culprit is tall.", "is short", "The culprit is short."),
    ("wears glasses", "The culprit wears glasses.", "does not wear glasses", "The culprit does not wear glasses."),
    ("has a local accent", "The culprit has a local accent.", "has a foreign accent", "The culprit has a foreign accent."),
]


@dataclass
class HypInst:
    attrs: list[list[int]]   # attrs[b][a] in {0,1} for each suspect b, attribute a
    culprit: int
    clue_order: list[int]    # permutation of attribute indices = reveal order
    q: tuple[int, int, int]  # (kind, b, t): kind 0 = "was suspect b viable after clue t", 1 = culprit


class HypTrack:
    """Semantic delayed-selector: track which suspects survive clues; the answer requires all threads."""
    name = "hyp_track"

    def __init__(self, n_suspects: int = 4, n_attrs: int = 5):
        assert n_suspects <= 2 ** n_attrs and n_suspects <= len(_HYP_NAMES)
        self.N, self.A = n_suspects, n_attrs
        self.M = n_attrs  # one latent step per clue
        self.names = _HYP_NAMES[:n_suspects]
        self.chance = 0.5 * (1 / n_suspects) + 0.5 * 0.5  # culprit (1/N) vs viability (binary)

    def sample(self, rng: random.Random) -> HypInst:
        # Staggered eliminations: assign each non-culprit a DISTINCT first-differing clue-step, so
        # the viability vector evolves non-trivially across all M steps (each thread is an extended
        # train, not isolated by clue 1). Requires N-1 <= M.
        assert self.N - 1 <= self.M, "need at least N-1 attributes to stagger eliminations"
        clue_order = list(range(self.A))
        rng.shuffle(clue_order)
        culprit_vals = [rng.randrange(2) for _ in range(self.A)]
        culprit = rng.randrange(self.N)
        attrs: list[list[int]] = [None] * self.N  # type: ignore
        attrs[culprit] = list(culprit_vals)
        others = [b for b in range(self.N) if b != culprit]
        kill_steps = rng.sample(range(1, self.M + 1), len(others))  # distinct, 1-indexed
        for b, k in zip(others, kill_steps):
            v = list(culprit_vals)               # match culprit on steps 1..k-1
            v[clue_order[k - 1]] ^= 1            # first difference exactly at clue-step k
            for i in range(k, self.A):           # randomize attributes revealed after the kill step
                v[clue_order[i]] = rng.randrange(2)
            attrs[b] = v
        if rng.random() < 0.5:
            q = (1, 0, 0)
        else:
            q = (0, rng.randrange(self.N), rng.randrange(self.A))
        return HypInst(attrs, culprit, clue_order, q)

    def _viable(self, p: HypInst, b: int, t: int) -> int:
        """1 iff suspect b matches the culprit on the first t revealed attributes (after clue t)."""
        return int(all(p.attrs[b][p.clue_order[i]] == p.attrs[p.culprit][p.clue_order[i]] for i in range(t)))

    def _clause(self, a: int, v: int) -> str:
        return _HYP_ATTRS[a][0 if v == 0 else 2]

    def _clue(self, a: int, v: int) -> str:
        return _HYP_ATTRS[a][1 if v == 0 else 3]

    def prompt(self, p: HypInst) -> str:
        lines = []
        for b in range(self.N):
            cl = [self._clause(a, p.attrs[b][a]) for a in range(self.A)]
            desc = ", ".join(cl[:-1]) + ", and " + cl[-1]
            lines.append(f"- {self.names[b]} {desc}.")
        clues = "\n".join(f"{i + 1}. {self._clue(a, p.attrs[p.culprit][a])}"
                          for i, a in enumerate(p.clue_order))
        return (
            f"A crime was committed by exactly one of {self.N} suspects. What is known about each:\n"
            + "\n".join(lines) +
            f"\n\nClues about the culprit are revealed in order:\n{clues}\n\n"
            f"Track which suspects stay consistent with the clues as each one is revealed. Only "
            f"AFTER your thinking, you will be asked either about one suspect at one clue, or for "
            f"the culprit. Answer with a single digit.\n\n"
            f"Reason inside <think> </think> -- after each clue, note which suspects remain "
            f"consistent."
        )

    def body_digits(self, p: HypInst) -> list[int]:
        # optional dot/surface form: number of suspects still consistent after each clue
        return [sum(self._viable(p, b, t + 1) for b in range(self.N)) for t in range(self.M)]

    def _qa(self, p: HypInst, q: tuple[int, int, int]) -> tuple[str, int]:
        kind, b, t = q
        if kind == 1:
            return f"Which numbered suspect (1-{self.N}) is the culprit? ", p.culprit + 1
        return (f"Was {self.names[b]} still consistent after clue {t + 1}? Answer 1 for yes, 0 for no: ",
                self._viable(p, b, t + 1))

    def query(self, p: HypInst) -> str:
        return self._qa(p, p.q)[0]

    def answer(self, p: HypInst) -> int:
        return self._qa(p, p.q)[1]

    def all_queries(self, p: HypInst) -> list[tuple[str, int]]:
        qs = [self._qa(p, (0, b, t)) for b in range(self.N) for t in range(self.A)]
        return qs + [self._qa(p, (1, 0, 0))]

    def step_states(self, p: HypInst) -> list[list[int]]:
        """Ground-truth latent thought per clue-step: viability bit of each suspect after clue t."""
        return [[self._viable(p, b, t + 1) for b in range(self.N)] for t in range(self.M)]

    def probe_targets(self, p: HypInst) -> dict[str, int]:
        t = {f"v{b + 1}_{s + 1}": self._viable(p, b, s + 1) for b in range(self.N) for s in range(self.A)}
        t["culprit"] = p.culprit + 1
        return t


# ---------------------------------------------------------------- 7. diffuse (coupled, parallel-necessary)
# 1-D coupled cellular automaton on a ring of K cells: x_i(t+1) = (x_{i-1}(t) + x_{i+1}(t)) mod 10.
# PARALLELISM IS NECESSARY TO SOLVE: cell j's value at step M depends on its light-cone of width
# 2M+1; with M >= K/2 every cell's final depends on ALL initial cells, so you cannot track one
# thread in isolation. MULTI-STEP: M coupled layers (no parallel shortcut for large M). TIGHT-
# FEASIBLE: the rule is fixed (learned, not per-instance), so the latent only carries the K-cell
# ROW forward — z_1 reads the small initial row, z_t applies the rule -> genuine load-bearing
# recurrence. z_t = the whole row at step t = K parallel threads superposed.
@dataclass
class CAInst:
    init: list[int]          # initial row, K cells
    rows: list[list[int]]    # rows[t] = state after t+1 steps; len == M
    q: int                   # queried cell index


class Diffuse:
    name = "diffuse"

    def __init__(self, k: int = 3, m: int = 4, mod: int = 10):
        # Coupled ring CA, parallelism PROVABLY necessary (light cone). For the DOT organism the
        # surface CoT (and thus the dot block) is the FULL ROW per step -> M = K*m positions, each
        # carrying one cell at one step (step-major). Keep K*m near ~12 (chase's proven dot count).
        self.K, self.m, self.mod = k, m, mod
        self.M = k * m            # dot-block length = full row x steps (step-major)
        self.chance = 1.0 / mod

    def _step(self, row: list[int]) -> list[int]:
        K = self.K
        return [(row[(i - 1) % K] + row[(i + 1) % K]) % self.mod for i in range(K)]

    def sample(self, rng: random.Random) -> CAInst:
        init = [rng.randrange(self.mod) for _ in range(self.K)]
        rows, cur = [], init
        for _ in range(self.m):
            cur = self._step(cur)
            rows.append(cur)
        return CAInst(init, rows, rng.randrange(self.K))

    def prompt(self, p: CAInst) -> str:
        cells = ", ".join(f"c{i + 1}={p.init[i]}" for i in range(self.K))
        return (
            f"{self.K} cells sit in a ring (c1..c{self.K}, and c{self.K} is adjacent to c1). Initial "
            f"values: {cells}.\n\nEach step, every cell SIMULTANEOUSLY becomes the sum modulo 10 of "
            f"its two ring neighbours (left and right). Apply this for {self.m} steps. Only AFTER "
            f"your thinking, you will be asked for one cell's final value -- answer with ONLY that "
            f"value as \\boxed{{d}} (a single digit).\n\n"
            f"Reason inside <think> </think> -- after each step write all {self.K} cell values in "
            f"order c1..c{self.K}."
        )

    def body_digits(self, p: CAInst) -> list[int]:
        # surface form = the FULL ROW per step, step-major (all K cells of step 1, then step 2, ...)
        # so the dot block carries every cell at every step -> parallelism necessary, K*m positions.
        return [p.rows[t][i] for t in range(self.m) for i in range(self.K)]

    def query(self, p: CAInst) -> str:
        return f"Final value of cell c{p.q + 1}: "

    def answer(self, p: CAInst) -> int:
        return p.rows[-1][p.q]

    def all_queries(self, p: CAInst) -> list[tuple[str, int]]:
        return [(f"Final value of cell c{i + 1}: ", p.rows[-1][i]) for i in range(self.K)]

    def step_states(self, p: CAInst) -> list[list[int]]:
        """Ground-truth parallel state per latent step: the whole row after step t (K threads)."""
        return [list(p.rows[t]) for t in range(self.m)]

    def probe_targets(self, p: CAInst) -> dict[str, int]:
        return {f"c{i + 1}_{t + 1}": p.rows[t][i] for i in range(self.K) for t in range(self.m)}


# ---------------------------------------------------------------- 8. journeys (cohesive NL trains)
# K travelers each walk a contiguous M-room PATH (a cohesive natural-language "train of thought"
# that extends over M positions). Fixed transition rule next = rooms[(7*idx+3) mod 10] (a fixed
# permutation -> no per-instance table -> the within-thread Markov chain has everything it needs).
# Thread-MAJOR: traveler b's M rooms are a contiguous latent span. Delayed query names one traveler
# -> every traveler's train must be maintained (each load-bearing). The latent symbols are ROOM
# WORDS (natural language), one per position; the per-thread span reads as a coherent journey.
_ROOMS = ["kitchen", "garden", "attic", "cellar", "hallway", "study", "bedroom", "bathroom",
          "garage", "library"]
_TRAVELERS = ["Anna", "Ben", "Cara", "Dan", "Eve"]
_J_MUL, _J_ADD = 7, 3  # gcd(7,10)=1 -> a permutation (non-degenerate, every room reachable)


@dataclass
class JInst:
    starts: list[int]          # start room index per traveler
    paths: list[list[int]]     # paths[b][t] = room index after step t+1; len == M
    q: int                     # queried traveler


class Journeys:
    name = "journeys"
    chance = 0.1               # 10 rooms

    def __init__(self, k: int = 3, m: int = 4):
        self.K, self.m = k, m
        self.M = k * m         # thread-major latent block length = K travelers x M rooms
        self.rooms = _ROOMS

    def _step(self, idx: int) -> int:
        return (_J_MUL * idx + _J_ADD) % 10

    def sample(self, rng: random.Random) -> JInst:
        starts = [rng.randrange(10) for _ in range(self.K)]
        paths = []
        for s in starts:
            cur, p = s, []
            for _ in range(self.m):
                cur = self._step(cur); p.append(cur)
            paths.append(p)
        return JInst(starts, paths, rng.randrange(self.K))

    def symbol_ids(self, tok) -> list[int]:
        if not hasattr(tok, "_journey_room_ids"):
            ids = []
            for r in self.rooms:
                t = tok(" " + r, add_special_tokens=False)["input_ids"]
                if len(t) != 1:
                    t = tok(r, add_special_tokens=False)["input_ids"]
                assert len(t) == 1, (r, t)
                ids.append(t[0])
            tok._journey_room_ids = ids
        return tok._journey_room_ids

    n_symbols = 10

    def prompt(self, p: JInst) -> str:
        starts = ", ".join(f"person {b+1} ({_TRAVELERS[b]}) in the {self.rooms[p.starts[b]]}" for b in range(self.K))
        return (
            f"{self.K} people wander a house with 10 rooms. Each minute, a person in room number i "
            f"moves to room number (7*i + 3) mod 10 (rooms are numbered 0-9 in the order: "
            f"{', '.join(f'{i}={r}' for i, r in enumerate(self.rooms))}). They all start: {starts}. "
            f"They each take {self.m} steps. Only AFTER your thinking, one person will be named -- "
            f"answer with ONLY the room they end in.\n\n"
            f"Reason inside <think> </think> -- write each person's full path, one person at a time."
        )

    def step_states(self, p: JInst) -> list[list[int]]:
        # thread-MAJOR grid: outer = traveler, inner = their M rooms (one cohesive span per thread)
        return [list(p.paths[b]) for b in range(self.K)]

    def query(self, p: JInst) -> str:
        return f"Which room does person number {p.q + 1} end in? "

    def answer(self, p: JInst) -> int:
        return p.paths[p.q][-1]

    def all_queries(self, p: JInst) -> list[tuple[str, int]]:
        return [(f"Which room does person number {b + 1} end in? ", p.paths[b][-1]) for b in range(self.K)]

    def probe_targets(self, p: JInst) -> dict[str, int]:
        return {f"{_TRAVELERS[b]}_{t + 1}": p.paths[b][t] for b in range(self.K) for t in range(self.m)}


# ---------------------------------------------------------------- 9. tales (statement-trains)
# Like journeys, but each latent STEP is a full natural-language STATEMENT (a multi-token clause)
# rather than a single logical token. Per character b the train reads as a cohesive mini-story:
#   "Anna entered the kitchen and found the key. Anna entered the garden and found the lamp. ..."
# Each statement is a fixed 9-token span with TWO computed slots (room@+3, item@+7) carrying TWO
# INDEPENDENT load-bearing chains (room: (7r+3)%10, item: (3i+1)%10 -- different permutations). So
# every statement bears parallel load in >1 token, and the delayed query selects one (character,
# attribute) of 2K chains -> all must be maintained. The template words are FIXED embeddings (so the
# latent span literally spells out a statement); the computed slots form per-attribute Markov chains
# (each attends only its same-attribute predecessor + the prompt at step 1) -> per-statement load-
# bearing, no prompt-recompute past step 1. See markov_tales.py.
_ITEMS = ["key", "map", "coin", "book", "lamp", "ring", "gem", "knife", "torch", "sword"]
_TALE_R_MUL, _TALE_R_ADD = 7, 3   # room rule (a permutation of 0..9)
_TALE_I_MUL, _TALE_I_ADD = 3, 1   # item rule (a DIFFERENT permutation -> chains diverge)


@dataclass
class TaInst:
    start_room: list[int]
    start_item: list[int]
    rooms: list[list[int]]     # rooms[b][t] = room after step t+1
    items: list[list[int]]     # items[b][t] = item after step t+1
    q: int                     # queried character
    attr: int                  # 0 = room, 1 = item


class Tales:
    name = "tales"
    chance = 0.1               # 10 rooms / 10 items
    STMT = 9                   # tokens per statement: [name, entered, the, ROOM, and, found, the, ITEM, .]
    ROOM_OFF, ITEM_OFF = 3, 7  # computed-slot offsets within a statement

    def __init__(self, k: int = 3, m: int = 4):
        self.K, self.m = k, m
        self.M = k * m * self.STMT   # latent statement-token count (thread-major)
        self.rooms_vocab = _ROOMS
        self.items_vocab = _ITEMS
        self.names = _TRAVELERS

    def _rstep(self, r: int) -> int:
        return (_TALE_R_MUL * r + _TALE_R_ADD) % 10

    def _istep(self, i: int) -> int:
        return (_TALE_I_MUL * i + _TALE_I_ADD) % 10

    def sample(self, rng: random.Random) -> TaInst:
        sr = [rng.randrange(10) for _ in range(self.K)]
        si = [rng.randrange(10) for _ in range(self.K)]
        rooms, items = [], []
        for b in range(self.K):
            cr, rl = sr[b], []
            for _ in range(self.m):
                cr = self._rstep(cr); rl.append(cr)
            ci, il = si[b], []
            for _ in range(self.m):
                ci = self._istep(ci); il.append(ci)
            rooms.append(rl); items.append(il)
        return TaInst(sr, si, rooms, items, rng.randrange(self.K), rng.randrange(2))

    def _single_ids(self, tok, words, attr):
        key = f"_tales_ids_{attr}"
        if not hasattr(tok, key):
            ids = []
            for w in words:
                t = tok(" " + w, add_special_tokens=False)["input_ids"]
                assert len(t) == 1, (w, t)
                ids.append(t[0])
            setattr(tok, key, ids)
        return getattr(tok, key)

    def room_ids(self, tok):
        return self._single_ids(tok, self.rooms_vocab, "room")

    def item_ids(self, tok):
        return self._single_ids(tok, self.items_vocab, "item")

    def union_ids(self, tok):
        return self.room_ids(tok) + self.item_ids(tok)   # [20]: rooms 0-9, items 10-19

    def template_row(self, tok, b):
        """9 token ids for character b's statement; computed slots (ROOM_OFF, ITEM_OFF) are pad."""
        if not hasattr(tok, "_tales_tmpl"):
            w = lambda s: tok(s, add_special_tokens=False)["input_ids"]
            assert all(len(w(" " + x)) == 1 for x in ("entered", "the", "and", "found"))
            tok._tales_tmpl = dict(entered=w(" entered")[0], the=w(" the")[0], andw=w(" and")[0],
                                   found=w(" found")[0], dot=w(".")[0],
                                   names=[w(" " + n)[0] for n in self.names])
        t = tok._tales_tmpl
        return [t["names"][b], t["entered"], t["the"], tok.pad_token_id, t["andw"],
                t["found"], t["the"], tok.pad_token_id, t["dot"]]

    def room_states(self, p):
        return [list(p.rooms[b]) for b in range(self.K)]   # [K][m]

    def item_states(self, p):
        return [list(p.items[b]) for b in range(self.K)]   # [K][m]

    def prompt(self, p: TaInst) -> str:
        rmap = ", ".join(f"{i}={r}" for i, r in enumerate(self.rooms_vocab))
        imap = ", ".join(f"{i}={r}" for i, r in enumerate(self.items_vocab))
        starts = "; ".join(
            f"{self.names[b]} starts in the {self.rooms_vocab[p.start_room[b]]} holding the "
            f"{self.items_vocab[p.start_item[b]]}" for b in range(self.K))
        return (
            f"{self.K} people explore a house. Rooms 0-9: {rmap}. Items 0-9: {imap}. Each minute, a "
            f"person in room i walks to room (7*i+3) mod 10, and the item j they hold turns into item "
            f"(3*j+1) mod 10. {starts}. They each take {self.m} steps. Only AFTER your thinking, one "
            f"person will be named and you'll be asked for EITHER their final room OR their final "
            f"item -- answer with ONLY that one word.\n\nReason inside <think> </think> -- for each "
            f"person in turn, write one sentence per step: '<name> entered the <room> and found the "
            f"<item>.'")

    def query(self, p: TaInst) -> str:
        what = "room" if p.attr == 0 else "item"
        return f"What {what} does person number {p.q + 1} end with? "

    def answer(self, p: TaInst) -> int:
        """Union index: room 0-9, or 10 + item 0-9."""
        return p.rooms[p.q][-1] if p.attr == 0 else 10 + p.items[p.q][-1]

    def answer_token(self, tok, p: TaInst) -> int:
        return self.union_ids(tok)[self.answer(p)]

    def all_queries(self, p: TaInst):
        qs = [(f"What room does person number {b + 1} end with? ", p.rooms[b][-1]) for b in range(self.K)]
        qs += [(f"What item does person number {b + 1} end with? ", 10 + p.items[b][-1]) for b in range(self.K)]
        return qs

    def probe_targets(self, p: TaInst) -> dict[str, int]:
        d = {}
        for b in range(self.K):
            for t in range(self.m):
                d[f"{self.names[b]}_room_{t + 1}"] = p.rooms[b][t]
                d[f"{self.names[b]}_item_{t + 1}"] = p.items[b][t]
        return d


# ---------------------------------------------------------------- 10. sagas (reasoning-path worlds)
# K parallel SIMULATED WORLDS, each a running quantity evolved over M steps by DIVERSE narrated
# operations. Unlike journeys/tales (each step a fixed lookup of a placeholder variable), every step
# is a genuine INFERENCE that combines TWO inputs — a per-step premise (op + operand, given in the
# statement) AND the running conclusion (from the predecessor) — via a real arithmetic operation:
#   "the merchant gained 3 coins, now 7."   (7 = 4 + 3, computed; not a placeholder)
# Ops: gained/earned (+k), lost/spent/drained (-k), scaled (*k, k in {2,3}); all mod 10. The * ops
# make the chain ORDER-DEPENDENT (shuffle breaks it) and no op wipes history (every step load-bearing).
# Each world has its own theme (subject+resource) and its own event sequence -> K diverse paths; the
# delayed query asks one world's final tally, so all K running conclusions must be carried forward.
_SAGA_THEMES = [("merchant", "coins"), ("general", "troops"), ("builder", "bricks"),
                ("captain", "crates"), ("smith", "scrolls")]
_SAGA_OPS = {"add": ["gained", "earned"], "sub": ["lost", "spent", "drained"], "mul": ["scaled"]}
_SAGA_CONNS = ["now", "leaving", "reaching", "totaling"]


@dataclass
class SagaInst:
    themes: list[int]           # theme index per world
    starts: list[int]           # starting tally per world (0-9)
    ops: list[list[str]]        # ops[b][t] in {"add","sub","mul"}
    verbs: list[list[int]]      # verbs[b][t] = index into _SAGA_OPS[op] (surface synonym)
    operands: list[list[int]]   # operands[b][t]
    conns: list[list[int]]      # connective index per step
    states: list[list[int]]     # states[b][t] = running tally AFTER step t (mod 10)
    q: int                      # queried world


class Sagas:
    name = "sagas"
    chance = 0.1
    STMT = 7                    # [subject, opverb, operand, resource, conn, STATE, .]
    SUBJ_OFF, OP_OFF, OPND_OFF, RES_OFF, CONN_OFF, STATE_OFF, DOT_OFF = 0, 1, 2, 3, 4, 5, 6

    def __init__(self, k: int = 3, m: int = 6):
        self.K, self.m = k, m
        self.M = k * m * self.STMT

    def _apply(self, op, prev, k):
        if op == "add": return (prev + k) % 10
        if op == "sub": return (prev - k) % 10
        return (prev * k) % 10            # mul

    def sample(self, rng: random.Random) -> SagaInst:
        themes = rng.sample(range(len(_SAGA_THEMES)), self.K)   # distinct themes -> diverse worlds
        starts = [rng.randrange(10) for _ in range(self.K)]
        ops, verbs, operands, conns, states = [], [], [], [], []
        for b in range(self.K):
            opl, vl, kl, cl, sl = [], [], [], [], []
            cur = starts[b]
            mul_pos = rng.randrange(self.m)                     # guarantee >=1 mul -> order-dependent
            for t in range(self.m):
                op = "mul" if t == mul_pos else rng.choice(["add", "sub", "mul"])
                k = rng.choice([2, 3]) if op == "mul" else rng.randint(1, 9)
                cur = self._apply(op, cur, k)
                opl.append(op); vl.append(rng.randrange(len(_SAGA_OPS[op]))); kl.append(k)
                cl.append(rng.randrange(len(_SAGA_CONNS))); sl.append(cur)
            ops.append(opl); verbs.append(vl); operands.append(kl); conns.append(cl); states.append(sl)
        return SagaInst(themes, starts, ops, verbs, operands, conns, states, rng.randrange(self.K))

    def verb_token(self, b, t, p):
        return _SAGA_OPS[p.ops[b][t]][p.verbs[b][t]]

    def _ids(self, tok, words, key):
        if not hasattr(tok, key):
            ids = []
            for w in words:
                tt = tok(" " + w, add_special_tokens=False)["input_ids"]; assert len(tt) == 1, (w, tt); ids.append(tt[0])
            setattr(tok, key, ids)
        return getattr(tok, key)

    def _verb_id(self, tok, word):
        if not hasattr(tok, "_saga_vmap"):
            flat = [v for op in ("add", "sub", "mul") for v in _SAGA_OPS[op]]
            tok._saga_vmap = {w: self._ids(tok, flat, "_saga_verbs")[i] for i, w in enumerate(flat)}
        return tok._saga_vmap[word]

    def conn_ids(self, tok):
        return self._ids(tok, _SAGA_CONNS, "_saga_conns")

    def subject_ids(self, tok):
        return self._ids(tok, [s for s, _ in _SAGA_THEMES], "_saga_subj")

    def resource_ids(self, tok):
        return self._ids(tok, [r for _, r in _SAGA_THEMES], "_saga_res")

    def union_ids(self, tok):
        return digit_ids(tok)

    def template_row(self, tok, b, t, p):
        """7 token ids for world b's statement t; STATE slot (offset 5) is pad (computed)."""
        if not hasattr(tok, "_saga_dot"):
            tok._saga_dot = tok(".", add_special_tokens=False)["input_ids"][0]
        di = digit_ids(tok)
        return [self.subject_ids(tok)[p.themes[b]],   # 0 subject
                self._verb_id(tok, self.verb_token(b, t, p)),  # 1 op verb (premise)
                di[p.operands[b][t]],                 # 2 operand (premise)
                self.resource_ids(tok)[p.themes[b]],  # 3 resource (flavor)
                self.conn_ids(tok)[p.conns[b][t]],    # 4 connective (flavor)
                tok.pad_token_id,                     # 5 STATE (computed; pad-overwritten)
                tok._saga_dot]                        # 6 period

    def state_states(self, p):
        return [list(p.states[b]) for b in range(self.K)]   # [K][m]

    def prompt(self, p: SagaInst) -> str:
        worlds = "; ".join(
            f"Trader {b+1} is the {_SAGA_THEMES[p.themes[b]][0]}, starting with {p.starts[b]} "
            f"{_SAGA_THEMES[p.themes[b]][1]}" for b in range(self.K))
        return (
            f"{self.K} traders keep separate running tallies (each 0-9, wrapping mod 10). Operations: "
            f"'gained'/'earned' k adds k; 'lost'/'spent'/'drained' k subtracts k; 'scaled' k multiplies "
            f"by k -- all mod 10. {worlds}. Each takes {self.m} steps. Only AFTER your thinking, one "
            f"trader (by NUMBER) is named and you must give their FINAL tally (a single digit).\n\n"
            f"Reason inside <think> </think> -- for each trader in turn, write one sentence per step: "
            f"'the <trader> <op> <k> <goods>, now <tally>.'")

    def query(self, p: SagaInst) -> str:
        return f"What is trader number {p.q + 1}'s final tally? "

    def answer(self, p: SagaInst) -> int:
        return p.states[p.q][-1]

    def all_queries(self, p: SagaInst):
        return [(f"What is trader number {b + 1}'s final tally? ", p.states[b][-1]) for b in range(self.K)]

    def probe_targets(self, p: SagaInst) -> dict[str, int]:
        return {f"{_SAGA_THEMES[p.themes[b]][0]}_{t+1}": p.states[b][t] for b in range(self.K) for t in range(self.m)}


TASKS = {t.name: t for t in (ParallelSelect, ChaseSelect, StepSelect, CoinTrack, AggSelect,
                             HypTrack, Diffuse, Journeys, Tales, Sagas)}


def make_task(name: str, **kwargs):
    return TASKS[name](**kwargs)