Learning Rules from Examples: How Can a Model Discover a Hidden Rule?

import os
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")   # keep the output readable

import json, math, random, statistics, itertools, zlib, sys, subprocess, hashlib, datetime
from collections import Counter, defaultdict
from pathlib import Path

# ── The models to run the experiments on ────────────────────────────────────
# Write the ones you want. A local name needs its weights already in your
# Hugging Face cache, an API name needs its key in the environment, and anything
# unreachable is reported and skipped rather than silently dropped.
MODELS = [
    "qwen3-8b",             # mlx-community/Qwen3-8B-4bit           ~4.5 GB
    "qwen3-30b-a3b",        # mlx-community/Qwen3-30B-A3B-...-4bit   ~17 GB
    "gpt-5.6-luna",         # OPENAI_API_KEY
    # "qwen3-32b",          # mlx-community/Qwen3-32B-4bit           ~18 GB
    # "gpt-5.6-terra",      # OPENAI_API_KEY
]

# Model replies are cached on disk so that re-running the notebook after an edit
# costs nothing. Delete the directory to force fresh calls; set CACHE_ENABLED to
# False to never write one.
CACHE, CACHE_ENABLED = Path(".lab_cache"), True

print("executed:", datetime.date.today().isoformat())

def identity(x):       return list(x)
def reverse(x):        return list(x[::-1])
def rotate_left(x):    return list(x[1:]) + list(x[:1])
def rotate_right(x):   return list(x[-1:]) + list(x[:-1])
def duplicate(x):      return list(x) + list(x)
def every_second(x):   return list(x[::2])              # 0-based: positions 0, 2, 4
def increment(x):      return [v + 1 for v in x]
def decrement(x):      return [v - 1 for v in x]
def double(x):         return [2 * v for v in x]
def swap_ends(x):      return [x[-1]] + list(x[1:-1]) + [x[0]] if len(x) > 1 else list(x)
def sort_asc(x):       return sorted(x)
def sort_desc(x):      return sorted(x, reverse=True)
def keep_last(x):      return list(x[-1:])
def keep_max(x):       return [max(x)] if x else []
def pairwise_swap(x):
    out = [v for a, b in zip(x[::2], x[1::2]) for v in (b, a)]
    return out + (list(x[-1:]) if len(x) % 2 else [])

OPS = {f.__name__: f for f in (identity, reverse, rotate_left, rotate_right, duplicate,
                               every_second, increment, decrement, double, swap_ends,
                               sort_asc, sort_desc, keep_last, keep_max, pairwise_swap)}

x = [3, 1, 4, 1, 5]
for name in ("reverse", "every_second", "pairwise_swap", "keep_max"):
    print(f"  {name:<15}{x} -> {OPS[name](x)}")

def run(program, x, ops=OPS):
    out = list(x)
    for name in program:
        out = ops[name](out)
    return out

def as_callable(program, ops=OPS):
    return lambda x: run(program, x, ops)

def make_task(seed, program, n_demos=3, n_tests=1, lengths=(3, 4, 5, 6), ops=OPS):
    """A hidden program, some demonstrations of it, and a held-out query."""
    rng = random.Random(seed)
    fn = as_callable(program, ops)
    xs = [[rng.randint(1, 9) for _ in range(lengths[i % len(lengths)])]
          for i in range(n_demos + n_tests)]
    return {"program": tuple(program), "fn": fn,
            "demos": [(v, fn(v)) for v in xs[:n_demos]],
            "tests": [(v, fn(v)) for v in xs[n_demos:]]}

def random_program(rng, depth, ops=OPS):
    names = [n for n in ops if n != "identity"]
    return tuple(rng.choice(names) for _ in range(depth))

t = make_task(2649, ("reverse",))
print("hidden:", " -> ".join(t["program"]))
for a, b in t["demos"]:
    print(f"    {a} -> {b}")

for seed, prim in ((2649, "reverse"), (5501, "rotate_left")):
    first, again = make_task(seed, (prim,)), make_task(seed, (prim,))
    print(f"seed {seed}, hidden rule {prim}   regenerates identically: "
          f"{first['demos'] == again['demos'] and first['tests'] == again['tests']}")

k = len(OPS)
for d in range(1, 7):
    print(f"  depth {d}:  {k**d:>12,} programs")

# Seven probe inputs, deliberately varied in length, order, ties and sign.
PROBES = ([1, 2, 3], [4, 7, 9, 2], [5, 8], [3, 3, 1, 9, 6], [7],
          [2, 6, 4, 8, 1, 5], [-2, 5, 0])
MAX_LEN = 64                      # `duplicate` doubles length; cap the frontier

def signature(fn, probes=PROBES):
    """A program's behavioural fingerprint, or None if it runs out of bounds."""
    out = []
    for p in probes:
        try:
            y = fn(list(p))
        except Exception:
            return None
        if len(y) > MAX_LEN:
            return None
        out.append(tuple(y))
    return tuple(out)

def enumerate_signatures(max_depth, ops):
    """Bottom-up, level by level, keeping the shortest program per signature."""
    names = list(ops)
    seen, frontier, rows = {}, [()], []
    for depth in range(1, max_depth + 1):
        nxt = []
        for prefix in frontier:
            for name in names:
                cand = prefix + (name,)
                sig = signature(as_callable(cand, ops))
                if sig is None or sig in seen:
                    continue
                seen[sig] = cand
                nxt.append(cand)
        frontier = nxt
        rows.append((depth, len(names) ** depth, len(seen)))
    return rows, seen

rows, seen = enumerate_signatures(6, OPS)
print(f"{'depth':>6}{'programs':>14}{'signatures':>13}{'overcount':>11}")
for d, syn, sig in rows:
    print(f"{d:>6}{syn:>14,}{sig:>13,}{syn/sig:>10.0f}x")

groups = defaultdict(list)
frontier = [()]
for _ in range(3):
    nxt = []
    for prefix in frontier:
        for name in OPS:
            cand = prefix + (name,)
            sig = signature(as_callable(cand))
            if sig is None:
                continue
            groups[sig].append(cand)
            nxt.append(cand)
    frontier = nxt

multi = {s: g for s, g in groups.items() if len(g) > 1}
print(f"signature groups holding more than one program: {len(multi)}")

rng = random.Random(4242)
IN_DIST = [[rng.randint(1, 9) for _ in range(rng.randint(3, 6))] for _ in range(400)]
LONGER  = [[rng.randint(1, 9) for _ in range(rng.randint(7, 12))] for _ in range(300)]

def sig_on(prog, pool):
    try:
        return tuple(tuple(run(prog, list(x))) for x in pool)
    except Exception:
        return None

for label, pool in (("same lengths as the demos (3-6)", IN_DIST), ("longer inputs (7-12)", LONGER)):
    split = sum(1 for g in multi.values()
                if any(sig_on(o, pool) != sig_on(g[0], pool) for o in g[1:]))
    print(f"  separable on {label:<32} {split:>3}/{len(multi)} = {split/len(multi):.2%}")

for sig, g in multi.items():
    a, b = g[0], g[1]
    for x in LONGER:
        if run(a, list(x)) != run(b, list(x)):
            print("merged as identical:", " -> ".join(a), " vs ", " -> ".join(b))
            print("  on probes: indistinguishable")
            print(f"  on {x}:")
            print(f"    {run(a, list(x))}")
            print(f"    {run(b, list(x))}")
            break
    else:
        continue
    break

def permutation_op(perm):
    """Output position i takes input element perm[i]. Defined for one arity only."""
    perm = tuple(perm)
    def op(x):
        if len(x) != len(perm):
            raise ValueError("wrong arity")
        return [x[i] for i in perm]
    return op

def permutation_hypotheses(n):
    return [(f"perm{list(p)}", permutation_op(p)) for p in itertools.permutations(range(n))]

def permutation_task(seed, n=5, n_demos=6, alphabet=3):
    rng = random.Random(seed)
    perm = list(range(n)); rng.shuffle(perm)
    op = permutation_op(perm)
    xs = [[rng.randint(1, alphabet) for _ in range(n)] for _ in range(n_demos + 1)]
    return {"perm": perm, "fn": op,
            "demos": [(v, op(v)) for v in xs[:-1]], "query": xs[-1]}

def consistent(hyps, examples):
    """Mitchell's version space: everything still compatible with the data."""
    out = []
    for label, fn in hyps:
        ok = True
        for x, y in examples:
            try:
                if fn(list(x)) != list(y):
                    ok = False; break
            except Exception:
                ok = False; break
        if ok:
            out.append((label, fn))
    return out

def bits(hyps):
    return math.log2(len(hyps)) if hyps else float("nan")

hyps = permutation_hypotheses(5)
task = permutation_task(11, n=5, n_demos=6, alphabet=3)

live = hyps
print(f"start: {len(live):>3} hypotheses  {bits(live):.2f} bits")
for x, y in task["demos"]:
    live = consistent(live, [(x, y)])
    print(f"  {x} -> {y}   {len(live):>3} left  {bits(live):.2f} bits")
print("\nhidden operator was perm" + str(task["perm"]))

def floor_curve(alphabet, arity=5, seeds=200, max_k=6):
    hyps = permutation_hypotheses(arity)
    curves = []
    for s in range(seeds):
        t = permutation_task(1000 * alphabet + s, n=arity, n_demos=max_k, alphabet=alphabet)
        live, row = hyps, []
        for (x, y) in t["demos"]:
            live = consistent(live, [(x, y)])
            row.append(bits(live))
        curves.append(row)
    return [statistics.fmean(c[i] for c in curves) for i in range(max_k)]

print(f"{'k':>3}" + "".join(f"{'a=' + str(a):>9}" for a in (2, 3, 4, 9)))
cur = {a: floor_curve(a) for a in (2, 3, 4, 9)}
for i in range(6):
    print(f"{i+1:>3}" + "".join(f"{cur[a][i]:>9.2f}" for a in (2, 3, 4, 9)))

def teaches(combo, hyps):
    sigs = set()
    for _, fn in hyps:
        s = tuple(tuple(fn(list(x))) for x in combo)
        if s in sigs:
            return False
        sigs.add(s)
    return True

def teaching_dimension(n, a, trials=3000):
    hyps = permutation_hypotheses(n)
    universe = list(itertools.product(range(1, a + 1), repeat=n))
    r = random.Random(0)
    for k in range(1, 5):
        cands = [(u,) for u in universe] if k == 1 else \
                [tuple(r.choice(universe) for _ in range(k)) for _ in range(trials)]
        if any(teaches(c, hyps) for c in cands):
            return k
    return None

print(f"{'arity':>6}{'alphabet':>10}{'measured':>10}{'ceil(log_a n)':>15}")
for n in (3, 4, 5, 6):
    for a in (2, 3, 9):
        m = teaching_dimension(n, a)
        pred = max(1, math.ceil(math.log(n, a)))
        print(f"{n:>6}{a:>10}{m:>10}{pred:>15}{'   ok' if m == pred else '   MISMATCH'}")

hyps5 = permutation_hypotheses(5)
universe = list(itertools.product(range(1, 3), repeat=5))
r = random.Random(0)
found = None
for _ in range(20000):
    combo = tuple(r.choice(universe) for _ in range(3))
    if teaches(combo, hyps5):
        found = combo
        break
for row in found:
    print("   ", list(row))
print("\n  position codes read down the columns:")
for pos in range(5):
    print(f"    position {pos}: {''.join(str(row[pos]) for row in found)}")

def eig(hyps, x):
    """Expected information gain: with deterministic hypotheses this is exactly
    the entropy of the partition the query induces on the survivors."""
    groups = defaultdict(int)
    for _, fn in hyps:
        try:
            groups[tuple(fn(list(x)))] += 1
        except Exception:
            groups[None] += 1
    n = len(hyps)
    return -sum((c / n) * math.log2(c / n) for c in groups.values()) if n > 1 else 0.0

_qr = random.Random(0)
cands = [[_qr.randint(1, 9) for _ in range(5)] for _ in range(300)]
live8 = consistent(hyps5, task["demos"][:1])
best = max(cands, key=lambda x: eig(live8, x))
worst = min(cands, key=lambda x: eig(live8, x))
print(f"survivors after one demonstration: {len(live8)}  ({bits(live8):.2f} bits)")
print(f"  best query  {best}  -> {eig(live8, best):.2f} bits")
print(f"  worst query {worst}  -> {eig(live8, worst):.2f} bits")
print(f"  a query everyone agrees on is worth exactly {eig(live8, [1,1,1,1,1]):.2f} bits")

def extract_json(text):
    """First well-formed JSON object anywhere in the reply — survives markdown
    fences and chatter. No regex; the decoder decides where the object ends."""
    dec = json.JSONDecoder()
    for i, ch in enumerate(text):
        if ch != "{":
            continue
        try:
            obj, _ = dec.raw_decode(text[i:])
        except ValueError:
            continue
        if isinstance(obj, dict):
            return obj
    return None

def compile_answer(ans, ops=OPS):
    """Turn what the model said into something executable.

    A named chain of primitives is the only representation accepted, and that
    is deliberate twice over. It keeps the answer inside the language the model
    was actually offered — validating against a wider dictionary would accept a
    primitive the prompt had removed and quietly destroy the ablation later on.
    And it means nothing the model writes is ever executed as code: the reply is
    a list of names, checked against `ops`, and the interpreter is the one
    written at the top of this lab. Running model-authored Python behind a
    trimmed `__builtins__` is not a sandbox, and there is nothing here it would
    buy.
    """
    if not isinstance(ans, dict):
        return None, "no-answer"
    names = ans.get("program")
    if isinstance(names, list) and names and all(str(n) in ops for n in names):
        prog = tuple(str(n) for n in names)
        return (lambda x: run(prog, list(x), ops)), "program"
    return None, "uncompilable"

def extensionally_equal(fn, truth, probes=PROBES):
    for pr in probes:
        try:
            if list(fn(list(pr))) != list(truth(list(pr))):
                return False
        except Exception:
            return False
    return True

replies = [
    '{"rule": "move the first element to the end", "program": ["rotate_left"]}',
    'Sure! ```json\n{"rule":"reverse then duplicate","program":["reverse","duplicate"]}\n```',
    '{"rule": "shift left", "python": "def f(x):\\n    return x[1:] + x[:1]"}',
    '{"rule": "move the last element to the front, no wait, the first to the end",'
    ' "program": ["rotate_left"]}',
]
truth = OPS["rotate_left"]
for r in replies:
    ans = extract_json(r) or {}
    fn, how = compile_answer(ans)
    ok = fn is not None and extensionally_equal(fn, truth)
    print(f"  compiled from {how:<9} equals rotate_left: {ok}   <- {ans.get('rule','')[:44]}")

SYSTEM = ("You infer a hidden transformation on lists of integers from examples. "
          "You are precise, you commit to an answer, and you reply with JSON only.")

def menu(ops=OPS):
    return "\n".join(f"  {n}" for n in ops)

def identification_prompt(demos, query, ops=OPS):
    shown = "\n".join(f"  {list(a)} -> {list(b)}" for a, b in demos)
    return (f"A hidden function phi maps a list of integers to a list of integers.\n"
            f"Observed:\n{shown}\n\nWhat is phi({list(query)})?\n\n"
            f"phi is a chain of one or more of these primitives, applied left to right:\n"
            f"{menu(ops)}\n\n"
            'Answer with one JSON object and nothing else:\n'
            '{"rule": "<one sentence>", "program": [<primitive names>], '
            '"output": [<the transformed list>]}')

def ask(model, prompt, max_tokens=900):
    return learner(model).ask(SYSTEM, prompt, max_tokens=max_tokens)

def identify(model, primitive, seed):
    """One task, one reply, graded on all four axes separately."""
    task = make_task(seed, (primitive,), n_demos=3, n_tests=1)
    query, truth = task["tests"][0]
    reply = ask(model, identification_prompt(task["demos"], query))
    ans = extract_json(reply.text)
    fn, how = compile_answer(ans or {})
    stated_out = (ans or {}).get("output")
    return {
        "primitive": primitive,
        "stated_program": (ans or {}).get("program"),
        # 1. did the reply come back as usable JSON at all
        "parsed": ans is not None,
        # 2. is the list it wrote out for the query right
        "output_correct": (list(stated_out) == list(truth)
                           if isinstance(stated_out, list) else False),
        # 3. is the rule it named the hidden function, on the probe set
        "rule_correct": fn is not None and extensionally_equal(fn, task["fn"]),
        # 4. checkable without an answer key: does its own rule reproduce the demos?
        "self_consistent": fn is not None and all(
            list(fn(list(a))) == list(b) for a, b in task["demos"]),
    }, task

E1 = {}
targets = [n for n in OPS if n != "identity"]
for m in MODELS:
    E1[m] = [identify(m, n, zlib.crc32(f"lab:{n}".encode()) % 10_000)[0] for n in targets]

if not MODELS:
    print("no learner available — skipping")
else:
    print(f"{len(targets)} primitives, one task each, one reply graded four ways\n")
    print(f"{'model':<16}{'parsed':>9}{'output':>9}{'rule':>8}{'fits demos':>13}"
          f"{'fits but wrong':>16}")
    for m, got in E1.items():
        n = len(got)
        col = lambda k: sum(bool(r[k]) for r in got) / n
        both = sum(1 for r in got if r["self_consistent"] and not r["rule_correct"])
        print(f"{m:<16}{col('parsed'):>9.2f}{col('output_correct'):>9.2f}"
              f"{col('rule_correct'):>8.2f}{col('self_consistent'):>13.2f}{both:>16}")
    print("\n  " + "  ".join(f"{r['primitive']}{'+' if r['rule_correct'] else '-'}"
                              for r in E1[MODELS[0]]))

if not MODELS:
    print("no learner available — skipping")
else:
    got = E1[MODELS[0]]
    miss = next((r for r in got if not r["rule_correct"] and r["stated_program"]), None)
    if miss is None:
        print("no compilable miss this run — try another model or more primitives")
    else:
        task = make_task(zlib.crc32(f"lab:{miss['primitive']}".encode()) % 10_000,
                         (miss["primitive"],))
        fn = (lambda prog: (lambda x: run(tuple(prog), list(x))))(miss['stated_program'])
        print(f"hidden rule: {miss['primitive']} | model said: {miss['stated_program']}")
        for x, y in task['demos']:
            out = fn(list(x))
            print(f"  {x} -> shown {y}, its rule gives {out}"
                  f"   {'ok' if out == y else '<-- breaks here'}")

def permutation_prompt(demos, query, arity, explicit=True):
    shown = "\n".join(f"  {list(a)} -> {list(b)}" for a, b in demos)
    convention = (f" Positions are 0-based, so p is a permutation of 0..{arity-1}."
                  if explicit else "")
    return (f"A hidden operator Psi rearranges lists of exactly {arity} integers: output "
            f"position i takes input position p[i], for some unknown permutation p."
            f"{convention}\nObserved:\n{shown}\n\nWhat is Psi({list(query)})?\n\n"
            'Answer with one JSON object and nothing else:\n'
            '{"permutation": [...], "output": [...]}')

def displacement(perm):
    """Distance from the identity — a simplicity proxy with no theory behind it
    beyond 'moving fewer things is a smaller change'."""
    return sum(1 for i, v in enumerate(perm) if i != v)

def evidence_trial(model, k, alphabet, arity=5, seeds=6):
    """One cell of the sweep: k demonstrations over an alphabet of that size."""
    hyps = permutation_hypotheses(arity)
    rows = []
    for s in range(seeds):
        t = permutation_task(7000 + 100 * k + s, n=arity, n_demos=k, alphabet=alphabet)
        live = consistent(hyps, t["demos"])
        live_perms = [json.loads(l[len("perm"):]) for l, _ in live]
        obj = extract_json(ask(
            model, permutation_prompt(t["demos"], t["query"], arity), 900).text) or {}
        perm = obj.get("permutation")
        valid = (isinstance(perm, list) and len(perm) == arity
                 and sorted(perm) == list(range(arity)))
        rows.append({
            "identifiable": len(live) == 1,
            "bits": bits(live),
            "exact": bool(valid and list(perm) == t["perm"]),
            # a weaker, fairer question when the evidence cannot pin it down:
            # did it at least name something the data still allows?
            "consistent_pick": bool(valid and list(perm) in live_perms),
            "named_disp": displacement(perm) if valid else None,
            "min_disp": min(displacement(p) for p in live_perms),
            # what picking uniformly among the survivors would score: the share
            # of survivors that are already least-moved. Without this the rate
            # below is a preference with nothing to be a preference over.
            "chance": (sum(1 for p in live_perms
                           if displacement(p) == min(displacement(q) for q in live_perms))
                       / len(live_perms)),
        })
    return rows

if not MODELS:
    print("no learner available — skipping")
else:
    print(f"{'model':<15}{'k':>3}{'alphabet':>10}{'mean bits left':>16}"
          f"{'exact':>8}{'consistent':>12}")
    E2 = {}
    for m in MODELS:
        E2[m] = []
        for alphabet, k in ((9, 3), (9, 6), (3, 3), (3, 6)):
            rows = evidence_trial(m, k, alphabet)
            E2[m] += rows
            n = len(rows)
            print(f"{m:<15}{k:>3}{alphabet:>10}"
                  f"{statistics.fmean(r['bits'] for r in rows):>16.2f}"
                  f"{sum(r['exact'] for r in rows) / n:>8.2f}"
                  f"{sum(r['consistent_pick'] for r in rows) / n:>12.2f}")

if not MODELS:
    print("no learner available — skipping")
else:
    print(f"{'model':<15}{'ambiguous & consistent':>24}{'least-moved':>13}"
          f"{'rate':>7}{'chance':>8}")
    for m in MODELS:
        amb = [r for r in E2[m] if not r["identifiable"] and r["consistent_pick"]]
        hit = sum(1 for r in amb if r["named_disp"] == r["min_disp"])
        rate = hit / len(amb) if amb else float("nan")
        chance = statistics.fmean(r["chance"] for r in amb) if amb else float("nan")
        cells = (f"{rate:>7.2f}{chance:>8.2f}" if amb else f"{'-':>7}{'-':>8}")
        print(f"{m:<15}{len(amb):>24}{hit:>13}{cells}")

PROBE = [9, 2, 4]                    # last element is deliberately not the maximum

# Every depth-<=2 program, as (label, callable). Small enough to write out.
DEPTH2 = [(" -> ".join(p), as_callable(p)) for p in
          [(a,) for a in OPS] + [(a, b) for a in OPS for b in OPS]]

def shortcut_demos(seed, aligned=True, n_demos=3):
    """True rule: keep_last. aligned=True sorts each input ascending, so the last
    element is also the maximum and the two rules cannot be told apart. aligned=False
    is the control — descending inputs, where they already disagree in the demos."""
    rng = random.Random(seed)
    demos = []
    for _ in range(n_demos):
        body = sorted(rng.sample(range(1, 10), rng.choice([3, 4])))
        x = body if aligned else body[::-1]
        demos.append((x, [x[-1]]))
    return demos

def shortcut_prompt(demos, condition, probe=None):
    probe = PROBE if probe is None else probe
    shown = "\n".join(f"  {list(a)} -> {list(b)}" for a, b in demos)
    head = (f"A hidden function phi maps a list of integers to a list of integers.\n"
            f"Observed:\n{shown}\n\nWhat is phi({probe})?\n\n"
            f"phi is a chain of one or more of these primitives, applied left to right:\n"
            f"{menu()}\n\n")
    if condition == "direct":
        return head + ('Answer with one JSON object and nothing else:\n'
                       '{"output": [<the transformed list>]}')
    if condition == "rule_first":
        return head + ("State the rule before answering.\n\n"
                       'Answer with one JSON object and nothing else:\n'
                       '{"rule": "<one sentence>", "program": [<primitive names>], '
                       '"output": [<the transformed list>]}')
    return head + ("First list every rule consistent with ALL the observations, not just "
                   "the first one you think of. Then answer.\n\n"
                   'Answer with one JSON object and nothing else:\n'
                   '{"candidates": [<one entry per consistent rule>], '
                   '"rule": "<the one you commit to>", "output": [<the transformed list>]}')

REPAIR_PROBE = [8, 3, 5]              # a second input where last is not the maximum

def verdict(out, probe=None):
    probe = PROBE if probe is None else probe
    if out == [probe[-1]]:  return "true rule"
    if out == [max(probe)]: return "shortcut"
    return "other"

demos = shortcut_demos(1, aligned=True)
live = consistent(DEPTH2, demos)
print("demonstrations:", [[x, y] for x, y in demos])
print(f"depth-<=2 behaviours consistent with them: {len(live)}   ({bits(live):.2f} bits)")
print("  including:", ",  ".join(l for l, _ in live[:5]))
print(f"\non the test {PROBE}:  keep_last -> {[PROBE[-1]]},  keep_max -> {[max(PROBE)]}")

CONDITIONS = ("direct", "rule_first", "hypotheses_first")

def shortcut_trial(model, condition, aligned, seeds=4):
    """The fourth condition, `repair`, is different in kind: the learner is shown
    the counterexample it just got wrong, added to the evidence, and asked again
    on a fresh input. It measures whether being corrected once is enough."""
    tally = Counter()
    for s in range(seeds):
        d = shortcut_demos(s, aligned=aligned)
        if condition == "repair":
            d = list(d) + [(PROBE, [PROBE[-1]])]      # hand it the counterexample
            probe, asked = REPAIR_PROBE, "rule_first"
        else:
            probe, asked = PROBE, condition
        obj = extract_json(ask(model, shortcut_prompt(d, asked, probe), 1500).text) or {}
        out = obj.get("output")
        tally[verdict(out, probe) if isinstance(out, list) else "other"] += 1
    return tally

if not MODELS:
    print("no learner available — skipping")
else:
    print(f"{'model':<15}{'asked':<19}{'demos':<10}{'true rule':>10}{'shortcut':>10}{'other':>7}")
    for m in MODELS:
        rows = [(c, True) for c in CONDITIONS] + [("repair", True), ("rule_first", False)]
        for cond, aligned in rows:
            t = shortcut_trial(m, cond, aligned=aligned)
            n = sum(t.values())
            print(f"{m:<15}{cond:<19}{'baited' if aligned else 'control':<10}"
                  f"{t['true rule']/n:>10.2f}{t['shortcut']/n:>10.2f}{t['other']/n:>7.2f}")

# One finite, shared query domain. The oracle maximises over exactly the inputs
# the learner is allowed to pick from, so "regret" is a real regret — with an
# open-ended input space the best-in-a-sample is not an oracle at all, and a
# learner can beat it, which is precisely what happened once in the note's sweep.
QUERY_DOMAIN = [[a, b, c] for a in (1, 2, 5, 9) for b in (1, 3, 7) for c in (2, 4, 8)] + \
               [[a, b, c, d] for a in (1, 6) for b in (2, 9) for c in (3, 5) for d in (4, 7)]

def query_prompt(labels, domain=QUERY_DOMAIN):
    """Both the oracle and the learner choose by index out of the same list.

    Describing the domain without printing it, and then accepting any list the
    learner writes, would leave the oracle restricted and the learner free —
    which is not a comparison at all. Numbering the queries removes the gap by
    construction: an answer is a `Q` number or it is nothing.
    """
    listing = "\n".join(f"  H{i+1}: {l}" for i, l in enumerate(labels))
    queries = "\n".join(f"  Q{i}: {x}" for i, x in enumerate(domain))
    return ("A hidden function is exactly one of these candidates:\n" + listing +
            "\n\nYou may observe the hidden function's output on ONE input, and it "
            "must be one of these:\n" + queries +
            "\n\nChoose the one whose output would rule out as many candidates as "
            "possible.\n\n"
            'Answer with one JSON object and nothing else:\n'
            '{"query_id": <the Q number>, "reason": "<one sentence>"}')

def parse_query_id(text, domain=QUERY_DOMAIN):
    """Accept 7, "7" and "Q7" alike.

    A reply that names the right query in a slightly different shape is a
    correct answer. Insisting on a bare integer manufactures exactly the kind of
    capability gap the traps below are about — and it did: on the first run of
    this section one model answered `"Q36"` every single time and scored zero
    bits in all twelve cells.
    """
    raw = (extract_json(text) or {}).get("query_id")
    if isinstance(raw, bool):
        return None
    if isinstance(raw, str):
        raw = raw.strip().lstrip("Qq")
    try:
        i = int(raw)
    except (TypeError, ValueError):
        return None
    return domain[i] if 0 <= i < len(domain) else None

QPROBE = QUERY_DOMAIN[::7]

def _qsig(fn):
    try:
        return tuple(tuple(fn(list(x))) for x in QPROBE)
    except Exception:
        return None

def candidate_pool():
    """One program per distinct QPROBE signature, as (label, callable).

    Distinct on the probe set, which is not the same as behaviourally distinct —
    the section on probe collisions measured exactly how much that costs.
    """
    sigs = {}
    for prog in [(a,) for a in OPS] + [(a, b) for a in OPS for b in OPS]:
        fn = as_callable(prog)
        sg = _qsig(fn)
        if sg is not None and sg not in sigs:
            sigs[sg] = (" -> ".join(prog), fn)
    return list(sigs.values())

POOL = candidate_pool()

def hypothesis_set(rng, size, condition):
    """Two ways to build a candidate set, and the difference is the experiment.

    `diverse` samples at random: most pairs already disagree on most inputs, so
    almost any query separates almost everything and a learner picking blindly
    scores near the ceiling. Such a task cannot tell you whether the learner is
    choosing. `confusable` ranks the pool by how often each program agrees with
    a random anchor and keeps the closest — there the choice of query decides
    how much you learn, which is the version the note reports.
    """
    if condition == "diverse":
        return rng.sample(POOL, size)
    anchor = POOL[rng.randrange(len(POOL))]
    a_sig = _qsig(anchor[1])
    agreement = lambda e: sum(u == v for u, v in zip(a_sig, _qsig(e[1])))
    return sorted(POOL, key=lambda e: (-agreement(e), len(e[0])))[:size]

def query_design(model, size, seed, condition="confusable"):
    """Returns (ceiling, oracle bits, mean bits of a blind pick, model bits, query)."""
    rng = random.Random(seed)
    hyps = hypothesis_set(rng, size, condition)
    oracle = max(eig(hyps, x) for x in QUERY_DOMAIN)
    chance = statistics.fmean(eig(hyps, x) for x in QUERY_DOMAIN)
    q = parse_query_id(ask(model, query_prompt([l for l, _ in hyps]), 900).text)
    return bits(hyps), oracle, chance, (eig(hyps, q) if q else 0.0), q

if not MODELS:
    print("no learner available — skipping")
else:
    print(f"{'model':<15}{'set':<12}{'n':>4}{'ceiling':>9}{'oracle':>8}{'model':>7}"
          f"{'blind':>7}{'regret':>8}   its query")
    for m in MODELS:
        for cond in ("diverse", "confusable"):
            for size in (5, 10, 20):
                ceil_, o, c, g, q = query_design(m, size, seed=size, condition=cond)
                print(f"{m:<15}{cond:<12}{size:>4}{ceil_:>9.2f}{o:>8.2f}{g:>7.2f}"
                      f"{c:>7.2f}{o - g:>8.2f}   {q}")

def sequential(model, size=20, seed=41, max_steps=6):
    """Query, observe, shrink, repeat. Compare against greedy and blind schedules."""
    rng = random.Random(seed)
    hyps = hypothesis_set(rng, size, "confusable")
    truth_label, truth_fn = hyps[rng.randrange(len(hyps))]

    def schedule(pick):
        live, steps = list(hyps), 0
        while len(live) > 1 and steps < max_steps:
            q = pick(live)
            if q is None:
                break
            live = consistent(live, [(q, list(truth_fn(list(q))))])
            steps += 1
        return steps, len(live)

    greedy = schedule(lambda live: max(QUERY_DOMAIN, key=lambda x: eig(live, x)))
    blind = schedule(lambda live: QUERY_DOMAIN[rng.randrange(len(QUERY_DOMAIN))])

    live, steps, history = list(hyps), 0, []
    while len(live) > 1 and steps < max_steps:
        prompt = query_prompt([l for l, _ in live])
        if history:
            prompt += "\n\nYou already observed:\n" + "\n".join(
                f"  {q} -> {y}" for q, y in history)
        q = parse_query_id(ask(model, prompt, 900).text)
        if q is None:
            break
        gain = eig(live, q)
        y = list(truth_fn(list(q)))
        live = consistent(live, [(q, y)])
        history.append((list(q), y))
        steps += 1
        if gain == 0.0 and steps >= 3:
            break            # it is asking uninformative questions; stop paying
    return {"truth": truth_label, "model": (steps, len(live)),
            "greedy": greedy, "blind": blind, "history": history}

if not MODELS:
    print("no learner available — skipping")
else:
    print(f"{'model':<15}{'questions asked':>16}{'greedy':>9}{'blind':>7}"
          f"{'left over':>11}")
    for m in MODELS:
        r = sequential(m)
        print(f"{m:<15}{r['model'][0]:>16}{r['greedy'][0]:>9}{r['blind'][0]:>7}"
              f"{r['model'][1]:>11}")
        for q, y in r["history"]:
            print(f"                 asked {str(q):<16} -> {y}")

def depth_trial(model, depth, n=5):
    """Hide a random chain of exactly `depth` primitives; count exact recoveries."""
    ok = 0
    for i in range(n):
        seed = 4000 + 100 * depth + i
        prog = random_program(random.Random(seed), depth)
        t = make_task(seed, prog, n_demos=4, n_tests=1)
        obj = extract_json(ask(
            model, identification_prompt(t["demos"], t["tests"][0][0]), 1500).text) or {}
        fn, _ = compile_answer(obj)
        ok += bool(fn is not None and extensionally_equal(fn, t["fn"]))
    return ok / n

if not MODELS:
    print("no learner available — skipping")
else:
    depths = (1, 2, 3, 4)
    print(f"{'model':<15}" + "".join(f"{'depth ' + str(d):>10}" for d in depths))
    for m in MODELS:
        print(f"{m:<15}" + "".join(f"{depth_trial(m, d):>10.2f}" for d in depths))

PAIRS = [("reverse", "increment"), ("sort_asc", "rotate_left"), ("double", "keep_last")]

def order_trial(model, a, b, n=2):
    """Accuracy on (a then b) and on (b then a), same pair, same seeds."""
    out = {}
    for tag, prog in (("ab", (a, b)), ("ba", (b, a))):
        ok = 0
        for i in range(n):
            seed = zlib.crc32(f"{a}|{b}|{tag}|{i}".encode()) % 10_000
            t = make_task(seed, prog, n_demos=4, n_tests=1)
            obj = extract_json(ask(
                model, identification_prompt(t["demos"], t["tests"][0][0]), 1500).text) or {}
            fn, _ = compile_answer(obj)
            ok += bool(fn is not None and extensionally_equal(fn, t["fn"]))
        out[tag] = ok / n
    return out

if not MODELS:
    print("no learner available — skipping")
else:
    ORDER = {}
    for m in MODELS:
        ORDER[m] = [order_trial(m, a, b) for a, b in PAIRS]
        for (a, b), r in zip(PAIRS, ORDER[m]):
            print(f"  {m:<15}{a} -> {b:<14} {r['ab']:.2f}    reversed {r['ba']:.2f}")
        gaps = [abs(r["ab"] - r["ba"]) for r in ORDER[m]]
        print(f"  {m:<15}mean |gap| over {len(PAIRS)} pairs: {statistics.fmean(gaps):.3f}\n")

LENGTH_RULES = ("reverse", "every_second", "rotate_left", "duplicate", "pairwise_swap")
TEST_LENGTHS = (8, 16, 32, 64, 128)

def length_prompt(demos, query):
    shown = "\n".join(f"  {list(a)} -> {list(b)}" for a, b in demos)
    return (f"A hidden function phi maps a list of integers to a list of integers. "
            f"The same rule applies at every length.\nObserved:\n{shown}\n\n"
            f"Apply phi to this list of {len(query)} integers:\n  {list(query)}\n\n"
            f"phi is a chain of one or more of these primitives, applied left to right:\n"
            f"{menu()}\n\n"
            'Answer with one JSON object and nothing else:\n'
            '{"rule": "<one sentence>", "program": [<primitive names>], '
            '"output": [<the full transformed list>]}')

def length_case(rule, seed):
    """One task, reused at every length. The demonstrations are fixed and the
    query at length L is the first L elements of one long sequence, so moving
    down the column changes the length and nothing else. Drawing a fresh task
    per length would confound the thing being measured with task difficulty."""
    rng = random.Random(seed)
    fn = OPS[rule]
    demos = [(x, fn(x)) for x in
             [[rng.randint(1, 9) for _ in range(L)] for L in (3, 4, 5, 6)]]
    long_query = [rng.randint(1, 9) for _ in range(max(TEST_LENGTHS))]
    return fn, demos, long_query

def length_trial(model, rule, length, seed):
    """Three scores from one reply: the rule, the written-out list, and the
    model's own program executed by Python at that length."""
    fn, demos, long_query = length_case(rule, seed)
    query = long_query[:length]
    truth = fn(query)
    obj = extract_json(ask(model, length_prompt(demos, query), 3000).text) or {}
    got, _ = compile_answer(obj)
    written = obj.get("output")
    executed = None
    if got is not None:
        try:
            executed = list(got(list(query))) == list(truth)
        except Exception:
            executed = False
    return {
        "rule_correct": got is not None and extensionally_equal(got, fn),
        "written_correct": isinstance(written, list) and list(written) == list(truth),
        "program_at_length": executed,
    }

if not MODELS:
    print("no learner available — skipping")
else:
    print(f"{'model':<15}{'length':>8}{'rule':>8}{'wrote it out':>14}{'its program run':>17}")
    for m in MODELS:
        for L in TEST_LENGTHS:
            rows = [length_trial(m, r, L, zlib.crc32(f"e6:{r}".encode()) % 10_000)
                    for r in LENGTH_RULES]
            n = len(rows)
            ran = [r for r in rows if r["program_at_length"] is not None]
            print(f"{m:<15}{L:>8}"
                  f"{sum(r['rule_correct'] for r in rows) / n:>8.2f}"
                  f"{sum(r['written_correct'] for r in rows) / n:>14.2f}"
                  + (f"{sum(bool(r['program_at_length']) for r in ran) / len(ran):>17.2f}"
                     if ran else f"{'-':>17}"))

if not MODELS:
    print("no learner available — skipping")
else:
    t = make_task(7, ("rotate_left",))
    prompt = identification_prompt(t["demos"], t["tests"][0][0])
    for budget in (24, 900):
        r = ask(MODELS[0], prompt, max_tokens=budget)
        ok = compile_answer(extract_json(r.text) or {})[0] is not None
        shown = r.text.strip().replace("\n", " ")[:60] or "(empty)"
        print(f"  max_tokens={budget:<4} parsed={str(ok):<5} out={r.completion_tokens:<4} {shown}")

prog = "import sys; print([hash(('reverse', s)) % 1000 for s in range(3)])"
runs = [subprocess.run([sys.executable, "-c", prog], capture_output=True, text=True).stdout.strip()
        for _ in range(3)]
for r in runs:
    print("  ", r)
print("\nsame expression, three processes:", "IDENTICAL" if len(set(runs)) == 1 else "ALL DIFFERENT")
print("crc32 instead:", [zlib.crc32(f"reverse:{s}".encode()) % 1000 for s in range(3)])

# `permutation_prompt` is the one defined above; the `explicit` flag is the
# whole experiment — it is the only thing that changes between the two rows.
def permutation_trial(model, explicit, seeds=6, arity=5):
    one_indexed = valid = 0
    for i in range(seeds):
        t = permutation_task(100 + i, n=arity, n_demos=3, alphabet=3)
        prompt = permutation_prompt(t["demos"], t["query"], arity, explicit)
        obj = extract_json(ask(model, prompt).text) or {}
        perm = obj.get("permutation")
        if isinstance(perm, list) and len(perm) == arity:
            if sorted(perm) == list(range(arity)):
                valid += 1
            elif sorted(perm) == list(range(1, arity + 1)):
                one_indexed += 1
    return valid, one_indexed

if not MODELS:
    print("no learner available — skipping")
else:
    for m in MODELS:
        for explicit in (False, True):
            v, o = permutation_trial(m, explicit)
            label = "0-based stated explicitly" if explicit else "convention left implicit"
            print(f"  {m:<14} {label:<28} valid: {v}   1-indexed: {o}")

rng = random.Random(3)
true_rates = [rng.uniform(0.05, 0.75) for _ in range(14)]   # pairs differ in difficulty
CELL = 6                                                     # seeds per direction

def sample_gap(rates):
    gaps = []
    for p in rates:
        ab = sum(rng.random() < p for _ in range(CELL)) / CELL
        ba = sum(rng.random() < p for _ in range(CELL)) / CELL
        gaps.append(abs(ab - ba))
    return statistics.fmean(gaps)

observed = sample_gap(true_rates)          # a world with NO order effect
overall = statistics.fmean(true_rates)

def null_band(rates, trials=4000):
    r = random.Random(0)
    out = []
    for _ in range(trials):
        gaps = []
        for p in rates:
            ab = sum(r.random() < p for _ in range(CELL)) / CELL
            ba = sum(r.random() < p for _ in range(CELL)) / CELL
            gaps.append(abs(ab - ba))
        out.append(statistics.fmean(gaps))
    out.sort()
    return statistics.fmean(out), out[int(0.05 * len(out))], out[int(0.95 * len(out))]

print(f"observed mean gap in a world with no order effect: {observed:.3f}\n")
for label, rates in (("one global rate", [overall] * len(true_rates)),
                     ("conditioned on the pair", true_rates)):
    mean, lo, hi = null_band(rates)
    verdict = "outside the null" if not (lo <= observed <= hi) else "inside the null"
    print(f"  {label:<26} null {mean:.3f}  [{lo:.3f}, {hi:.3f}]  -> {verdict}")

def conditioned_null(pairs, cell, trials=4000, seed=0):
    """Null gap distribution with each pair's difficulty held fixed.

    The pooled rate for a pair is estimated from both directions together, so
    the simulated world has that pair's difficulty and no order effect at all.
    Everything the resampling produces is therefore noise by construction.
    """
    r = random.Random(seed)
    rates = [(p["ab"] + p["ba"]) / 2 for p in pairs]
    out = []
    for _ in range(trials):
        gaps = []
        for q in rates:
            ab = sum(r.random() < q for _ in range(cell)) / cell
            ba = sum(r.random() < q for _ in range(cell)) / cell
            gaps.append(abs(ab - ba))
        out.append(statistics.fmean(gaps))
    out.sort()
    return statistics.fmean(out), out[int(0.95 * len(out))]

if not MODELS:
    print("no learner available — skipping")
else:
    print(f"{'model':<15}{'observed gap':>14}{'null mean':>11}{'null 95th':>11}   verdict")
    for m in MODELS:
        obs = statistics.fmean(abs(p["ab"] - p["ba"]) for p in ORDER[m])
        mean, p95 = conditioned_null(ORDER[m], cell=2)
        verdict = "exceeds the null" if obs > p95 else "inside the null"
        print(f"{m:<15}{obs:>14.3f}{mean:>11.3f}{p95:>11.3f}   {verdict}")

def enumerate_programs(max_depth, ops=OPS):
    """Shortest program per behaviour signature, bottom-up, level by level."""
    names, seen, frontier, out = list(ops), {}, [()], []
    for _ in range(max_depth):
        nxt = []
        for prefix in frontier:
            for name in names:
                cand = prefix + (name,)
                sig = signature(as_callable(cand, ops))
                if sig is None or sig in seen:
                    continue
                seen[sig] = cand
                out.append(cand); nxt.append(cand)
        frontier = nxt
    return out

SPACE = [(pr, as_callable(pr)) for pr in enumerate_programs(3)]

def first_hit(order, demos):
    for i, idx in enumerate(order, 1):
        prog, fn = SPACE[idx]
        try:
            if all(list(fn(list(x))) == list(y) for x, y in demos):
                return i
        except Exception:
            pass
    return None

r = random.Random(5)
prog = random_program(r, 3)
t = make_task(5, prog, n_demos=4, n_tests=1)
fixed = list(range(len(SPACE)))
shuffled = fixed[:]; random.Random(1).shuffle(shuffled)

nb, nr = first_hit(fixed, t['demos']), first_hit(shuffled, t['demos'])
print(f"hidden program: {' -> '.join(prog)}")
print(f"  shortest-first order : {nb:>4} candidates evaluated")
print(f"  shuffled order       : {nr:>4}")
print(f"  the fixed order was worth log2({nr}/{nb}) = {math.log2(nr/nb):+.2f} bits, "
      f"before any model")

def parse_ranking(text):
    """Pull the model's ordering over primitives out of its reply."""
    obj = extract_json(text)
    if not isinstance(obj, dict):
        return None
    names = [str(v) for v in (obj.get("ranking") or []) if str(v) in OPS]
    return names or None

def order_by_ranking(names):
    """Best-first: cost is the mean rank of a program's primitives, so a chain of
    well-rated operations is tried early. Unlisted primitives sink to the bottom."""
    score = {n: -1.0 for n in OPS}
    for i, n in enumerate(names):
        score[n] = float(len(names) - i)
    def cost(i):
        prog = SPACE[i][0]
        return (-sum(score.get(n, 0.0) for n in prog) / len(prog), len(prog))
    return sorted(range(len(SPACE)), key=cost)

def demo_features(demos):
    """Cheap, model-free descriptors of what the transformation did."""
    ratios = [len(y) / max(len(x), 1) for x, y in demos]
    return {"ratio": statistics.fmean(ratios),
            "multiset_preserved": all(sorted(x) == sorted(y) for x, y in demos),
            "values_changed": all(any(v not in x for v in y) for x, y in demos),
            "sorted_asc": all(y == sorted(y) for _, y in demos),
            "sorted_desc": all(y == sorted(y, reverse=True) for _, y in demos),
            "singleton": all(len(y) == 1 for _, y in demos)}

def heuristic_scores(demos):
    """A hand-written primitive plausibility rule, thirty lines and no model.

    This is the baseline that decides what the model's ranking is worth. Beating
    brute force is easy; beating a morning's worth of hand-written rules is the
    claim actually being tested.
    """
    f = demo_features(demos)
    s = {n: 0.0 for n in OPS}
    s["duplicate"] += 3 if f["ratio"] > 1.5 else -3
    for n in ("every_second", "keep_last", "keep_max"):
        s[n] += 2 if f["ratio"] < 0.75 else -2
    if f["singleton"]:
        for n in ("keep_last", "keep_max"):
            s[n] += 3
    if f["multiset_preserved"]:
        for n in ("increment", "decrement", "double"):
            s[n] -= 4
        for n in ("reverse", "rotate_left", "rotate_right", "swap_ends",
                  "pairwise_swap", "sort_asc", "sort_desc"):
            s[n] += 1
    if f["values_changed"]:
        for n in ("increment", "decrement", "double"):
            s[n] += 3
    if f["sorted_asc"]:
        s["sort_asc"] += 2
    if f["sorted_desc"]:
        s["sort_desc"] += 2
    return s

def order_by_scores(scores):
    def cost(i):
        prog = SPACE[i][0]
        return (-sum(scores.get(n, 0.0) for n in prog) / len(prog), len(prog))
    return sorted(range(len(SPACE)), key=cost)

def guided_search(model, n_tasks=8):
    """Four solvers on the same tasks: brute, blind shuffle, hand heuristic, model."""
    fixed = list(range(len(SPACE)))
    got = {k: [] for k in ("brute", "random", "heuristic", "model")}
    bits = {k: [] for k in ("random", "heuristic", "model")}
    for i in range(n_tasks):
        r = random.Random(900 + i)
        prog = random_program(r, 3)
        t = make_task(900 + i, prog, n_demos=4, n_tests=1)
        shown = "\n".join(f"  {list(a)} -> {list(b)}" for a, b in t["demos"])
        reply = ask(model,
                    f"A hidden program is a chain of primitives applied left to right.\n"
                    f"Observed:\n{shown}\n\nAvailable primitives:\n{menu()}\n\n"
                    "Rank the primitives by how likely each is to appear. This ranking "
                    "will order an exhaustive search.\n\n"
                    'Answer with one JSON object: {"ranking": [<names, most likely first>]}',
                    max_tokens=1500)
        ranking = parse_ranking(reply.text)

        shuffled = list(range(len(SPACE)))
        random.Random(1000 + i).shuffle(shuffled)
        orders = {"brute": fixed, "random": shuffled,
                  "heuristic": order_by_scores(heuristic_scores(t["demos"])),
                  "model": order_by_ranking(ranking) if ranking else None}
        hits = {k: (first_hit(o, t["demos"]) if o is not None else None)
                for k, o in orders.items()}
        if not hits["brute"] or any(v is None for v in hits.values()):
            continue
        for k, v in hits.items():
            got[k].append(v)
        for k in bits:
            bits[k].append(math.log2(hits["brute"] / hits[k]))
    return got, bits

if not MODELS:
    print("no learner available — skipping")
else:
    print(f"{'model':<15}{'brute':>8}{'shuffled':>10}{'heuristic':>11}{'model':>8}"
          f"     guidance vs brute (bits)")
    for m in MODELS:
        got, bits = guided_search(m)
        if got["brute"]:
            med = lambda k: statistics.median(got[k])
            b = lambda k: statistics.fmean(bits[k])
            print(f"{m:<15}{med('brute'):>8.0f}{med('random'):>10.0f}"
                  f"{med('heuristic'):>11.0f}{med('model'):>8.0f}"
                  f"     shuffle {b('random'):+.2f}   hand {b('heuristic'):+.2f}"
                  f"   model {b('model'):+.2f}")

def _safe(fn, x):
    try:
        return list(fn(list(x)))
    except Exception:
        return None

def reachability(ablated, n_tasks=16):
    """Classify each task by exhaustive search *before* any model is asked.

    Using the removed primitive in the generating program does not make a rule
    inexpressible — another composition may compute exactly the same function.
    Grading a decline as correct without checking that would score the model on
    a label I guessed. So enumerate the reduced language and look.
    """
    ops = {n: f for n, f in OPS.items() if n != ablated}
    # Every raw chain, with no probe-signature pruning. The pruning is not
    # sound — measured near the top of this lab — and a classification that
    # calls a task impossible must not rest on it. Fourteen primitives to
    # depth 3 is under three thousand programs, so soundness is free here.
    space = [(c, as_callable(c, ops)) for d in (1, 2, 3)
             for c in itertools.product(list(ops), repeat=d)]
    out = []
    for i in range(n_tasks):
        r = random.Random(500 + i)
        # Random hidden programs over the FULL language, not programs forced to
        # contain the ablated primitive. Forcing it guarantees the answer and
        # turns the classification into a restatement of how the task was built.
        prog = random_program(r, r.choice([1, 2, 3]))
        t = make_task(500 + i, prog, n_demos=4, n_tests=1)
        fitting = [(pr, fn) for pr, fn in space
                   if all(_safe(fn, a) == list(b) for a, b in t["demos"])]
        exact = [pr for pr, fn in fitting if extensionally_equal(fn, t["fn"])]
        # A: the truth is expressible. B: something fits the demonstrations but
        # nothing computes the truth. C: nothing in the reduced language even
        # fits the demonstrations — the one class established exactly, since it
        # is a statement about a finite space and finitely many equations.
        cls = "A reachable" if exact else ("B fits but wrong" if fitting
                                           else "C no fitting program")
        out.append({"task": t, "ops": ops, "hidden": prog, "class": cls,
                    "uses_ablated": ablated in prog,
                    "reachable": len(exact) > 0, "n_fitting": len(fitting)})
    return out

def ablated_trial(model, ablated, n_tasks=16):
    """Offer the reduced language, allow a decline, and grade against the
    exhaustive verdict rather than against the shape of the hidden program."""
    cases = reachability(ablated, n_tasks)
    tally = Counter()
    for c in cases:
        t, ops = c["task"], c["ops"]
        prompt = identification_prompt(t["demos"], t["tests"][0][0], ops)
        prompt += ('\n\nIf no chain of the listed primitives can express the rule, answer '
                   '{"inexpressible": true} instead of guessing.')
        obj = extract_json(ask(model, prompt, max_tokens=1500).text) or {}
        tag = "reachable" if c["reachable"] else "unreachable"
        if obj.get("inexpressible"):
            tally[f"{tag}/declined"] += 1
            continue
        # graded against the language it was shown, not the full one
        fn, _ = compile_answer(obj, ops)
        fits = fn is not None and all(_safe(fn, a) == list(b) for a, b in t["demos"])
        if fits and fn is not None and extensionally_equal(fn, t["fn"]):
            tally[f"{tag}/correct"] += 1
        elif fits:
            tally[f"{tag}/fits but wrong"] += 1
        else:
            tally[f"{tag}/fails its own demos"] += 1
    return cases, tally

if not MODELS:
    print("no learner available — skipping")
else:
    for m in MODELS:
        for ablated in ("duplicate", "increment"):
            cases, tally = ablated_trial(m, ablated)
            by_class = Counter(c["class"] for c in cases)
            used = sum(1 for c in cases if c["uses_ablated"])
            still = sum(1 for c in cases if c["uses_ablated"] and c["reachable"])
            print(f"  {m} without `{ablated}`: " +
                  ", ".join(f"{by_class[k]} {k}" for k in sorted(by_class)) +
                  f"   ({used} tasks were generated using it, {still} survive anyway)")
            for k in sorted(tally):
                print(f"      {k:<34}{tally[k]}")
        print()

def decoys_per_task(n_demos, ablated=None, n_tasks=24, depth=3):
    """A decoy fits every demonstration and is not the hidden function.
    Purely a property of the language and the evidence — no model involved."""
    ops = OPS if ablated is None else {n: f for n, f in OPS.items() if n != ablated}
    space = [(pr, as_callable(pr, ops)) for pr in enumerate_programs(depth, ops)]
    counts = []
    for i in range(n_tasks):
        r = random.Random(7000 + i)
        prog = random_program(r, r.choice([1, 2, 3]))
        t = make_task(7000 + i, prog, n_demos=n_demos, n_tests=1)
        fits = [fn for _, fn in space
                if all(_safe(fn, a) == list(b) for a, b in t["demos"])]
        exact = [fn for fn in fits if extensionally_equal(fn, t["fn"])]
        counts.append(len(fits) - len(exact))
    return statistics.fmean(counts)

print(f"{'demos':>6}{'full language':>16}{'keep_last removed':>20}")
for nd in (1, 2, 3, 4, 6):
    print(f"{nd:>6}{decoys_per_task(nd):>16.1f}{decoys_per_task(nd, 'keep_last'):>20.1f}")

BASE = {n: f for n, f in OPS.items() if n != "identity"}
MOTIFS = [("reverse", "duplicate"), ("sort_asc", "every_second"),
          ("rotate_left", "increment"), ("swap_ends", "double")]

def expand(chain, lib):
    """A chain of library names, flattened back to base primitives."""
    out = []
    for name in chain:
        out.extend(lib.get(name, (name,)))
    return tuple(out)

def build_stream(n_tasks, seed=7):
    """Tasks that share recurring substructure without ever repeating.

    Each task is a motif wrapped in one arbitrary primitive, so the motifs
    genuinely recur — otherwise there is nothing for a library learner to find,
    and the experiment would be rigged in the opposite direction.
    """
    rng = random.Random(seed)
    tasks = []
    for i in range(n_tasks):
        motif = MOTIFS[i % len(MOTIFS)]
        wrap = rng.choice(list(BASE))
        chain = (wrap,) + motif if rng.random() < 0.5 else motif + (wrap,)
        fn = lambda x, c=chain: run(c, list(x), BASE)
        xs = [[rng.randint(1, 9) for _ in range(rng.choice([3, 4, 5]))] for _ in range(4)]
        tasks.append({"target": chain, "demos": [(x, fn(x)) for x in xs]})
    return tasks

def solve(demos, lib, max_depth=3):
    """Bottom-up search over base primitives plus the current library.

    A library entry costs depth 1 no matter how long it expands to — that
    discount is the entire benefit an abstraction is supposed to buy.
    """
    names = list(BASE) + list(lib)
    seen, frontier, nodes = set(), [()], 0
    for _ in range(max_depth):
        nxt = []
        for prefix in frontier:
            for name in names:
                cand = prefix + (name,)
                flat = expand(cand, lib)
                if len(flat) > 9:
                    continue
                fn = lambda x, f=flat: run(f, list(x), BASE)
                sig = signature(fn)
                if sig is None or sig in seen:
                    continue
                seen.add(sig)
                nodes += 1
                if all(_safe(fn, x) == list(y) for x, y in demos):
                    return nodes, cand
                nxt.append(cand)
        frontier = nxt
    return nodes, None

def mine(solutions, lib):
    """Best fragment by description-length saving: (uses - 1) * (len - 1).

    Not by frequency. A two-step fragment used twice saves one step and costs a
    dictionary entry; the objective has to say so. This is the DreamCoder
    library-learning criterion in miniature.
    """
    counts = Counter()
    for chain in solutions:
        flat = expand(chain, lib)
        for size in (2, 3):
            for i in range(len(flat) - size + 1):
                counts[flat[i:i + size]] += 1
    known = {tuple(v) for v in lib.values()}
    best, best_score = None, 0
    for frag, uses in counts.items():
        if frag in known:                 # already promoted; do not rediscover it
            continue
        score = (uses - 1) * (len(frag) - 1)
        if score > best_score:
            best, best_score = frag, score
    return best, best_score

def campaign(tasks, grow, rounds=5):
    lib, solved_chains, log = {}, [], []
    batch = max(1, len(tasks) // rounds)
    for r in range(rounds):
        chunk = tasks[r * batch:(r + 1) * batch]
        nodes, lengths, solved, exact = [], [], 0, 0
        for t in chunk:
            n, chain = solve(t["demos"], lib)
            nodes.append(n)
            if chain is not None:
                solved += 1
                lengths.append(len(chain))
                solved_chains.append(chain)
                # Search returns the *shortest* program that fits, which need not
                # be the one the task was built from. Mining only ever sees what
                # search returns, so this fraction bounds what it can discover.
                exact += (expand(chain, lib) == t["target"])
        log.append({"round": r, "solved": solved,
                    "median_nodes": statistics.median(nodes) if nodes else None,
                    "program_length": round(statistics.fmean(lengths), 2) if lengths else None,
                    "found_target": round(exact / max(solved, 1), 2),
                    "branching": len(BASE) + len(lib), "library": list(lib)})
        if grow and solved_chains:
            frag, score = mine(solved_chains, lib)
            if frag and score >= 2:
                lib[f"P{len(lib) + 1}"] = frag
    return log, lib

tasks = build_stream(60)
fixed_log, _ = campaign(tasks, grow=False)
grow_log, lib = campaign(tasks, grow=True)

print(f"{'round':>6}{'fixed nodes':>13}{'growing nodes':>15}{'speedup':>9}"
      f"{'prog len':>10}{'found target':>14}{'branching':>11}")
for a, b in zip(fixed_log, grow_log):
    sp = a["median_nodes"] / b["median_nodes"] if b["median_nodes"] else float("nan")
    print(f"{a['round']:>6}{a['median_nodes']:>13.0f}{b['median_nodes']:>15.0f}"
          f"{sp:>8.2f}x{b['program_length']:>10.2f}{b['found_target']:>14.2f}"
          f"{b['branching']:>11}")
def recovered(motif, lib):
    """Compare behaviour, not spelling. `double -> swap_ends` computes exactly the
    same function as `swap_ends -> double`, because one acts on values and the
    other on positions, so an equality test on tuples would report a motif as
    missed when the solver had in fact found it."""
    want = signature(lambda x: run(motif, list(x), BASE))
    return any(signature(lambda x, f=tuple(v): run(f, list(x), BASE)) == want
               for v in lib.values())

hits = sum(recovered(m, lib) for m in MOTIFS)
print(f"\nplanted motifs recovered: {hits} / {len(MOTIFS)}")
print("learned:", {k: " -> ".join(v) for k, v in lib.items()} or "nothing promoted")