"""Net effect of adding new tokens on normal text: total nats per CHARACTER over whole held-out passages, with the new
words merged wherever they occur (everything else unchanged) vs the same passages without new tokens. Catches harm
elsewhere (a high prior makes new tokens over-predicted)."""
import numpy as np, torch, torch.nn.functional as F
import new_words as nw
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("google/gemma-3-270m")
counts = np.load("data/unigram_counts.npy"); val = np.load("data/val.npy", mmap_mode="r")
names = nw.make_names(tok, 12); words = nw.corpus_words(tok, val, 100, 5)
np_ = [tok(" " + w, add_special_tokens=False).input_ids for w in names]
wp = [tok(" " + w, add_special_tokens=False).input_ids for w in words]
E = nw.gemma_embedding()
m, _ = nw.load("B", "finals/B/step_07820.pt", counts)
T0, S0 = m.table.clone(), m.seal.detach().clone(); V = T0.shape[0]
NSEQ, L = 300, 512
seqs = [np.asarray(val[i * L: (i + 1) * L]).tolist() for i in range(NSEQ)]   # inside the text the words were picked from
chars = sum(len(tok.decode(s)) for s in seqs)
def merge(s, new_ids):
    out, i = [], 0
    while i < len(s):
        for p, nid in sorted(zip(wp, new_ids), key=lambda x: -len(x[0])):
            if s[i:i + len(p)] == p: out.append(nid); i += len(p); break
        else: out.append(s[i]); i += 1
    return out
@torch.no_grad()
def total_nats(seq_list):
    tot = 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()
        tot += float(F.cross_entropy(lg[0], x[0, 1:], reduction="sum"))
    return tot
m.table, m.seal = T0, S0
base = total_nats(seqs); print(f"no new tokens: {base / chars:.4f} nats/char", flush=True)
def codes(kind, pl):
    if kind == "first": c = T0[[p[0] for p in pl]].clone()
    else:
        v = torch.stack([(E[p].sum(0) if kind == "sum" else E[p].mean(0)) for p in pl]).cuda(); c = torch.sign((v - m.table_mu) @ m.table_W)
    c[c == 0] = 1; return c
tr = np.load("data/train_000.npy", mmap_mode="r")[:100_000_000]
def seq_count(p):
    k = len(p); h = np.flatnonzero(tr[:len(tr) - k + 1] == p[0])
    for j in range(1, k): h = h[tr[h + j] == p[j]]
    return len(h)
wc = np.array([seq_count(p) for p in wp]) * (2.05e9 / 1e8)
freq_prior = torch.tensor(np.log((wc + 1) / (counts.sum() + 262144)), dtype=torch.float32, device="cuda")
print("word counts (scaled to 2.05B): median", int(np.median(wc)), "min", int(wc.min()), "max", int(wc.max()), flush=True)
for kind in ("first", "sum", "mean"):
    for pct in (75, 99, "freq"):
        if pct == "freq":
            pw = freq_prior; pn = torch.quantile(S0.float(), 0.75).repeat(len(np_))
        else:
            pr = torch.quantile(S0.float(), pct / 100); pw = pr.repeat(len(wp)); pn = pr.repeat(len(np_))
        m.table = torch.cat([T0, codes(kind, np_), codes(kind, wp)]); m.seal = torch.cat([S0, pn, pw])
        ids_w = list(range(V + len(np_), V + len(np_) + len(wp)))
        merged = [merge(s, ids_w) for s in seqs]
        n_merged = sum(len(s) for s in seqs) - sum(len(s) for s in merged)
        t = total_nats(merged)
        print(f"code={kind:5s} prior {pct:6}: {t / chars:.4f} nats/char ({(t - base) / chars:+.4f} vs none; {n_merged} piece-tokens merged)", flush=True)