"""THE CONTROL AXIS B WAS MISSING: how much does the author collide with HIMSELF?

memorization_check.py compares each arm against the TRAIN corpus and uses the
base-unadapted arm as the negative control. On Hemingway that control is weak in a way
it was not on Brontë, and the weakness runs one way only -- it makes an innocent arm
look guilty:

  * base-unadapted writes 18,035 words of summary prose; the adapted arms write 27,413
    of pastiche. Text that does not imitate the register cannot collide with its
    n-grams, so 0.00 measures "different register", not "did not memorise".
  * Hemingway's register IS short, plain, high-frequency English with heavy unattributed
    dialogue. An arm that SUCCEEDS at the voice task must start colliding with 8-grams
    built out of the commonest word sequences in the language.

So the honest reference is not the base arm. It is HELD-OUT HEMINGWAY -- val text no arm
trained on, written by the author himself, which by construction did not memorise the
train split. Whatever rate that produces is the floor the metric returns for innocent
text in this register, and only an excess over THAT is evidence of copying.

Chunks are cut to the generations' own length so the comparison is like-for-like: a
longer sample has more chances to collide.
"""
import json, re, sys, pathlib, statistics as st
CORP = pathlib.Path(sys.argv[1]); EVAL = pathlib.Path(sys.argv[2]); N = 8

def norm(t): return re.findall(r"[a-z']+", t.lower())

train_words, val_texts = [], []
for f in sorted(CORP.glob("*.copy0.jsonl")):
    for l in f.read_text(encoding="utf-8").splitlines():
        r = json.loads(l)
        (val_texts.append(r["text"]) if r.get("split") == "val" else train_words.extend(norm(r["text"])))
grams = {" ".join(train_words[i:i+N]) for i in range(len(train_words)-N+1)}
print(f"train (copy0, split=train): {len(train_words):,} words, {len(grams):,} distinct {N}-grams")

def longest(w):
    best = 0; i = 0
    while i <= len(w)-N:
        if " ".join(w[i:i+N]) in grams:
            k = N
            while i+k < len(w) and " ".join(w[i+k-N+1:i+k+1]) in grams: k += 1
            best = max(best, k); i += 1
        else: i += 1
    return best

arm_lens = []
for arm in ("ckpt1750", "ckpt850", "base"):
    p = EVAL/f"beats5.{arm}.jsonl"
    if p.exists():
        arm_lens += [len(norm(json.loads(l)["raw"]))
                     for l in p.read_text(encoding="utf-8").splitlines() if l.strip()]
CHUNK = int(st.median(arm_lens))
print(f"median generation length across arms: {CHUNK} words -- val is chunked to match\n")

vw = norm("\n".join(val_texts))
chunks = [vw[i:i+CHUNK] for i in range(0, len(vw)-CHUNK+1, CHUNK)]
longs = [longest(c) for c in chunks]
hits = sum(1 for x in longs if x >= N)
print(f"{'sample':<34} {'n':>5} {'hit-rate':>9} {'mean-longest':>13} {'max':>5}")
print("-"*70)
print(f"{'HELD-OUT HEMINGWAY (never trained)':<34} {len(chunks):>5} {hits/len(chunks):>9.2f} "
      f"{sum(longs)/len(longs):>13.1f} {max(longs):>5}")
for arm in ("base", "ckpt1750", "ckpt850"):
    p = EVAL/f"beats5.{arm}.jsonl"
    if not p.exists(): continue
    rows = [json.loads(l) for l in p.read_text(encoding="utf-8").splitlines() if l.strip()]
    L = [longest(norm(r["raw"])) for r in rows]
    h = sum(1 for x in L if x >= N)
    print(f"{arm:<34} {len(rows):>5} {h/len(rows):>9.2f} {sum(L)/len(L):>13.1f} {max(L):>5}")
print(f"\npositive control (a train slice vs train): longest = {longest(train_words[1000:1160])} "
      f"(must be large, else blind)")