Blame for research/scripts/curriculum.py

"""Procedural curriculum: generators of verifiable reasoning problems with exact answers.

The families are deliberately NOT the families of the eval suites (no BBH task types, no MMLU questions): training on
them and then measuring BBH / MMLU-Pro / math measures transfer of the skill (state tracking, systematic search,
exact computation, careful reading of code), not memorised task formats.

Every item: {id, task: "math", src: "cur/<family>", prompt, answer} with the answer graded by grade.py's math checker.
usage: curriculum.py <n_per_family> <out.jsonl> [seed]
"""
import itertools
import json
import math
import random
import sys
from fractions import Fraction

TAIL = "\n\nPlease reason step by step, and put your final answer within \\boxed{}."
NAMES = ["A", "B", "C", "D", "E", "F", "G", "H", "I"]


def code_trace(r):
    """Short Python program; ask for the printed value (single line)."""
    a, b, n = r.randint(2, 9), r.randint(1, 7), r.randint(4, 9)
    kind = r.randint(0, 4)
    if kind == 0:
        src = f"x = {a}\ny = {b}\nfor i in range({n}):\n    if i % 2 == 0:\n        x = x + y * i\n    else:\n        y = (y + x) % 11\nprint(x - y)"
    elif kind == 1:
        lst = [r.randint(-9, 20) for _ in range(n)]
        src = f"nums = {lst}\nbest, cur = nums[0], 0\nfor v in nums:\n    cur = max(v, cur + v)\n    best = max(best, cur)\nprint(best)"
    elif kind == 2:
        w = "".join(r.choice("abcde") for _ in range(n + 3))
        src = f"s = '{w}'\nseen = {{}}\nout = 0\nfor i, ch in enumerate(s):\n    if ch in seen:\n        out += i - seen[ch]\n    seen[ch] = i\nprint(out)"
    elif kind == 3:
        src = f"def f(k):\n    if k <= 1:\n        return {b}\n    return f(k - 1) + {a} * f(k - 2) % 7\nprint(f({n}))"
    else:
        lst = [r.randint(1, 30) for _ in range(n)]
        src = f"a = {lst}\nfor i in range(len(a)):\n    for j in range(len(a) - 1 - i):\n        if a[j] % {a} > a[j + 1] % {a}:\n            a[j], a[j + 1] = a[j + 1], a[j]\nprint(a[{r.randint(0, n - 1)}])"
    env = {}
    import io, contextlib
    buf = io.StringIO()
    with contextlib.redirect_stdout(buf):
        exec(src, env)  # noqa: S102  (our own generated code)
    return f"What does the following Python program print?\n\n```python\n{src}\n```", buf.getvalue().strip()


def linear_system(r):
    while True:
        x, y, z = (r.randint(-9, 9) for _ in range(3))
        rows = [[r.randint(-5, 5) for _ in range(3)] for _ in range(3)]
        det = (rows[0][0] * (rows[1][1] * rows[2][2] - rows[1][2] * rows[2][1])
               - rows[0][1] * (rows[1][0] * rows[2][2] - rows[1][2] * rows[2][0])
               + rows[0][2] * (rows[1][0] * rows[2][1] - rows[1][1] * rows[2][0]))
        if det != 0:
            break
    eqs = []
    for c in rows:
        rhs = c[0] * x + c[1] * y + c[2] * z
        terms = " ".join(f"{'+' if v >= 0 else '-'} {abs(v)}{s}" for v, s in zip(c, "xyz") if v != 0).lstrip("+ ")
        eqs.append(f"{terms} = {rhs}")
    p, q = r.randint(1, 5), r.randint(1, 5)
    return ("Solve the system of equations\n" + "\n".join(eqs) + f"\nand find the value of {p}x + {q}y - z.",
            str(p * x + q * y - z))


def number_theory(r):
    k = r.randint(0, 2)
    if k == 0:
        a, e, m = r.randint(2, 50), r.randint(50, 2000), r.choice([7, 9, 11, 13, 17, 19, 23, 100, 1000])
        return f"Find the remainder when ${a}^{{{e}}}$ is divided by ${m}$.", str(pow(a, e, m))
    if k == 1:
        nums = [r.randint(2, 60) * r.choice([1, 2, 3, 6]) for _ in range(3)]
        g = math.gcd(*nums)
        l = math.lcm(*nums)
        return f"Let $g = \\gcd({nums[0]}, {nums[1]}, {nums[2]})$ and $L = \\operatorname{{lcm}}({nums[0]}, {nums[1]}, {nums[2]})$. Compute $L / g$.", str(l // g)
    n = r.randint(200, 5000)
    if n < 300:
        return f"How many positive divisors does ${n}^2$ have?", str(sum(1 for d in range(1, n * n + 1) if (n * n) % d == 0))
    return f"How many positive divisors does ${n}$ have?", str(sum(1 for d in range(1, n + 1) if n % d == 0))


def graph_path(r):
    n = r.randint(6, 9)
    nodes = NAMES[:n]
    edges = {}
    for i in range(1, n):  # spanning tree keeps it connected
        j = r.randint(0, i - 1)
        edges[(nodes[j], nodes[i])] = r.randint(1, 12)
    for _ in range(r.randint(n, 2 * n)):
        a, b = r.sample(nodes, 2)
        if (a, b) not in edges and (b, a) not in edges:
            edges[(a, b)] = r.randint(1, 12)
    dist = {v: math.inf for v in nodes}
    dist[nodes[0]] = 0
    todo = set(nodes)
    while todo:
        u = min(todo, key=lambda v: dist[v])
        todo.remove(u)
        for (a, b), w in edges.items():
            for s, t in ((a, b), (b, a)):
                if s == u and dist[u] + w < dist[t]:
                    dist[t] = dist[u] + w
    target = max(nodes[1:], key=lambda v: (dist[v], r.random()))
    lines = [f"{a} - {b}: {w}" for (a, b), w in r.sample(list(edges.items()), len(edges))]
    return (f"Roads connect towns {', '.join(nodes)} (both directions, with lengths):\n" + "\n".join(lines) +
            f"\nWhat is the length of the shortest route from {nodes[0]} to {target}?", str(dist[target]))


def scheduling(r):
    n = r.randint(5, 8)
    tasks = NAMES[:n]
    dur = {t: r.randint(1, 9) for t in tasks}
    pre = {t: sorted(r.sample(tasks[:i], r.randint(0, min(2, i)))) if i else [] for i, t in enumerate(tasks)}
    fin = {}
    for t in tasks:
        fin[t] = max((fin[p] for p in pre[t]), default=0) + dur[t]
    desc = "\n".join(f"Task {t} takes {dur[t]} hours" + (f" and can start only after {' and '.join(pre[t])} "
                                                        f"{'is' if len(pre[t]) == 1 else 'are'} finished." if pre[t] else ".")
                     for t in r.sample(tasks, n))
    return ("A project consists of tasks; any number of tasks can run at the same time.\n" + desc +
            "\nWhat is the minimum number of hours needed to finish all tasks?", str(max(fin.values())))


def counting(r):
    N = r.randint(100, 3000)
    a, b = r.sample([2, 3, 4, 5, 6, 7, 9, 11], 2)
    c = r.choice([x for x in [3, 4, 5, 7, 8, 10, 13] if x not in (a, b)])
    cnt = sum(1 for k in range(1, N + 1) if (k % a == 0 or k % b == 0) and k % c != 0)
    return f"How many integers from 1 to {N} (inclusive) are divisible by {a} or by {b}, but not by {c}?", str(cnt)


def probability(r):
    k = r.randint(0, 1)
    if k == 0:
        s = r.randint(5, 10)
        faces = r.choice([4, 6, 8])
        good = sum(1 for x, y, z in itertools.product(range(1, faces + 1), repeat=3) if x + y + z == s)
        p = Fraction(good, faces ** 3)
        return f"Three fair {faces}-sided dice (faces 1 to {faces}) are rolled. What is the probability that the sum is {s}? Give the answer as a fraction in lowest terms.", f"\\frac{{{p.numerator}}}{{{p.denominator}}}"
    red, blue = r.randint(2, 7), r.randint(2, 7)
    d = r.randint(2, 3)
    tot = math.comb(red + blue, d)
    good = sum(math.comb(red, i) * math.comb(blue, d - i) for i in range(1, d + 1) if i <= red and d - i <= blue)
    p = Fraction(good, tot)
    return f"A bag has {red} red and {blue} blue marbles. {d} marbles are drawn at random without replacement. What is the probability that at least one of them is red? Give the answer as a fraction in lowest terms.", f"\\frac{{{p.numerator}}}{{{p.denominator}}}"


def cipher(r):
    words = ["planet", "garden", "silver", "rocket", "window", "puzzle", "forest", "bridge", "candle", "violin",
             "harbor", "meadow", "tunnel", "lantern", "compass"]
    w = r.choice(words)
    s = r.randint(3, 22)
    enc = "".join(chr((ord(ch) - 97 + s) % 26 + 97) for ch in w)
    w2 = r.choice([x for x in words if x != w])
    enc2 = "".join(chr((ord(ch) - 97 + s) % 26 + 97) for ch in w2)
    return (f"A Caesar cipher shifts every letter forward by the same amount (wrapping z to a). The word '{w}' is "
            f"encrypted as '{enc}'. Using the same shift, decrypt '{enc2}'. Give the decrypted word.",
            f"\\text{{{w2}}}")


def determinant(r):
    m = [[r.randint(-6, 6) for _ in range(3)] for _ in range(3)]
    d = (m[0][0] * (m[1][1] * m[2][2] - m[1][2] * m[2][1]) - m[0][1] * (m[1][0] * m[2][2] - m[1][2] * m[2][0])
         + m[0][2] * (m[1][0] * m[2][1] - m[1][1] * m[2][0]))
    rows = " \\\\ ".join(" & ".join(str(v) for v in row) for row in m)
    return f"Compute the determinant of the matrix $\\begin{{pmatrix}} {rows} \\end{{pmatrix}}$.", str(d)


def rates(r):
    a, b = r.randint(2, 12), r.randint(2, 12)
    c = r.randint(a * b // (a + b) + 2, 40) if r.random() < 0.5 else None
    if c:
        t = Fraction(1, Fraction(1, a) + Fraction(1, b) - Fraction(1, c))
        q = f"Pipe 1 fills a tank in {a} hours and pipe 2 fills it in {b} hours, while a drain empties the full tank in {c} hours. Starting empty with both pipes and the drain open, how many hours until the tank is full?"
    else:
        t = Fraction(1, Fraction(1, a) + Fraction(1, b))
        q = f"Alice paints a fence in {a} hours and Bob paints the same fence in {b} hours. Working together, how many hours do they need?"
    return q + " Give an exact answer.", (str(t.numerator) if t.denominator == 1 else f"\\frac{{{t.numerator}}}{{{t.denominator}}}")


def recurrence(r):
    a, b, c = r.randint(1, 4), r.randint(-3, 3), r.randint(-5, 5)
    x0, x1, n = r.randint(0, 5), r.randint(0, 5), r.randint(8, 14)
    xs = [x0, x1]
    for _ in range(n - 1):
        xs.append(a * xs[-1] + b * xs[-2] + c)
    return (f"A sequence is defined by $x_0 = {x0}$, $x_1 = {x1}$ and $x_{{n+1}} = {a}x_n {'+' if b >= 0 else '-'} "
            f"{abs(b)}x_{{n-1}} {'+' if c >= 0 else '-'} {abs(c)}$ for $n \\ge 1$. Find $x_{{{n}}}$.", str(xs[n]))


def string_ops(r):
    """Hard: state tracking over a sequence of string operations."""
    s = "".join(r.choice("abcdefghij") for _ in range(r.randint(7, 10)))
    ops, cur = [], list(s)
    for _ in range(r.randint(7, 11)):
        k = r.randint(0, 4)
        if k == 0:
            i, j = sorted(r.sample(range(len(cur)), 2))
            ops.append(f"swap the letters at positions {i + 1} and {j + 1}")
            cur[i], cur[j] = cur[j], cur[i]
        elif k == 1:
            t = r.randint(1, len(cur) - 1)
            ops.append(f"rotate the string left by {t}")
            cur = cur[t:] + cur[:t]
        elif k == 2:
            i, j = sorted(r.sample(range(len(cur)), 2))
            ops.append(f"reverse the part from position {i + 1} to position {j + 1}")
            cur[i:j + 1] = cur[i:j + 1][::-1]
        elif k == 3:
            a, b = r.sample("abcdefghij", 2)
            ops.append(f"replace every '{a}' with '{b}'")
            cur = [b if c == a else c for c in cur]
        else:
            i = r.randint(0, len(cur) - 1)
            ops.append(f"delete the letter at position {i + 1}")
            del cur[i]
    steps = "\n".join(f"{n + 1}. {o}" for n, o in enumerate(ops))
    return (f"Start with the string '{s}' (positions are counted from 1). Apply these operations in order:\n{steps}\n"
            "What is the final string?", f"\\text{{{''.join(cur)}}}")


def stack_sim(r):
    st, ops, total = [], [], 0
    for _ in range(r.randint(12, 18)):
        if st and r.random() < 0.4:
            k = r.choice(["pop", "dup", "add"] if len(st) >= 2 else ["pop", "dup"])
            if k == "pop":
                total += st.pop()
                ops.append("pop (add the popped value to the running total)")
            elif k == "dup":
                st.append(st[-1])
                ops.append("duplicate the top")
            else:
                a, b = st.pop(), st.pop()
                st.append(a + b)
                ops.append("pop two values and push their sum")
        else:
            v = r.randint(1, 20)
            st.append(v)
            ops.append(f"push {v}")
    steps = "\n".join(f"{n + 1}. {o}" for n, o in enumerate(ops))
    return (f"A stack starts empty and a running total starts at 0. Perform:\n{steps}\nAt the end, what is the running "
            f"total plus the sum of all values left on the stack times 2?", str(total + 2 * sum(st)))


def crt(r):
    ms = r.sample([3, 4, 5, 7, 9, 11, 13], 3)
    while math.gcd(ms[0], ms[1]) != 1 or math.gcd(ms[0], ms[2]) != 1 or math.gcd(ms[1], ms[2]) != 1:
        ms = r.sample([3, 4, 5, 7, 9, 11, 13], 3)
    x = r.randint(0, ms[0] * ms[1] * ms[2] - 1)
    cond = ", ".join(f"$x \\equiv {x % m} \\pmod{{{m}}}$" for m in ms)
    return f"Find the smallest non-negative integer $x$ such that {cond}.", str(x)


def graph_path_big(r):
    q, a = graph_path(_Big(r))
    return q, a


class _Big:
    """Random proxy that makes graph_path draw 10-12 nodes."""

    def __init__(self, r):
        self.r = r
        self.first = True

    def randint(self, a, b):
        if self.first and (a, b) == (6, 9):
            self.first = False
            return self.r.randint(9, 9)
        return self.r.randint(a, b)

    def __getattr__(self, k):
        return getattr(self.r, k)


def counting_big(r):
    N = r.randint(10 ** 5, 10 ** 6)
    a, b, c = r.sample([2, 3, 5, 7, 11, 13], 3)
    cnt = N // a + N // b + N // c - N // (a * b) - N // (a * c) - N // (b * c) + N // (a * b * c)
    return f"How many integers from 1 to {N} are divisible by at least one of {a}, {b}, {c}?", str(cnt)


def linear4(r):
    while True:
        sol = [r.randint(-6, 6) for _ in range(4)]
        M = [[r.randint(-4, 4) for _ in range(4)] for _ in range(4)]
        A = [row[:] for row in M]
        det = Fraction(1)
        for i in range(4):  # Gaussian elimination determinant
            p = next((k for k in range(i, 4) if A[k][i] != 0), None)
            if p is None:
                det = 0
                break
            if p != i:
                A[i], A[p] = A[p], A[i]
                det = -det
            det *= A[i][i]
            for k in range(i + 1, 4):
                f = Fraction(A[k][i], A[i][i])
                A[k] = [A[k][j] - f * A[i][j] for j in range(4)]
        if det != 0:
            break
    eqs = []
    for row in M:
        rhs = sum(c * v for c, v in zip(row, sol))
        terms = " ".join(f"{'+' if c >= 0 else '-'} {abs(c)}{s}" for c, s in zip(row, "wxyz") if c).lstrip("+ ")
        eqs.append(f"{terms} = {rhs}")
    return "Solve the system\n" + "\n".join(eqs) + "\nand give the value of w + 2x + 3y + 4z.", \
        str(sol[0] + 2 * sol[1] + 3 * sol[2] + 4 * sol[3])


FAMILIES = [code_trace, linear_system, number_theory, graph_path, scheduling, counting, probability, cipher,
            determinant, rates, recurrence]
HARD = [string_ops, stack_sim, crt, graph_path_big, counting_big, linear4]

if __name__ == "__main__":
    n, out = int(sys.argv[1]), sys.argv[2]
    rng = random.Random(int(sys.argv[3]) if len(sys.argv) > 3 else 0)
    fams = HARD if len(sys.argv) > 4 and sys.argv[4] == "hard" else FAMILIES
    rows = []
    for fam in fams:
        seen = set()
        while sum(1 for x in rows if x["src"] == "cur/" + fam.__name__) < n:
            q, a = fam(rng)
            if q in seen:
                continue
            seen.add(q)
            rows.append({"id": f"cur/{fam.__name__}/{len(seen)}", "task": "math", "src": "cur/" + fam.__name__,
                         "prompt": q + TAIL, "answer": a})
    rng.shuffle(rows)
    with open(out, "w", encoding="utf-8") as f:
        for x in rows:
            f.write(json.dumps(x, ensure_ascii=False) + "\n")
    print(len(rows), "items")