"""Add words from EXAMPLES, no gradient: a new word's frozen code = sign of the average readout direction the
trained model produces just before the word, over up to K example sentences from the TRAINING text; prior = log of the
word's training frequency. Evaluated on HELD-OUT text: total nats per character over whole passages (words merged where
they occur) vs no new words. Also reports tokens saved."""
import sys, numpy as np, torch, torch.nn.functional as F
import new_words as nw
from transformers import AutoTokenizer
ARM, CKPT = (sys.argv[1], sys.argv[2]) if len(sys.argv) > 2 else ("B", "finals/B/step_07820.pt")
NW, K, CTX = 200, 50, 128
tok = AutoTokenizer.from_pretrained("google/gemma-3-270m")
counts = np.load("data/unigram_counts.npy"); val = np.load("data/val.npy", mmap_mode="r")
tr = np.load("data/train_000.npy", mmap_mode="r")[:100_000_000]
words = nw.corpus_words(tok, val, NW, 5)
wp = [tok(" " + w, add_special_tokens=False).input_ids for w in words]
m, tag = nw.load(ARM, CKPT, counts)
T0, S0 = m.table.clone(), m.seal.detach().clone(); V = T0.shape[0]
def occ(p, arr, limit):
    k = len(p); h = np.flatnonzero(arr[:len(arr) - k + 1] == p[0])
    for j in range(1, k): h = h[arr[h + j] == p[j]]
    return h, len(h)
codes, priors = [], []
with torch.no_grad():
    for p in wp:
        h, n = occ(p, tr, K)
        h = [x for x in h if x >= CTX][:K]
        if not h:
            v = torch.zeros(512, device="cuda")
        else:
            ctx = torch.tensor(np.stack([tr[x - CTX:x] for x in h]).astype(np.int64), device="cuda")
            with torch.autocast("cuda", dtype=torch.bfloat16):
                z = m.out(m.hidden(ctx)[:, -1]).float()          # readout direction just before the word
            v = z.mean(0)
        c = torch.sign(v); c[c == 0] = 1; codes.append(c)
        priors.append(np.log((n * 20.5 + 1) / (counts.sum() + 262144)))   # 100M-token shard scaled to 2.05B
codes = torch.stack(codes); priors = torch.tensor(priors, dtype=torch.float32, device="cuda")
NSEQ, L = 600, 512
seqs = [np.asarray(val[i * L:(i + 1) * L]).tolist() for i in range(NSEQ)]
chars = sum(len(tok.decode(s)) for s in seqs)
order = sorted(range(len(wp)), key=lambda i: -len(wp[i]))
def merge(s, ids):
    out, i = [], 0
    while i < len(s):
        for j in order:
            p = wp[j]
            if s[i:i + len(p)] == p: out.append(ids[j]); i += len(p); break
        else: out.append(s[i]); i += 1
    return out
@torch.no_grad()
def total(seq_list):
    t = 0.0
    for s in seq_list:
        x = torch.tensor([[1] + s], device="cuda")
        with torch.autocast("cuda", dtype=torch.bfloat16):
            lg = m.logits(m.hidden(x[:, :-1])).float()
        t += float(F.cross_entropy(lg[0], x[0, 1:], reduction="sum"))
    return t
m.table, m.seal = T0, S0
base = total(seqs)
ids = list(range(V, V + len(wp)))
merged = [merge(s, ids) for s in seqs]
saved = sum(map(len, seqs)) - sum(map(len, merged))
for name, cd, pr in (("examples code + frequency prior", codes, priors),
                     ("piece-mean code + frequency prior", None, priors)):
    if cd is None:
        E = nw.gemma_embedding(); vv = torch.stack([E[p].mean(0) for p in wp]).cuda()
        cd = torch.sign((vv - m.table_mu) @ m.table_W); cd[cd == 0] = 1
    m.table = torch.cat([T0, cd]); m.seal = torch.cat([S0, pr])
    t = total(merged)
    print(f"{tag} | {name}: {t / chars:.4f} vs {base / chars:.4f} nats/char ({(t - base) / chars:+.4f}, {(t - base) / base * 100:+.2f}%) | "
          f"{saved} tokens saved of {sum(map(len, seqs))}", flush=True)