#!/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]



# ---------------------------------------------------------------------------
# RULER PARITY ADDITIONS (2026-09-17)
#
# RULER is 13 tasks in 4 categories: 8 NIAH retrieval variants, 1 multi-hop
# tracing, 2 aggregation, 2 QA. The six tasks above cover 4 of the 8 NIAH slots
# plus VT. Everything below fills the rest.
#
# The NIAH variants differ only in NEEDLE TYPE (7-digit number vs UUID) and
# HAYSTACK TYPE (real text vs a repeated noise sentence), so they are
# parameterisations rather than new mechanisms -- which is exactly how RULER
# defines them.
#
# CWE and FWE are structurally different: the body IS the data, not a haystack
# with needles planted in it. They are listed in BODY_TASKS and build_one takes
# a separate path for them.
# ---------------------------------------------------------------------------

def rnd_uuid(rng):
    h = "0123456789abcdef"
    f = lambda n: "".join(rng.choice(h) for _ in range(n))
    return f"{f(8)}-{f(4)}-{f(4)}-{f(4)}-{f(12)}"


# Tasks whose haystack must be the repeated-noise filler rather than real text.
NOISE_TASKS = {"niah_single_1"}
# Tasks that generate their own body instead of planting needles in a haystack.
BODY_TASKS = {"cwe", "fwe", "qa_1", "qa_2"}


def t_niah_single_1(rng, depth):
    """S-NIAH-1: word-number needle in REPEATED NOISE (see NOISE_TASKS)."""
    return t_niah_single(rng, depth)


def t_niah_single_3(rng, depth):
    """S-NIAH-3: word-UUID needle. UUIDs are far longer and less token-friendly
    than 7-digit numbers, so this is the hardest single-needle variant."""
    k, v = rng.choice(WORDS), rnd_uuid(rng)
    s = f" The special magic uuid for {k} is {v}. "
    q = f"\nQuestion: What is the special magic uuid for {k}?\n"
    a = f"Answer: The special magic uuid for {k} is"
    return [(s, depth)], q, a, [v]


def t_niah_multikey_2(rng, depth):
    """MK-NIAH-2: 8 distractor keys instead of 4."""
    return t_niah_multikey(rng, depth, n_keys=8)


def t_niah_multikey_3(rng, depth, n_keys=8):
    """MK-NIAH-3: 8 keys AND UUID values."""
    keys = rng.sample(WORDS, n_keys)
    vals = [rnd_uuid(rng) for _ in keys]
    tgt = rng.randrange(n_keys)
    sents = [(f" The special magic uuid for {k} is {v}. ",
              depth if i == tgt else rng.random())
             for i, (k, v) in enumerate(zip(keys, vals))]
    k, v = keys[tgt], vals[tgt]
    q = f"\nQuestion: What is the special magic uuid for {k}?\n"
    a = f"Answer: The special magic uuid for {k} is"
    return sents, q, a, [v]


def t_niah_multiquery(rng, depth, n_keys=4):
    """MQ-NIAH: several keys planted, ALL of them queried at once. Distinct from
    multikey (one target among distractors) and multivalue (one key, many
    values): here every planted pair must be retrieved."""
    keys = rng.sample(WORDS, n_keys)
    vals = [rnd_number(rng) for _ in keys]
    sents = [(f" The special magic number for {k} is {v}. ",
              depth if i == 0 else rng.random())
             for i, (k, v) in enumerate(zip(keys, vals))]
    q = ("\nQuestion: What are the special magic numbers for "
         + ", ".join(keys) + "?\n")
    a = "Answer: The special magic numbers are"
    return sents, q, a, vals          # every value must appear


# ---- aggregation ----------------------------------------------------------
# These return (body_text, question, answer_prefix, expected). build_one
# tokenises body_text directly and does NOT use the haystack.

CWE_COMMON, CWE_TOPK = 10, 10
FWE_TOPK = 3


def t_cwe(rng, depth, body_budget_tok=4096):
    """CWE: a word list in which CWE_COMMON words recur many times and the rest
    appear once. The model must name the top-K most common.

    Word pool is built from WORDS x a digit suffix so every item is a single
    common English token plus a digit -- keeps tokenisation tight and avoids
    the model recognising real-word frequency from pretraining."""
    pool = [f"{w}{i}" for w in WORDS for i in range(10)]
    rng.shuffle(pool)
    common = pool[:CWE_COMMON]
    rest = pool[CWE_COMMON:]
    # ~4 tokens per "word," item; reserve a little headroom.
    n_items = max(CWE_COMMON * 6, int(body_budget_tok / 4.2))
    reps = max(3, (n_items // 2) // max(1, CWE_COMMON))
    items = [c for c in common for _ in range(reps)]
    singles = rest * (max(0, n_items - len(items)) // max(1, len(rest)) + 1)
    items += singles[: max(0, n_items - len(items))]
    rng.shuffle(items)
    body = " " + ", ".join(items) + ". "
    q = f"\nQuestion: What are the {CWE_TOPK} most common words in the list above?\n"
    a = "Answer: The most common words are"
    return body, q, a, common          # all K must appear


def t_fwe(rng, depth, body_budget_tok=4096, alpha=2.0):
    """FWE: words drawn from a Zeta(alpha) distribution over a shuffled pool;
    the model must name the top-3 most frequent. Unlike CWE the frequencies are
    graded rather than two-tier, so ranking matters."""
    pool = [f"{w}{i}" for w in WORDS for i in range(10)]
    rng.shuffle(pool)
    n_items = max(60, int(body_budget_tok / 4.2))
    # Zeta: P(rank r) ~ r^-alpha over the shuffled pool.
    wts = [(r + 1) ** (-alpha) for r in range(len(pool))]
    tot = sum(wts)
    cum, acc = [], 0.0
    for w in wts:
        acc += w / tot
        cum.append(acc)
    def draw():
        u = rng.random()
        lo, hi = 0, len(cum) - 1
        while lo < hi:
            mid = (lo + hi) // 2
            if cum[mid] < u: lo = mid + 1
            else: hi = mid
        return pool[lo]
    items = [draw() for _ in range(n_items)]
    counts = {}
    for it in items: counts[it] = counts.get(it, 0) + 1
    top = [w for w, _ in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))][:FWE_TOPK]
    rng.shuffle(items)
    body = " " + ", ".join(items) + ". "
    q = f"\nQuestion: What are the {FWE_TOPK} most frequently appearing words in the list above?\n"
    a = "Answer: The most frequent words are"
    return body, q, a, top



# ---- QA with distractors (RULER qa_1 / qa_2) ------------------------------
# RULER "extends traditional QA datasets by adding distracting paragraphs".
# The GOLD paragraph sits at `depth`; the rest of the body is distractor
# paragraphs drawn from other questions in the same dataset.
#
# HotpotQA yes/no answers are DROPPED: scoring is substring match on normalised
# text, and "yes" matches almost any generation, which would score noise as a hit.

QA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "qa_data")
_QA_CACHE = {}


def _load_qa(kind):
    if kind in _QA_CACHE:
        return _QA_CACHE[kind]
    import pyarrow.parquet as pq
    if kind == "squad":
        t = pq.read_table(os.path.join(QA_DIR, "squad_val.parquet")).to_pydict()
        rows = []
        for ctx, qn, ans in zip(t["context"], t["question"], t["answers"]):
            a = ans["text"][0] if ans["text"] else ""
            if len(a) >= 3:
                rows.append((ctx, qn, a))
    else:
        t = pq.read_table(os.path.join(QA_DIR, "hotpot_val.parquet")).to_pydict()
        rows = []
        for ctx, qn, a in zip(t["context"], t["question"], t["answer"]):
            if a.lower() in ("yes", "no") or len(a) < 3:
                continue                      # see note above
            paras = [" ".join(sents) for sents in ctx["sentences"]]
            rows.append(("\n\n".join(paras), qn, a))
    if not rows:
        raise SystemExit(f"no usable {kind} rows in {QA_DIR}")
    _QA_CACHE[kind] = rows
    return rows


def _t_qa(kind, rng, depth, body_budget_tok):
    rows = _load_qa(kind)
    gold_ctx, question, answer = rows[rng.randrange(len(rows))]
    gold = " " + gold_ctx.strip() + " "
    # distractors from OTHER questions, appended until the budget is full
    need = max(0, body_budget_tok - len(ENC.encode(gold)))
    pre, post, got = [], [], 0
    want_pre = int(max(0.0, min(1.0, depth)) * need)
    while got < need:
        d = rows[rng.randrange(len(rows))][0]
        if d == gold_ctx:
            continue
        piece = " " + d.strip() + " "
        n = len(ENC.encode(piece))
        (pre if sum(len(ENC.encode(x)) for x in pre) < want_pre else post).append(piece)
        got += n
    body = "".join(pre) + gold + "".join(post)
    q = f"\nQuestion: {question}\n"
    a = "Answer:"
    return body, q, a, [answer]


def t_qa_1(rng, depth, body_budget_tok=4096):
    """SQuAD paragraph + SQuAD distractors."""
    return _t_qa("squad", rng, depth, body_budget_tok)


def t_qa_2(rng, depth, body_budget_tok=4096):
    """HotpotQA (multi-hop, 2 gold paragraphs) + distractors."""
    return _t_qa("hotpot", rng, depth, body_budget_tok)

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,
    # --- RULER parity (2026-09-17) ---
    "niah_single_1":     t_niah_single_1,
    "niah_single_3":     t_niah_single_3,
    "niah_multikey_2":   t_niah_multikey_2,
    "niah_multikey_3":   t_niah_multikey_3,
    "niah_multiquery":   t_niah_multiquery,
    "cwe":               t_cwe,
    "fwe":               t_fwe,
    "qa_1":              t_qa_1,
    "qa_2":              t_qa_2,
}


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.
    """
    # BODY_TASKS (cwe/fwe) generate their own body and need to know the budget,
    # so they are called with it; everything else keeps the original signature.
    if task in BODY_TASKS:
        probe_q, probe_a = TASKS[task](random.Random(0), depth, 64)[1:3]
        tail0 = len(ENC.encode(probe_q)) + len(ENC.encode(probe_a))
        body_text, q, a, expected = TASKS[task](rng, depth, max(64, prompt_len - tail0))
    else:
        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}")

    if task in BODY_TASKS:
        # The body IS the data. Repeat-or-truncate to hit the exact length; the
        # C++ reader slices at a fixed stride so drift is fatal.
        b = ENC.encode(body_text)
        if len(b) < body_budget:
            b = (b * (body_budget // max(1, len(b)) + 2))[:body_budget]
        out = b[:body_budget] + q_ids + a_ids
        assert len(out) == prompt_len, (len(out), prompt_len)
        return out, expected

    # 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.
        # NOISE_TASKS (S-NIAH-1) deliberately use the repeated-noise filler
        # instead of real text -- that IS the variant's definition in RULER.
        hay_path = "" if task in NOISE_TASKS else args.haystack
        hay = load_haystack(hay_path, 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()