#!/usr/bin/env python3
"""
needle_gen.py -- synthesize retrieval-probe token streams for the existing
CP_EVAL_PPL / CP_EVAL_TOKENS_BIN path. No C++ changes are needed: the eval
already windows a raw int32 token file and already reports PER-POSITION
mean_nll, so planting an answer span at a known offset lets us read retrieval
NLL straight out of the CSV.

Stdlib only (array + random) -- deliberately no numpy, because the local pyenv
python3 lacks it while /usr/bin/python3 has it, and this must run identically on
the local box and the server regardless of which `python3` resolves first.

FIXED-SLOT LAYOUT (the construction that makes the differencing valid)
---------------------------------------------------------------------
Every variant of a trial has byte-identical structure except for WHERE the
needle sits (or whether it is present). The probe prefix and the answer span are
always at the SAME absolute offsets, so we always difference NLLs measured at
the same point on a position-varying curve:

  trial (T tokens):
    [ filler .......... N0..N(k-1) ......... filler ][ probe prefix ][ ANSWER ][pad]
                        ^needle, k tokens                 k/2           k/2     >=2
                        slot_start = round(p_frac * max_slot_start)

    p_frac = 0.0 -> needle as far from the probe as possible
    p_frac = 1.0 -> needle immediately adjacent to the probe (the RETRIEVAL FLOOR)

VARIANTS (7 per model/T/cache)
  p0, p0.25, p0.5, p0.75, p1   needle planted at that slot; p1 is the floor
  absent                       no needle planted; probe prefix and answer tokens
                               still present (the no-information ceiling)
  null                         needle planted at p_frac=0 but the probe prefix is
                               an UNRELATED k/2-gram, so retrieval is impossible
                               -> supplies the significance threshold (RAW NLL)

CRITICAL OFF-BY-ONE
  The eval scores input[t] = toks[base+t] against target[t] = toks[base+t+1] and
  accumulates the NLL at the position of the INPUT. So the NLL of the answer token
  living at trial offset `a` appears in the CSV at row `a-1`. The manifest records
  BOTH; needle_score.py must read `csv_positions`, never `answer_offsets`.

Usage
  python3 needle_gen.py --T 32768 --variant p0.5 --block 0 --out-dir needles
  python3 needle_gen.py --T 32768 --all --blocks 5 --out-dir needles
"""
import argparse, array, json, os, random, sys

VOCAB = 50304          # must match cfg.vocab_size; the C++ side bound-checks too
PAD_TAIL = 4           # >=2: windows stride by exactly T, so the target at T-1
                       # belongs to the NEXT trial. Keep the answer clear of it.
P_FRACS = [0.0, 0.25, 0.5, 0.75, 1.0]
MIX = 1_000_003        # odd multiplier for deterministic seed derivation


# Active p_frac grid. Defaults to P_FRACS; --p-fracs overrides it for the S>R
# ablation, whose predicted dead gap [R,S) falls BETWEEN the default points and
# would otherwise be unobservable. Overriding changes only which slots are probed,
# never the layout, so answer-span offsets stay invariant.
_ACTIVE_P_FRACS = list(P_FRACS)


def set_p_fracs(vals):
    global _ACTIVE_P_FRACS
    _ACTIVE_P_FRACS = list(vals)


def variant_names():
    return [f"p{p:g}" for p in _ACTIVE_P_FRACS] + ["absent", "null"]


def parse_variant(name):
    """-> (p_frac or None, is_null). absent -> (None, False); null -> (0.0, True)."""
    if name == "absent":
        return None, False
    if name == "null":
        return 0.0, True
    if name.startswith("p"):
        p = float(name[1:])
        if not 0.0 <= p <= 1.0:
            raise SystemExit(f"p_frac must be in [0,1] (got {p})")
        return p, False
    raise SystemExit(f"unknown variant '{name}' (expected one of {variant_names()})")


def check_platform():
    if sys.byteorder != "little":
        raise SystemExit("CP_EVAL_TOKENS_BIN expects int32 LITTLE-endian; this host is big-endian")
    if array.array("i").itemsize != 4 or array.array("H").itemsize != 2:
        raise SystemExit("unexpected array itemsizes ('i' must be 4B, 'H' must be 2B)")


def load_filler(path, need, rng):
    """Real tokens, so the LM is not off-distribution (fabricated filler would
    inflate NLL everywhere and mask the effect). Source is uint16; the eval file
    must be int32 LE, so widen explicitly and bound-check."""
    ntok = os.path.getsize(path) // 2
    if ntok < need + 1:
        raise SystemExit(f"{path} has {ntok} tokens, need >= {need + 1}")
    # One contiguous run per trial, so filler keeps local textual statistics
    # instead of being a shuffled bag of tokens.
    start = rng.randrange(0, ntok - need)
    buf = array.array("H")
    with open(path, "rb") as f:
        f.seek(start * 2)
        buf.fromfile(f, need)
    out = array.array("i", buf)                       # C-level uint16 -> int32
    lo, hi = min(out), max(out)
    if lo < 0 or hi >= VOCAB:
        raise SystemExit(f"{path} has ids outside [0,{VOCAB}): min={lo} max={hi}")
    return out


def build_layout(T, k, probe_pos=None):
    """Offsets are a pure function of (T, k, probe_pos) -- identical for every variant.

    probe_pos is None (default): probe+answer sit at the END of the trial, so the
    needle-to-probe DISTANCE is T - slot. Distance and the needle's ABSOLUTE
    position are then perfectly anti-correlated -- one variable, not two -- which
    makes "far from the question" and "early in the document" indistinguishable.

    probe_pos = P: probe+answer are planted at absolute offset P and everything
    after them is ordinary filler. Pinning the needle's slot while varying P (or
    vice versa) breaks that confound. Legal because the eval reports NLL at EVERY
    position, so the scored span does not have to be at the end -- and the
    manifest's csv_positions carries the offsets, so needle_score.py follows
    automatically with no change.
    """
    if k % 2:
        raise SystemExit(f"--needle-len must be even (got {k})")
    if probe_pos is None:
        answer_end = T - PAD_TAIL
        answer_start = answer_end - k // 2
        probe_start = answer_start - k // 2           # probe prefix: [probe_start, answer_start)
    else:
        probe_start = probe_pos
        answer_start = probe_start + k // 2
        answer_end = probe_start + k
        if answer_end > T - PAD_TAIL:
            raise SystemExit(f"--probe-pos {probe_pos} puts the answer span end at "
                             f"{answer_end}, past T-PAD_TAIL={T - PAD_TAIL}; the last "
                             "position's target belongs to the next trial")
    max_slot_start = probe_start - k                   # needle must fit before the probe
    if max_slot_start < 0:
        raise SystemExit(f"probe at {probe_start} leaves no room for a needle of "
                         f"{k} tokens before it (T={T})")
    return dict(
        T=T, needle_len=k, pad_tail=PAD_TAIL, probe_pos=probe_pos,
        probe_start=probe_start, answer_start=answer_start, answer_end=answer_end,
        max_slot_start=max_slot_start,
        answer_offsets=[answer_start, answer_end],
        # NLL is reported at the INPUT position, so the answer token at offset a is
        # scored on CSV row a-1. See "CRITICAL OFF-BY-ONE" above.
        csv_positions=[answer_start - 1, answer_end - 1],
    )


def make_trial(lay, p_frac, is_null, filler, rng, reverse=False):
    """One trial of exactly T int32 tokens. `filler` is consumed (modified)."""
    k = lay["needle_len"]
    half = k // 2
    trial = filler                                    # already a fresh copy per trial
    needle = array.array("i", [rng.randrange(VOCAB) for _ in range(k)])

    # The answer span is the needle's second half, at fixed offsets -- unless
    # `reverse`, which swaps the halves: probe = needle's SECOND half, answer = its
    # FIRST, so the model must retrieve what PRECEDES the probe. Forward retrieval
    # (answer following the keyword) is an induction-head copy and "an easily
    # learned pattern" (arXiv:2404.16811 s3.1), so it is exactly what training on
    # forward data would Goodhart; backward retrieval cannot be satisfied that way,
    # which makes it the transfer test. Layout OFFSETS are untouched, so
    # csv_positions stays valid and needle_score.py needs no change.
    trial[lay["answer_start"]:lay["answer_end"]] = needle[:half] if reverse else needle[half:]

    # Probe prefix = the needle's other half, EXCEPT for `null`, where an unrelated
    # k/2-gram makes retrieval impossible while the answer TARGETS stay identical.
    if is_null:
        prefix = array.array("i", [rng.randrange(VOCAB) for _ in range(half)])
    else:
        prefix = needle[half:] if reverse else needle[:half]
    trial[lay["probe_start"]:lay["answer_start"]] = prefix

    # Plant the needle in exactly one slot (or none, for `absent`).
    slot_start = None
    if p_frac is not None:
        slot_start = int(round(p_frac * lay["max_slot_start"]))
        trial[slot_start:slot_start + k] = needle
    return trial, slot_start


def gen(args, variant, block):
    p_frac, is_null = parse_variant(variant)
    lay = build_layout(args.T, args.needle_len, args.probe_pos)
    # Seed depends on (seed, T, block, trial) and NOT on the variant, so every
    # variant of a block shares identical filler and identical needle tokens --
    # the whole point of the fixed-slot construction.
    base_seed = (args.seed * MIX + args.T) * MIX + block

    stream = array.array("i")
    slots = []
    for i in range(args.trials):
        rng = random.Random(base_seed * MIX + i)
        filler = load_filler(args.filler, args.T, rng)
        trial, slot = make_trial(lay, p_frac, is_null, filler, rng, args.reverse)
        if len(trial) != args.T:
            raise SystemExit(f"trial length {len(trial)} != T={args.T} -- bug")
        stream.extend(trial)
        slots.append(slot)

    # +1 token: window w reads toks[base+t] and toks[base+t+1] for t in [0,T), so
    # the final window needs index N*T to exist.
    stream.append(stream[-1])
    if len(stream) != args.trials * args.T + 1:
        raise SystemExit("stream length wrong -- bug")

    os.makedirs(args.out_dir, exist_ok=True)
    tag = f"T{args.T}_{variant}_b{block}"
    bin_path = os.path.join(args.out_dir, f"needle_{tag}.bin")
    with open(bin_path, "wb") as f:
        stream.tofile(f)

    man = dict(bin=bin_path, variant=variant, block=block, p_frac=p_frac, is_null=is_null,
               trials=args.trials, seed=args.seed, vocab=VOCAB, reverse=args.reverse,
           slot_starts=slots, **lay)
    man_path = os.path.join(args.out_dir, f"needle_{tag}.json")
    with open(man_path, "w") as f:
        json.dump(man, f, indent=2)
    return bin_path, man_path, man


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--T", type=int, required=True, help="tokens per trial (= CP_T)")
    ap.add_argument("--trials", type=int, default=None,
                    help="trials per block (= CP_EVAL_WINDOWS). Default is the plan's "
                         "pinned N: 8 at T<=131072, 4 above.")
    ap.add_argument("--blocks", type=int, default=5,
                    help="blocks = SEPARATE files, each needing its own eval invocation; "
                         "the only source of spread (CP_EVAL_PPL emits no variance)")
    ap.add_argument("--block", type=int, default=None, help="generate just this block index")
    ap.add_argument("--variant", default=None, help=f"one of {variant_names()}")
    ap.add_argument("--all", action="store_true", help="every variant x every block")
    ap.add_argument("--needle-len", type=int, default=8, help="k; probe=k/2, answer=k/2")
    ap.add_argument("--filler", default="Data_Loader/Data/edufineweb_val_00000.bin")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--probe-pos", type=int, default=None,
                    help="absolute offset for the probe+answer span. Default: the end "
                         "of the trial. Setting it lets DISTANCE vary independently of "
                         "the needle's absolute position -- the confound the default "
                         "layout cannot separate. Isolate runs by --out-dir: filenames "
                         "are unchanged, so needle_score.py needs no flag.")
    ap.add_argument("--out-dir", default="needles")
    ap.add_argument("--reverse", action="store_true",
                    help="BACKWARD retrieval (transfer test): probe = needle's 2nd half, "
                         "answer = its 1st. Use a SEPARATE --out-dir -- filenames do not "
                         "encode it, and absent/null must be regenerated in reverse too.")
    ap.add_argument("--p-fracs", default=None,
                    help="comma-separated p_frac list overriding the default "
                         f"{P_FRACS}. Needed by the STRING S>R ablation, whose predicted "
                         "dead gap lies between the default points.")
    args = ap.parse_args()

    if args.p_fracs:
        vals = [float(x) for x in args.p_fracs.split(",") if x.strip()]
        if not vals or any(not 0.0 <= v <= 1.0 for v in vals):
            raise SystemExit(f"--p-fracs must be values in [0,1] (got {args.p_fracs})")
        set_p_fracs(vals)

    check_platform()
    if args.trials is None:                            # pinned N from the plan (A1)
        args.trials = 8 if args.T <= 131072 else 4

    if args.all:
        jobs = [(v, b) for v in variant_names() for b in range(args.blocks)]
    else:
        if args.variant is None:
            raise SystemExit("give --variant NAME (and optionally --block), or --all")
        jobs = [(args.variant, args.block if args.block is not None else 0)]

    for v, b in jobs:
        bp, _, man = gen(args, v, b)
        slot = man["slot_starts"][0]
        print(f"{v:8s} block={b}  {bp}  answer_offsets={man['answer_offsets']} "
              f"csv_positions={man['csv_positions']}  slot={slot}")


if __name__ == "__main__":
    main()