#!/usr/bin/env python3
"""
in2_gate.py -- the EXACT layout gate for in2_gen.py output. No GPU, no CUDA.

This is the only exact pre-flight check in the IN2 pilot. The step-0 loss check
was demoted to a diagnostic because its two hypotheses sit only ~0.22 nats apart
(no-induction 4.147 vs strong-induction 3.925, against a 3.901 plain baseline),
which cannot validate a layout. These assertions can.

Checks, per arm dir:
  1  file size is exactly examples*T*2 bytes (headerless uint16)
  2  every token in [0, VOCAB)
  3  no answer span intrudes into the last PAD_TAIL offsets of its example
  4  planted spans (needle / probe+answer / decoy) never overlap
  5  RETRIEVAL IS ACTUALLY POSSIBLE: the tokens at the probe equal the needle's
     key, and the tokens at the answer equal the needle's value. If this fails
     the task is unsolvable and the whole experiment is void.
  6  batch alignment: re-derive x/y exactly as DataLoaderLite::next_batch does
     (x=buf[i], y=buf[i+1], stride B*T*dp) and assert every example boundary
     lands on a row boundary
  7  arm-specific distance invariants (b: all == span, c: all >= d_min)

Cross-arm:
  8  arms B and C have BIT-IDENTICAL answer_start histograms -- the discipline
     the whole B-vs-C contrast rests on

Usage
  python3 in2_gate.py --dirs /tmp/in2b /tmp/in2c --B 2 --dp 2
"""
import argparse, array, json, os, sys

VOCAB = 50304
fails = []


def check(cond, msg):
    print(("  ok   " if cond else "  FAIL ") + msg)
    if not cond:
        fails.append(msg)


def load(d, split="train"):
    js = [f for f in os.listdir(d) if f.endswith(".json") and f"_{split}_" in f]
    if len(js) != 1:
        raise SystemExit(f"{d}: expected exactly one {split} manifest, found {js}")
    man = json.load(open(os.path.join(d, js[0])))
    toks = array.array("H")
    for sp in man.get("shards", [man["bin"]]):
        with open(sp, "rb") as f:
            toks.fromfile(f, os.path.getsize(sp) // 2)
    return man, toks


def gate_one(d, B, dp, split="train"):
    man, toks = load(d, split)
    T, n_ex, arm, span = man["T"], man["examples"], man["arm"], man["span"]
    kl, vl, pad = man["key_len"], man["val_len"], man["pad_tail"]
    print(f"\n=== arm {arm}  T={T}  examples={n_ex}  K={man['K']} ===")

    tot = sum(os.path.getsize(sp) for sp in man.get("shards", [man["bin"]]))
    check(tot == n_ex * T * 2,
          f"1 total size across {len(man.get('shards',[1]))} shard(s) == examples*T*2 "
          f"({n_ex*T*2} bytes, got {tot})")
    lo, hi = min(toks), max(toks)
    check(lo >= 0 and hi < VOCAB, f"2 tokens in [0,{VOCAB}) (min {lo}, max {hi})")

    bad_pad = bad_ovl = bad_ret = bad_dist = 0
    for ei, recs in enumerate(man["records"]):
        base = ei * T
        spans = []
        for r in recs:
            a0, a1 = r["answer_start"], r["answer_start"] + vl
            if a1 > T - pad:
                bad_pad += 1
            spans.append((r["needle_start"], r["needle_start"] + span))
            spans.append((r["probe_start"], r["probe_start"] + span))
            # 5: the probe must reproduce the needle's key, the answer its value
            nk = toks[base + r["needle_start"]: base + r["needle_start"] + kl]
            nv = toks[base + r["needle_start"] + kl: base + r["needle_start"] + span]
            pk = toks[base + r["probe_start"]: base + r["probe_start"] + kl]
            av = toks[base + a0: base + a1]
            if nk != pk or nv != av:
                bad_ret += 1
            if arm == "b" and r["distance"] != span:
                bad_dist += 1
            if arm == "c" and r["distance"] < man["d_min"]:
                bad_dist += 1
        spans.sort()
        for (s0, e0), (s1, _) in zip(spans, spans[1:]):
            if s1 < e0:
                bad_ovl += 1
    check(bad_pad == 0, f"3 no answer within the last {pad} offsets ({bad_pad} bad)")
    check(bad_ovl == 0, f"4 planted spans never overlap ({bad_ovl} bad)")
    check(bad_ret == 0, f"5 RETRIEVAL POSSIBLE: probe==key and answer==value ({bad_ret} bad)")
    check(bad_dist == 0, f"7 arm-{arm} distance invariant ({bad_dist} bad)")

    # 6: DataLoaderLite windowing. next_batch reads B*T+1 tokens and sets
    # y[i]=buf[i+1]; pos_ advances by B*T*world. Example boundaries must land on
    # row boundaries or an example would straddle two rows.
    stride = B * T * dp
    check(stride % T == 0, f"6a batch stride {stride} is a multiple of T={T}")
    check((n_ex * T) % stride == 0,
          f"6b total {n_ex*T} tokens is a whole number of batches (stride {stride})")
    if len(toks) >= B * T + 1:
        x0 = toks[0:B * T]
        y0 = toks[1:B * T + 1]
        rows_ok = all(x0[r * T] == toks[r * T] for r in range(B))
        check(rows_ok, f"6c row r starts at token r*T for all B={B} rows")
        check(y0[T - 1] == toks[T], "6d y[T-1] targets the next example's first token "
                                    "(why PAD_TAIL exists)")
    return man


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dirs", nargs="+", required=True)
    ap.add_argument("--B", type=int, default=2)
    ap.add_argument("--dp", type=int, default=2)
    ap.add_argument("--split", default="train", choices=["train", "val"])
    args = ap.parse_args()

    mans = {}
    for d in args.dirs:
        m = gate_one(d, args.B, args.dp, args.split)
        mans[m["arm"]] = m

    if "b" in mans and "c" in mans:
        print("\n=== cross-arm ===")
        hb, hc = mans["b"]["answer_start_hist"], mans["c"]["answer_start_hist"]
        check(hb == hc, "8 arms B and C have BIT-IDENTICAL answer_start histograms "
                        f"({len(hb)} vs {len(hc)} distinct offsets)")
        check(mans["b"]["seed"] == mans["c"]["seed"], "8b same seed")
        check(mans["b"]["answer_tokens"] == mans["c"]["answer_tokens"],
              "8c same answer-token count")

    print(f"\n{'ALL PASS' if not fails else str(len(fails)) + ' FAILURE(S)'}")
    return 1 if fails else 0


if __name__ == "__main__":
    sys.exit(main())