"""Verifiable synthetic task generators with programmatic ground truth.

Families (each item is a dict: family, template, depth, prompt, answer, text, meta):
  arith   multi-step arithmetic chains            (depth = number of operations)
  bind    variable binding + dependent substitution (depth = chain length)
  fsm     register / finite-state updates          (depth = number of updates)
  excl    conjunctive constraints with distractors (depth = number of constraints)
  sched   rule-based selection under a budget and exclusions (depth = number of candidates)
  compose held-out composition: bind chain feeding an arithmetic expression (depth = total steps)

Splitting is done BEFORE sampling: by (family, template id, depth regime, seed stream). Train uses templates
T0-T2 and depths <= max_train_depth; held-out "template" items use templates T3-T4; held-out "depth" items use
depths > max_train_depth; the compose family is entirely held out from training unless explicitly requested.
Text format: "<prompt> Answer: <answer>" ; the answer is the only supervised span for exact-match evaluation.
"""
from __future__ import annotations

import hashlib
import json
import random
from dataclasses import dataclass, asdict
from typing import Callable

OPS = {"+": lambda a, b: a + b, "-": lambda a, b: a - b, "*": lambda a, b: a * b}

# Phrasing templates: index 0..2 = train, 3..4 = held-out paraphrases.
TEMPLATES = {
    "arith": ["Compute {expr}.", "What is {expr}?", "Evaluate the expression {expr}.", "Calculate {expr} step by step and give the result.", "Find the value of {expr}."],
    "bind": ["Let {chain}. What is {var}?", "Given {chain}, compute {var}.", "{chain}. Find {var}.", "Suppose {chain}. Then {var} equals what?", "With {chain}, the value of {var} is?"],
    "fsm": ["Start with register {init}. {ops}. What is the register?", "The counter begins at {init}. {ops}. Final counter?", "Initial value {init}. {ops}. Result?", "A register holds {init}. Apply: {ops}. What does it hold now?", "Begin at {init}; {ops}; final value?"],
    "excl": ["Items: {items}. Keep the items that are {rules}. Which items remain?", "Consider {items}. Select those that are {rules}. List them.", "{items}. Filter: {rules}. Remaining items?", "Given the list {items}, keep only the ones {rules}. Answer with the names.", "From {items}, which satisfy: {rules}?"],
    "sched": ["Budget {budget}. Tasks: {tasks}. Rule: {rule}. Choose the affordable task with the highest value.", "You have {budget} units. Options: {tasks}. Constraint: {rule}. Which affordable option has the largest value?", "Capacity {budget}. Candidates: {tasks}. Exclusion: {rule}. Pick the best affordable candidate.", "With a budget of {budget}, among {tasks}, under the rule {rule}, which task maximises value while staying affordable?", "Budget is {budget}. {tasks}. {rule}. Best affordable task by value?"],
    "compose": ["Let {chain}. Compute {expr}.", "Given {chain}, what is {expr}?", "{chain}. Evaluate {expr}.", "Suppose {chain}. Then find {expr}.", "With {chain}, calculate {expr}."],
}


def _rng(family: str, template: int, depth: int, seed: int, idx: int) -> random.Random:
    key = f"{family}|{template}|{depth}|{seed}|{idx}"
    return random.Random(int(hashlib.sha256(key.encode()).hexdigest()[:16], 16))


def gen_arith(r: random.Random, depth: int, template: int):
    vals = [r.randint(1, 20)]
    ops = []
    expr = str(vals[0])
    cur = vals[0]
    for _ in range(depth):
        op = r.choice(["+", "-", "*"])
        b = r.randint(1, 9) if op == "*" else r.randint(1, 30)
        cur = OPS[op](cur, b)  # left-to-right evaluation (stated below)
        expr += f" {op} {b}"
        ops.append(op)
    prompt = TEMPLATES["arith"][template].format(expr=expr) + " Evaluate left to right."
    return prompt, str(cur), {"ops": "".join(ops)}


def gen_bind(r: random.Random, depth: int, template: int):
    names = r.sample(list("abcdefghjkmnpqrstuvwxyz"), depth + 1)
    vals = {names[0]: r.randint(1, 15)}
    parts = [f"{names[0]} = {vals[names[0]]}"]
    for i in range(1, depth + 1):
        op = r.choice(["+", "-", "*"])
        b = r.randint(1, 5) if op == "*" else r.randint(1, 12)
        src = names[i - 1]
        vals[names[i]] = OPS[op](vals[src], b)
        parts.append(f"{names[i]} = {src} {op} {b}")
    # distractor: an unused variable
    d = r.choice([n for n in "ABCDEFG"])
    parts.insert(r.randint(1, len(parts)), f"{d} = {r.randint(1, 40)}")
    chain = ", ".join(parts)
    target = names[depth]
    prompt = TEMPLATES["bind"][template].format(chain=chain, var=target)
    return prompt, str(vals[target]), {"target": target}


def gen_fsm(r: random.Random, depth: int, template: int):
    init = r.randint(0, 10)
    cur = init
    steps = []
    for _ in range(depth):
        op = r.choice(["add", "sub", "double", "reset", "halve"])
        if op == "add":
            b = r.randint(1, 9); cur += b; steps.append(f"add {b}")
        elif op == "sub":
            b = r.randint(1, 9); cur -= b; steps.append(f"subtract {b}")
        elif op == "double":
            cur *= 2; steps.append("double it")
        elif op == "halve":
            cur = cur // 2; steps.append("halve it (round down)")
        else:
            b = r.randint(0, 9); cur = b; steps.append(f"reset to {b}")
    prompt = TEMPLATES["fsm"][template].format(init=init, ops=", then ".join(steps))
    return prompt, str(cur), {}


def gen_excl(r: random.Random, depth: int, template: int):
    n = r.randint(4, 7)
    names = r.sample(["apple", "pear", "plum", "fig", "kiwi", "lime", "date", "yam", "nut", "bean"], n)
    items = {nm: (r.randint(1, 30), r.choice(["red", "green", "blue"])) for nm in names}
    rules = []
    def keep(v, c):
        return True
    preds: list[Callable] = []
    for _ in range(depth):
        kind = r.choice(["gt", "lt", "color", "even"])
        if kind == "gt":
            t = r.randint(5, 25); rules.append(f"greater than {t}"); preds.append(lambda v, c, t=t: v > t)
        elif kind == "lt":
            t = r.randint(5, 25); rules.append(f"smaller than {t}"); preds.append(lambda v, c, t=t: v < t)
        elif kind == "color":
            col = r.choice(["red", "green", "blue"]); rules.append(f"not {col}"); preds.append(lambda v, c, col=col: c != col)
        else:
            rules.append("even"); preds.append(lambda v, c: v % 2 == 0)
    remaining = [nm for nm in names if all(p(*items[nm]) for p in preds)]
    items_s = ", ".join(f"{nm} ({v}, {c})" for nm, (v, c) in items.items())
    prompt = TEMPLATES["excl"][template].format(items=items_s, rules=" and ".join(rules))
    ans = " ".join(remaining) if remaining else "none"
    return prompt, ans, {"n_items": n, "n_remaining": len(remaining)}


def gen_sched(r: random.Random, depth: int, template: int):
    names = [chr(ord("A") + i) for i in range(depth)]
    tasks = {nm: (r.randint(1, 12), r.randint(1, 20)) for nm in names}
    budget = r.randint(4, 12)
    kind = r.choice(["cost_gt", "value_lt", "none"])
    if kind == "cost_gt":
        t = r.randint(3, 10); rule = f"exclude tasks with cost greater than {t}"; ok = lambda c, v: c <= t
    elif kind == "value_lt":
        t = r.randint(3, 10); rule = f"exclude tasks with value below {t}"; ok = lambda c, v: v >= t
    else:
        rule = "no exclusion"; ok = lambda c, v: True
    feas = [(v, nm) for nm, (c, v) in tasks.items() if c <= budget and ok(c, v)]
    if not feas:
        ans = "none"
    else:
        best = max(v for v, _ in feas)
        ans = sorted(nm for v, nm in feas if v == best)[0]
    tasks_s = ", ".join(f"{nm} (cost {c}, value {v})" for nm, (c, v) in tasks.items())
    prompt = TEMPLATES["sched"][template].format(budget=budget, tasks=tasks_s, rule=rule)
    return prompt, ans, {"n_feasible": len(feas)}


def gen_compose(r: random.Random, depth: int, template: int):
    # bind chain of length depth-2 feeding a 2-op arithmetic expression over two chain variables
    k = max(1, depth - 2)
    names = r.sample(list("abcdefghjkmnpqrstuvwxyz"), k + 1)
    vals = {names[0]: r.randint(1, 10)}
    parts = [f"{names[0]} = {vals[names[0]]}"]
    for i in range(1, k + 1):
        op = r.choice(["+", "-"]); b = r.randint(1, 9); src = names[i - 1]
        vals[names[i]] = OPS[op](vals[src], b); parts.append(f"{names[i]} = {src} {op} {b}")
    v1, v2 = r.sample(names, 2) if len(names) > 1 else (names[0], names[0])
    op1, op2 = r.choice(["+", "-", "*"]), r.choice(["+", "-"])
    c = r.randint(1, 9)
    val = OPS[op2](OPS[op1](vals[v1], vals[v2]), c)
    expr = f"{v1} {op1} {v2} {op2} {c} (left to right)"
    prompt = TEMPLATES["compose"][template].format(chain=", ".join(parts), expr=expr)
    return prompt, str(val), {}


GENERATORS = {"arith": gen_arith, "bind": gen_bind, "fsm": gen_fsm, "excl": gen_excl, "sched": gen_sched, "compose": gen_compose}


def make_item(family: str, template: int, depth: int, seed: int, idx: int) -> dict:
    r = _rng(family, template, depth, seed, idx)
    prompt, answer, meta = GENERATORS[family](r, depth, template)
    text = f"{prompt} Answer: {answer}"
    return {"family": family, "template": template, "depth": depth, "seed": seed, "idx": idx,
            "prompt": prompt + " Answer:", "answer": " " + answer, "text": text, "meta": meta,
            "uid": hashlib.sha256(text.encode()).hexdigest()[:16]}


@dataclass
class SplitSpec:
    name: str
    families: tuple[str, ...]
    templates: tuple[int, ...]
    depths: tuple[int, ...]
    seed: int
    n_per_cell: int


def generate_split(spec: SplitSpec) -> list[dict]:
    items = []
    for fam in spec.families:
        for tpl in spec.templates:
            for dep in spec.depths:
                for i in range(spec.n_per_cell):
                    items.append(make_item(fam, tpl, dep, spec.seed, i))
    return items


def default_splits(n_train_per_cell: int = 400, n_eval_per_cell: int = 60) -> dict[str, SplitSpec]:
    """Train: 5 families, templates 0-2, depths 1-4, seed 1. Dev: seed 2 same cells (IID generalisation).
    Held-out template: templates 3-4. Held-out depth: depths 5-6. Compose: fully held-out family.
    Final test: seed 9 across all eval regimes (LOCKED until preregistration is frozen)."""
    fams = ("arith", "bind", "fsm", "excl", "sched")
    return {
        "train": SplitSpec("train", fams, (0, 1, 2), (1, 2, 3, 4), 1, n_train_per_cell),
        "calib": SplitSpec("calib", fams, (0, 1, 2), (1, 2, 3, 4), 3, n_eval_per_cell // 2),
        "dev_iid": SplitSpec("dev_iid", fams, (0, 1, 2), (1, 2, 3, 4), 2, n_eval_per_cell),
        "dev_template": SplitSpec("dev_template", fams, (3, 4), (1, 2, 3, 4), 2, n_eval_per_cell),
        "dev_depth": SplitSpec("dev_depth", fams, (0, 1, 2), (5, 6), 2, n_eval_per_cell),
        "dev_compose": SplitSpec("dev_compose", ("compose",), (0, 1, 2), (3, 4, 5), 2, n_eval_per_cell),
        "test_iid": SplitSpec("test_iid", fams, (0, 1, 2), (1, 2, 3, 4), 9, n_eval_per_cell),
        "test_template": SplitSpec("test_template", fams, (3, 4), (1, 2, 3, 4), 9, n_eval_per_cell),
        "test_depth": SplitSpec("test_depth", fams, (0, 1, 2), (5, 6), 9, n_eval_per_cell),
        "test_compose": SplitSpec("test_compose", ("compose",), (0, 1, 2), (3, 4, 5), 9, n_eval_per_cell),
        # controller-label stream (D11), written LAST so that every earlier split is byte-identical to v1
        "ctrl_train": SplitSpec("ctrl_train", fams, (0, 1, 2), (1, 2, 3, 4), 4, n_train_per_cell // 2),
    }


def write_splits(out_dir: str, n_train_per_cell: int = 400, n_eval_per_cell: int = 60) -> dict:
    from pathlib import Path
    out = Path(out_dir); out.mkdir(parents=True, exist_ok=True)
    manifest = {}
    all_uids: dict[str, str] = {}
    for name, spec in default_splits(n_train_per_cell, n_eval_per_cell).items():
        items = generate_split(spec)
        # cross-split exact-duplicate removal: an item whose full text already appeared in an earlier split
        # (train is written first) is DROPPED from the later split; within-split duplicates are dropped too.
        kept, dups = [], 0
        for it in items:
            if it["uid"] in all_uids:
                dups += 1
                continue
            all_uids[it["uid"]] = name
            kept.append(it)
        items = kept
        p = out / f"{name}.jsonl"
        with open(p, "w") as f:
            for it in items:
                f.write(json.dumps(it) + "\n")
        manifest[name] = {"spec": asdict(spec), "n": len(items), "duplicates_dropped": dups,
                          "sha256": hashlib.sha256(p.read_bytes()).hexdigest(), "locked": name.startswith("test")}
    (out / "manifest.json").write_text(json.dumps(manifest, indent=2))
    return manifest