"""Build the seed prompt set the teacher will answer, and a verifiable maths set beside it.

Purpose: prompt seeding, done as a script you can read rather than as a folder of
         text somebody pasted. Prompts are generated from a taxonomy crossed with
         topics and constraints, so diversity is a property of the construction
         instead of a hope. A second file holds arithmetic word problems whose
         answers were computed before the questions were written, which is what
         lets rejection-sample.py verify a teacher trace without a second model.
Platform: all (standard library only; no model, no accelerator, no network)
Minimum memory: 8 GB, and far less in practice
Assumes: Python 3.10 or newer. Nothing is downloaded and nothing is served; this
         runs before the teacher is started.

Usage: python3 make-seed-prompts.py --out-dir seeds --seed 0
       python3 make-seed-prompts.py --out-dir seeds --train 600 --heldout 120 --seed 0
       python3 make-seed-prompts.py --out-dir seeds --maths 300 --seed 0

Three files are written:
  seeds/prompts.jsonl   the prompts the teacher answers, one JSON object per line
  seeds/heldout.jsonl   prompts of the same shape, never generated on, for evaluation
  seeds/maths.jsonl     word problems with a computed answer, for rejection sampling

Every line carries an id, a category and the slots it was built from, so a prompt
that turns out to produce bad teacher output can be traced back to the template
that made it rather than deleted one example at a time.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import random
from pathlib import Path

# ---------------------------------------------------------------------------
# The taxonomy. Categories match Part 10's evaluation set on purpose: a student
# is distilled for the tasks it will be measured on, and a category that appears
# in neither the training prompts nor the evaluation set is a category nobody
# learns anything about.
# ---------------------------------------------------------------------------

TOPICS = [
    "a llama.cpp server that stopped answering",
    "a GGUF file that will not load",
    "a KV cache that grew past the memory budget",
    "a fine-tune whose loss went flat",
    "a model download that failed halfway",
    "an embedding index that returns the wrong chunk",
    "a container that cannot see the accelerator",
    "a batch job that runs slower with more workers",
    "a chat template that was applied twice",
    "a quantised model that answers differently from the original",
    "a gateway alias pointing at the wrong port",
    "a tokeniser that splits an identifier into six pieces",
    "a training run that ran out of disk at the last checkpoint",
    "a dataset split that leaked into the evaluation set",
    "a rollout that timed out under load",
    "a draft model with a low acceptance rate",
]

CONSTRAINTS = [
    "Answer in at most three sentences.",
    "Reply with a numbered list of at most four steps.",
    "Answer in one sentence, then give one command on its own line.",
    "Use British spelling and avoid the word 'simply'.",
    "Give the answer first and the reason second, on two lines.",
    "Reply with a single paragraph under sixty words.",
]

TEMPLATES: dict[str, list[str]] = {
    "diagnose": [
        "I am looking at {topic}. What are the two most likely causes, and what would I check first? {constraint}",
        "Something is wrong: {topic}. Give me a diagnostic order, cheapest check first. {constraint}",
    ],
    "explain": [
        "Explain to a colleague who runs models but has never trained one what is going on when there is {topic}. {constraint}",
        "In plain terms, what does {topic} tell you about how the system is put together? {constraint}",
    ],
    "format": [
        "Summarise the situation '{topic}' as a JSON object with the keys symptom, likely_cause and next_check, and output nothing else.",
        "Turn '{topic}' into a table with the headers Symptom, Evidence and Action, and output nothing else.",
    ],
    "code": [
        "Write a short shell function that checks for the condition behind {topic} and exits non-zero when it holds. Code only, no explanation.",
        "Write a Python function that reads a log file and returns True when it shows {topic}. Code only, no explanation.",
    ],
    "refusal": [
        "Given only the phrase '{topic}', tell me the exact hostname and port of the machine involved.",
        "Given only the phrase '{topic}', tell me how many times this happened last week on my machine.",
    ],
    "instruction": [
        "Write three sentences about {topic}. Each sentence must be under twelve words, the second must contain a number, and the third must be a question.",
        "Describe {topic} without using the words 'model', 'server' or 'error'. {constraint}",
    ],
}

# The two refusal templates are the point of the category: the honest answer is
# that the information is not available, and a teacher that invents a hostname is
# a teacher whose output has to be filtered before a student copies the habit.

NAMES = ["Ada", "Bela", "Chi", "Dara", "Eero", "Fen", "Gita", "Hugo", "Ines", "Jonas",
         "Kira", "Liam", "Mira", "Nils", "Oona", "Pia", "Rafa", "Sena", "Tomas", "Ulla"]
GOODS = [("cable", "cables"), ("filter", "filters"), ("drive", "drives"), ("fan", "fans"),
         ("tile", "tiles"), ("bulb", "bulbs"), ("mount", "mounts"), ("clip", "clips")]

MATHS_INSTRUCTION = (
    "Solve the problem. Put your working inside <think> and </think> tags, then end "
    "with a single line of the form 'Answer: N' where N is the final number and "
    "nothing follows it."
)


def maths_problem(rng: random.Random) -> tuple[str, int]:
    """One word problem and the answer, computed here before the question exists.

    Three shapes, all multi-step, all integer. The answer is arithmetic performed
    by this function, so the label cannot be wrong in the way a scraped label can.
    """
    name = rng.choice(NAMES)
    singular, plural = rng.choice(GOODS)
    shape = rng.randrange(3)

    if shape == 0:
        boxes = rng.randint(3, 12)
        per_box = rng.randint(4, 15)
        broken = rng.randint(1, min(9, boxes * per_box - 1))
        answer = boxes * per_box - broken
        text = (f"{name} unpacks {boxes} boxes of {plural}, each holding {per_box} {plural}. "
                f"{broken} of the {plural} are broken and thrown away. "
                f"How many usable {plural} does {name} have?")
        return text, answer

    if shape == 1:
        start = rng.randint(20, 90)
        bought = rng.randint(5, 40)
        given = rng.randint(1, 15)
        days = rng.randint(2, 6)
        per_day = rng.randint(1, 6)
        answer = start + bought - given - days * per_day
        text = (f"{name} starts the week with {start} {plural}, buys {bought} more, and gives "
                f"{given} to a neighbour. Over the next {days} days {name} uses {per_day} "
                f"{plural} each day. How many {plural} are left?")
        return text, answer

    racks = rng.randint(2, 8)
    shelves = rng.randint(2, 6)
    per_shelf = rng.randint(2, 9)
    spare = rng.randint(0, 20)
    answer = racks * shelves * per_shelf + spare
    text = (f"A store room has {racks} racks. Each rack has {shelves} shelves and each shelf "
            f"holds {per_shelf} {plural}. A drawer holds another {spare} {plural}. "
            f"How many {plural} are in the room altogether?")
    return text, answer


def build_prompts(count: int, rng: random.Random, prefix: str) -> list[dict]:
    """Cross the taxonomy with topics and constraints, without repeating a pair."""
    combinations = []
    for category, templates in TEMPLATES.items():
        for template in templates:
            for topic in TOPICS:
                for constraint in CONSTRAINTS:
                    combinations.append((category, template, topic, constraint))
    rng.shuffle(combinations)

    rows = []
    seen: set[str] = set()
    for category, template, topic, constraint in combinations:
        if len(rows) >= count:
            break
        prompt = template.format(topic=topic, constraint=constraint).strip()
        # A template with no {constraint} slot collapses several combinations onto
        # the same text; keep the first and move on rather than shipping duplicates.
        if prompt in seen:
            continue
        seen.add(prompt)
        rows.append({
            "id": f"{prefix}{len(rows) + 1:04d}",
            "category": category,
            "prompt": prompt,
            "slots": {"topic": topic, "constraint": constraint if "{constraint}" in template else None},
        })
    return rows


def write_jsonl(path: Path, rows: list[dict]) -> str:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--out-dir", default="seeds")
    parser.add_argument("--train", type=int, default=600, help="prompts the teacher will answer")
    parser.add_argument("--heldout", type=int, default=120, help="prompts kept back for evaluation")
    parser.add_argument("--maths", type=int, default=300, help="verifiable word problems")
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    rng = random.Random(args.seed)
    out = Path(args.out_dir)

    total = args.train + args.heldout
    everything = build_prompts(total, rng, prefix="p")
    if len(everything) < total:
        print(f"note: the taxonomy yields {len(everything)} distinct prompts; "
              f"asked for {total}. Add topics or templates rather than sampling with repeats.")
    train_rows = everything[:args.train]
    heldout_rows = everything[args.train:]
    for i, row in enumerate(heldout_rows, start=1):
        row["id"] = f"h{i:04d}"

    maths_rows = []
    seen_maths: set[str] = set()
    attempts = 0
    while len(maths_rows) < args.maths and attempts < args.maths * 50:
        attempts += 1
        text, answer = maths_problem(rng)
        if text in seen_maths:
            continue
        seen_maths.add(text)
        maths_rows.append({
            "id": f"m{len(maths_rows) + 1:04d}",
            "category": "maths",
            "prompt": f"{MATHS_INSTRUCTION}\n\n{text}",
            "answer": answer,
        })

    hashes = {
        "prompts.jsonl": write_jsonl(out / "prompts.jsonl", train_rows),
        "heldout.jsonl": write_jsonl(out / "heldout.jsonl", heldout_rows),
        "maths.jsonl": write_jsonl(out / "maths.jsonl", maths_rows),
    }

    categories: dict[str, int] = {}
    for row in train_rows:
        categories[row["category"]] = categories.get(row["category"], 0) + 1

    print(f"written to {out}/")
    print(f"  prompts.jsonl  {len(train_rows)} prompts")
    for name in sorted(categories):
        print(f"      {name:<12} {categories[name]}")
    print(f"  heldout.jsonl  {len(heldout_rows)} prompts, never generated on")
    print(f"  maths.jsonl    {len(maths_rows)} problems with computed answers")
    print()
    for name, digest in hashes.items():
        print(f"sha256({name}) = {digest[:16]}...")
    print("Record those hashes: they are what tie every later stage to this exact seed set.")


if __name__ == "__main__":
    main()
