"""Write the NextTerm-440M card's polynomial-continuation eval sets as jsonl.

Reproduces arithmetic_eval.py's sampling exactly (same tasks, coefficient ranges, per-task
RNG random.Random(seed + degree), same draw order): for each prior-term count k from the
task's minimum up to --max-terms, --samples-per-k sequences f(0), ..., f(k) of a random
non-negative-integer polynomial; the prompt is the first k terms and the answer is f(k).
Each row is {"id", "seq", "k"} with seq = prompt + [answer], so scripts/eval_neo.py
(last-term exact match; teacher-forced greedy) scores it directly.

Degrees 5 and 6 (quintic, sextic; --extended) are an optional extension, not part of the
published polynomial evals: the x^5 and
x^6 coefficients are drawn from [0, 9] like x^4's, before the card's draws (highest degree
first, as the card orders them), so degrees 1-4 stay byte-identical.

Usage: python gen_poly_evals.py <out_dir> [--extended] [--seed 0 --max-terms 25 --samples-per-k 200]
"""
import argparse
import json
import random
from pathlib import Path

COEFF_MAX = {"g": 9, "f": 9, "a": 9, "b": 99, "c": 999, "d": 9999, "e": 99999}  # g: x^6, f: x^5
CARD_TASKS = (("arithmetic", 1, 2), ("quadratic", 2, 3), ("cubic", 3, 4), ("quartic", 4, 5))
EXTENDED_TASKS = (("quintic", 5, 6), ("sextic", 6, 7))

p = argparse.ArgumentParser()
p.add_argument("out_dir")
p.add_argument("--seed", type=int, default=0)
p.add_argument("--max-terms", type=int, default=25)
p.add_argument("--samples-per-k", type=int, default=200)
p.add_argument("--extended", action="store_true", help="write quintic and sextic instead")
args = p.parse_args()
out = Path(args.out_dir)
out.mkdir(parents=True, exist_ok=True)

for name, degree, min_terms in (EXTENDED_TASKS if args.extended else CARD_TASKS):
    rng = random.Random(args.seed + degree)
    with open(out / f"poly_{name}.jsonl", "w") as f:
        for k in range(min_terms, args.max_terms + 1):
            for i in range(args.samples_per_k):
                g = rng.randint(0, COEFF_MAX["g"]) if degree >= 6 else 0
                f5 = rng.randint(0, COEFF_MAX["f"]) if degree >= 5 else 0
                # identical draw order to arithmetic_eval.sample_coefficients
                a = rng.randint(0, COEFF_MAX["a"]) if degree >= 4 else 0
                b = rng.randint(0, COEFF_MAX["b"]) if degree >= 3 else 0
                c = rng.randint(0, COEFF_MAX["c"]) if degree >= 2 else 0
                d = rng.randint(0, COEFF_MAX["d"]) if degree >= 1 else 0
                e = rng.randint(0, COEFF_MAX["e"])
                seq = [g * x**6 + f5 * x**5 + a * x**4 + b * x**3 + c * x**2 + d * x + e
                       for x in range(k + 1)]
                f.write(json.dumps({"id": f"{name}_k{k}_{i}", "seq": seq, "k": k}) + "\n")
    print(name, (args.max_terms - min_terms + 1) * args.samples_per_k, "rows")