#!/usr/bin/env python3
"""
longeval_gen.py -- build GENERATIVE long-context benchmark prompts for bluscriptCP.

Unlike needle_gen.py (which plants a random-token key/value pair and is scored by
teacher-forced NLL over an answer span that is ALREADY IN THE INPUT), everything
here is scored by what the model actually GENERATES. That is what RULER, NIAH and
VAL Probing measure, and it is the only form comparable to published numbers.

WHY THE PROMPTS LOOK LIKE THIS
  The target is a ~1B BASE language model with no instruction tuning, so a bare
  "Question: ...\\nAnswer:" often fails for reasons that have nothing to do with
  retrieval. Every task therefore ends with an answer PREFIX that makes the
  correct answer the natural continuation, e.g.

      ... haystack ...
      Question: What is the special magic number for aardvark?
      Answer: The special magic number for aardvark is

  The model only has to emit the digits. Scoring looks for the expected string
  anywhere in the generated tokens, so trailing chatter does not cost a point.

TASKS
  niah_single      RULER NIAH-single      one needle, one key, retrieve its value
  niah_multikey    RULER NIAH-multikey    K needles, retrieve ONE named key
  niah_multivalue  RULER NIAH-multivalue  one key with V values, retrieve all
  vt               RULER VT               chain of variable assignments, resolve
  val_forward      VAL forward retrieval  key stated before the query (default)
  val_backward     VAL backward retrieval value stated first, retrieve its KEY

OUTPUT (per task, into --out-dir)
  <task>_prompts.bin   int32 token stream, N prompts x --prompt-len tokens, end to end
                       (exactly what CP_GEN_TOKENS_BIN expects)
  <task>_meta.json     per-prompt expected answers + placement metadata

The prompt length is EXACT: the haystack is trimmed or extended so every prompt is
--prompt-len tokens, because the C++ reader slices the stream at a fixed stride.
"""

import argparse
import json
import os
import random
import struct
import sys

try:
    import tiktoken
except ImportError:
    sys.exit("longeval_gen.py needs tiktoken:  pip install tiktoken")

ENC = tiktoken.get_encoding("gpt2")


# ---------------------------------------------------------------------------
# Haystack: real text from the flux shards (uint16 GPT-2 token stream), so the
# distractor context has the same distribution the model was trained on. Falls
# back to a synthetic filler sentence when no shard is given, which keeps the
# generator usable on a box with no data.
# ---------------------------------------------------------------------------
FALLBACK = ENC.encode(
    "The grass swayed in the long afternoon light while the river ran on "
    "unhurried past the old stone bridge and the fields beyond. "
)


def load_haystack(path, need):
    """Return >= `need` GPT-2 token ids of filler.

    SLICE BEFORE CONVERTING. The flux val shard is 200 MB = 100M uint16 tokens;
    `toks.astype(np.int64).tolist()` over the whole array costs ~800 MB of numpy
    plus a ~2.8 GB Python list. Callers now ask for prompt_len*(n+2) so each
    prompt can start at its own offset, which makes the slice essential.
    """
    if not path:
        reps = need // len(FALLBACK) + 2
        return (FALLBACK * reps)[:need + len(FALLBACK)]
    import numpy as np
    toks = np.fromfile(path, dtype=np.uint16)
    if toks.size < need:
        reps = int(need // toks.size) + 2
        toks = np.tile(toks, reps)
    return toks[:need].astype(np.int64).tolist()


# vt's chain needs room AFTER its anchor for the remaining links, so the anchor
# cannot sit at depth 1.0. Both t_vt and the metadata writer call this, so the
# recorded depth can never drift from the depth actually built.
VT_ANCHOR_MAX = 0.8


def effective_depth(task, depth):
    return min(depth, VT_ANCHOR_MAX) if task == "vt" else depth


# ---------------------------------------------------------------------------
# Needle vocabulary. Words are common English so they tokenize compactly and the
# model has seen them; numbers are 7 digits so a lucky guess is ~1e-7.
# ---------------------------------------------------------------------------
WORDS = [
    "aardvark", "buffalo", "cactus", "dolphin", "emerald", "falcon", "glacier",
    "harbor", "iceberg", "jaguar", "kestrel", "lantern", "meadow", "nutmeg",
    "orchid", "penguin", "quartz", "raven", "sapphire", "tundra", "umbrella",
    "violet", "walnut", "xylophone", "yarrow", "zephyr", "anchor", "bramble",
    "cinder", "driftwood", "ember", "fjord", "granite", "hollow", "ivory",
]


def rnd_number(rng):
    return str(rng.randint(1000000, 9999999))


# ---------------------------------------------------------------------------
# Task builders. Each returns (needle_sentences, question, answer_prefix, expected)
#   needle_sentences : list of (sentence_text, depth_fraction) to plant
#   expected         : list of strings; ALL must appear in the generation to score 1
# ---------------------------------------------------------------------------
def t_niah_single(rng, depth):
    k, v = rng.choice(WORDS), rnd_number(rng)
    s = f" The special magic number for {k} is {v}. "
    q = f"\nQuestion: What is the special magic number for {k}?\n"
    a = f"Answer: The special magic number for {k} is"
    return [(s, depth)], q, a, [v]


def t_niah_multikey(rng, depth, n_keys=4):
    keys = rng.sample(WORDS, n_keys)
    vals = [rnd_number(rng) for _ in keys]
    # Distractors spread across the whole context; the TARGET sits at `depth`.
    sents = []
    tgt = rng.randrange(n_keys)
    for i, (k, v) in enumerate(zip(keys, vals)):
        d = depth if i == tgt else rng.random()
        sents.append((f" The special magic number for {k} is {v}. ", d))
    k, v = keys[tgt], vals[tgt]
    q = f"\nQuestion: What is the special magic number for {k}?\n"
    a = f"Answer: The special magic number for {k} is"
    return sents, q, a, [v]


def t_niah_multivalue(rng, depth, n_vals=3):
    k = rng.choice(WORDS)
    vals = [rnd_number(rng) for _ in range(n_vals)]
    # All values share one key, scattered; depth anchors the FIRST one.
    sents = [(f" One of the special magic numbers for {k} is {v}. ",
              depth if i == 0 else rng.random())
             for i, v in enumerate(vals)]
    q = f"\nQuestion: What are all the special magic numbers for {k}?\n"
    a = f"Answer: The special magic numbers for {k} are"
    return sents, q, a, vals          # every value must appear


def t_vt(rng, depth, hops=4):
    """RULER variable tracking: X1 = <num>, X2 = X1, X3 = X2 ... resolve the chain.

    Depths must INCREASE along the chain. build_one sorts planted sentences by
    depth, so random depths would scatter the chain in random document order --
    that is backward-reference resolution, strictly harder than RULER's ordered
    VT and not comparable to published numbers. Monotone depths put document
    order = chain order (Python's sort is stable, so ties keep chain order).
    """
    names = rng.sample([w.upper()[:3] + str(rng.randint(10, 99)) for w in WORDS], hops)
    val = rnd_number(rng)
    anchor = effective_depth("vt", depth)
    sents = [(f" VAR {names[0]} = {val}. ", anchor)]
    for i in range(1, hops):
        sents.append((f" VAR {names[i]} = {names[i-1]}. ",
                      anchor + (1.0 - anchor) * i / hops))
    q = f"\nQuestion: What is the value of VAR {names[-1]}?\n"
    a = f"Answer: VAR {names[-1]} ="
    return sents, q, a, [val]


def t_val_forward(rng, depth):
    return t_niah_single(rng, depth)


def t_val_backward(rng, depth):
    """VAL backward retrieval: the VALUE is the query, the KEY is the answer."""
    k, v = rng.choice(WORDS), rnd_number(rng)
    s = f" The special magic number for {k} is {v}. "
    q = f"\nQuestion: Which item has the special magic number {v}?\n"
    a = f"Answer: The item with the special magic number {v} is"
    return [(s, depth)], q, a, [k]


TASKS = {
    "niah_single":     t_niah_single,
    "niah_multikey":   t_niah_multikey,
    "niah_multivalue": t_niah_multivalue,
    "vt":              t_vt,
    "val_forward":     t_val_forward,
    "val_backward":    t_val_backward,
}


def build_one(task, rng, prompt_len, hay, depth, hay_off=0):
    """One prompt of EXACTLY prompt_len tokens.

    hay_off shifts this prompt's filler window. Without it every prompt starts
    the haystack at 0 and differs only in the needle, making the N prompts
    replicates rather than independent draws -- which inflates apparent
    precision because the effective sample size is far below N.
    """
    sents, q, a, expected = TASKS[task](rng, depth)
    q_ids, a_ids = ENC.encode(q), ENC.encode(a)
    tail = len(q_ids) + len(a_ids)
    body_budget = prompt_len - tail
    if body_budget <= 0:
        raise SystemExit(f"--prompt-len {prompt_len} too small for task {task}")

    # Plant each needle at its depth fraction of the BODY, then fill with haystack.
    planted = sorted(((max(0.0, min(1.0, d)), ENC.encode(s)) for s, d in sents),
                     key=lambda x: x[0])
    total_needle = sum(len(ids) for _, ids in planted)
    fill = body_budget - total_needle
    if fill <= 0:
        raise SystemExit(f"needles do not fit in --prompt-len {prompt_len}")

    H = lambda lo, n: hay[hay_off + lo: hay_off + lo + max(0, n)]

    body, used, hay_pos = [], 0, 0
    for frac, ids in planted:
        want = int(frac * fill)
        take = max(0, want - used)
        body.extend(H(hay_pos, take)); hay_pos += take; used += take
        body.extend(ids)
    body.extend(H(hay_pos, fill - used))

    out = body + q_ids + a_ids
    # Exact length: the C++ reader slices at a fixed stride, so drift is fatal.
    if len(out) < prompt_len:
        out = body + H(hay_pos, prompt_len - len(out)) + q_ids + a_ids
    out = out[:prompt_len - tail] + q_ids + a_ids
    assert len(out) == prompt_len, (len(out), prompt_len)
    return out, expected


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--task", required=True, choices=sorted(TASKS) + ["all"])
    ap.add_argument("--prompt-len", type=int, required=True,
                    help="exact tokens per prompt; must leave room for CP_GEN_MAX_NEW under CP_T")
    ap.add_argument("--n", type=int, default=50, help="prompts per task")
    ap.add_argument("--depths", default="0.1,0.3,0.5,0.7,0.9",
                    help="needle depth fractions, cycled across prompts")
    ap.add_argument("--haystack", default="", help="flux .bin shard for filler text")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out-dir", required=True)
    args = ap.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)
    depths = [float(d) for d in args.depths.split(",")]
    tasks = sorted(TASKS) if args.task == "all" else [args.task]

    for task in tasks:
        rng = random.Random(args.seed)          # per-task, so tasks are independent
        # (n + 2) windows so every prompt gets its own filler offset
        hay = load_haystack(args.haystack, args.prompt_len * (args.n + 2))
        bin_p = os.path.join(args.out_dir, f"{task}_prompts.bin")
        meta = []
        with open(bin_p, "wb") as f:
            for i in range(args.n):
                want = depths[i % len(depths)]
                eff = effective_depth(task, want)
                ids, expected = build_one(task, rng, args.prompt_len, hay, want,
                                          hay_off=i * args.prompt_len)
                f.write(struct.pack(f"<{len(ids)}i", *ids))
                # Record the EFFECTIVE depth: longeval_score.py groups on this
                # field, and depth-pooling is the primary unit of analysis, so a
                # clamped vt row must not be counted in the depth-1.0 cell.
                meta.append({"i": i, "depth": eff, "depth_requested": want,
                             "expected": expected})
        with open(os.path.join(args.out_dir, f"{task}_meta.json"), "w") as f:
            json.dump({"task": task, "prompt_len": args.prompt_len, "n": args.n,
                       "seed": args.seed, "items": meta}, f, indent=1)
        print(f"  {task:<16} {args.n} x {args.prompt_len} tok -> {bin_p}")


if __name__ == "__main__":
    main()