"""Adding words after training, without retraining (the main test).

A frozen-table model (B, C, B2) can take a NEW vocabulary row at any time: the new word gets a frozen 512-bit code made by
the same transform as the table (Gemma 3 270M embedding space -> centre -> the saved 512 principal directions -> sign),
here from the mean of the Gemma embeddings of the pieces the word is currently split into, and a frequency prior at the
75th percentile of the existing rows (the SHADOW-50M v1.1 rule). Nothing is trained. A learned table (A) has no way to
add a row without training, so A keeps reading the word piece by piece.

Two measurements, each for: A pieces | arm pieces | arm with the new token
  1 in-context recall: "My name is <w>. ... What is my name? My name is" -> is <w> produced? (greedy; pieces must all
    be right; the new token must be the single next token)
  2 natural text: words that Gemma splits into 2+ pieces, at their real occurrences in held-out FineWeb-Edu text:
    nats per word occurrence (sum of piece losses) vs nats for the single new token in the same context

  python new_words.py --arm B --ckpt finals/B/step_07820.pt [--tag B] [--ref-ckpt finals/A/step_07820.pt]
"""
import argparse, json, pathlib, random, re, time
import numpy as np
import torch
import torch.nn.functional as F

SEQ = 1024


def gemma_embedding():
    from safetensors import safe_open
    from huggingface_hub import hf_hub_download
    with safe_open(hf_hub_download("google/gemma-3-270m", "model.safetensors"), "pt") as f:
        return f.get_tensor([k for k in f.keys() if k.endswith("embed_tokens.weight")][0]).float()[:262144]


def load(arm, ckpt, counts):
    import train
    m = train.Model(arm, counts).cuda().eval()
    st = torch.load(ckpt, map_location="cuda", weights_only=False)
    m.load_state_dict(st["model"])
    return m, st.get("tag", arm)


def make_names(tok, n, seed=0):
    """made-up names that the tokenizer splits into 2-4 pieces (fresh: never in any training text)"""
    rng = random.Random(seed); syl = ["ka", "zor", "mel", "vit", "ran", "dul", "pex", "sho", "lin", "qua", "ber", "tis",
                                     "mon", "gra", "fel", "yun", "dra", "vok", "ser", "nim", "tal", "rux", "jeb", "wen"]
    out = set()
    while len(out) < n:
        w = "".join(rng.choice(syl) for _ in range(rng.randint(2, 3))).capitalize()
        k = len(tok(" " + w, add_special_tokens=False).input_ids)
        if 2 <= k <= 4:
            out.add(w)
    return sorted(out)


def corpus_words(tok, val, n_words, min_count):
    """real English words that the tokenizer splits into 2+ pieces, by frequency in held-out text"""
    text = tok.decode(val[:3_000_000].tolist())
    cnt = {}
    for w in re.findall(r"(?<= )[a-z]{6,}", text):
        cnt[w] = cnt.get(w, 0) + 1
    cands = [w for w, c in sorted(cnt.items(), key=lambda x: -x[1]) if c >= min_count]
    out = []
    for w in cands:
        if len(tok(" " + w, add_special_tokens=False).input_ids) >= 2:
            out.append(w)
        if len(out) >= n_words:
            break
    return out


@torch.no_grad()
def add_rows(m, E, piece_lists, prior_pct=75):
    """append one frozen code + prior per new word; returns the new ids"""
    V = m.table.shape[0]
    vecs = torch.stack([E[p].mean(0) for p in piece_lists]).cuda()
    codes = torch.sign((vecs - m.table_mu) @ m.table_W); codes[codes == 0] = 1
    prior = torch.quantile(m.seal.detach().float(), prior_pct / 100)
    m.table = torch.cat([m.table, codes])
    seal = torch.cat([m.seal.detach(), prior.repeat(len(piece_lists))])
    if isinstance(m.seal, torch.nn.Parameter):
        del m.seal
    m.seal = seal
    return list(range(V, V + len(piece_lists)))


@torch.no_grad()
def logp_next(m, ids):
    x = torch.tensor([ids[-SEQ:]], device="cuda")
    with torch.autocast("cuda", dtype=torch.bfloat16):
        return F.log_softmax(m.logits(m.hidden(x))[0, -1].float(), -1)


@torch.no_grad()
def greedy(m, ids, k):
    out = []
    for _ in range(k):
        nxt = int(logp_next(m, ids + out).argmax()); out.append(nxt)
    return out


RECALL = ["My name is {w}. I live in a small town near the sea. What is my name? My name is",
          "The new ship was called the {w}. Everyone in the harbour admired the",
          "Our team leader is {w}. Tomorrow the report will be presented by",
          "The rare mineral {w} was found in the cave. Scientists are now studying"]


def recall_test(m, tok, sep, words, new_ids=None):
    hits = 0; tot = 0
    for i, w in enumerate(words):
        pieces = tok(" " + w, add_special_tokens=False).input_ids
        for t in RECALL:
            pre, post = t.split("{w}")
            a = tok(pre.rstrip(), add_special_tokens=False).input_ids if pre.strip() else []
            b = tok(post, add_special_tokens=False).input_ids
            word = [new_ids[i]] if new_ids else pieces
            ids = [sep] + a + word + b
            got = greedy(m, ids, len(word))
            hits += int(got == word); tot += 1
    return hits / tot


def text_test(m, val, words_pieces, new_ids=None, max_occ=40, ctx=256):
    """nats per word occurrence at real occurrences in held-out text"""
    arr = np.asarray(val[:20_000_000]); res = []
    for i, p in enumerate(words_pieces):
        p = np.array(p); k = len(p)
        hits = np.where(np.all(np.stack([arr[j:len(arr) - k + j + 1] == p[j] for j in range(k)]), 0))[0]
        hits = [h for h in hits if h >= ctx][:max_occ]
        for h in hits:
            context = arr[h - ctx:h].tolist()
            if new_ids is None:
                lp = 0.0; seq = context[:]
                for t in p.tolist():
                    lp += float(logp_next(m, seq)[t]); seq.append(t)
            else:
                lp = float(logp_next(m, context)[new_ids[i]])
            res.append(-lp)
    return float(np.mean(res)), len(res)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", required=True, choices=["B", "C"]); ap.add_argument("--ckpt", required=True)
    ap.add_argument("--tag", default=""); ap.add_argument("--ref-ckpt", default="", help="arm A checkpoint (pieces only)")
    ap.add_argument("--data", default="data"); ap.add_argument("--names", type=int, default=60)
    ap.add_argument("--words", type=int, default=100)
    a = ap.parse_args()
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained("google/gemma-3-270m"); sep = tok.eos_token_id
    counts = np.load(pathlib.Path(a.data) / "unigram_counts.npy")
    val = np.load(pathlib.Path(a.data) / "val.npy", mmap_mode="r")
    names = make_names(tok, a.names)
    words = corpus_words(tok, val, a.words, min_count=5)
    wp = [tok(" " + w, add_special_tokens=False).input_ids for w in words]
    np_ = [tok(" " + w, add_special_tokens=False).input_ids for w in names]
    R = {"names": names, "words": words, "date": time.strftime("%F %H:%M")}
    if a.ref_ckpt:
        A, _ = load("A", a.ref_ckpt, counts)
        R["A_pieces"] = {"recall": recall_test(A, tok, sep, names), "text_nats_per_word": text_test(A, val, wp)}
        print("A pieces:", R["A_pieces"], flush=True); del A; torch.cuda.empty_cache()
    m, tag = load(a.arm, a.ckpt, counts); tag = a.tag or tag
    R[f"{tag}_pieces"] = {"recall": recall_test(m, tok, sep, names), "text_nats_per_word": text_test(m, val, wp)}
    print(f"{tag} pieces:", R[f"{tag}_pieces"], flush=True)
    E = gemma_embedding()
    ids_n = add_rows(m, E, np_); ids_w = add_rows(m, E, wp)
    R[f"{tag}_new_tokens"] = {"recall": recall_test(m, tok, sep, names, ids_n),
                              "text_nats_per_word": text_test(m, val, wp, ids_w), "rows_added": len(ids_n) + len(ids_w)}
    print(f"{tag} new tokens:", R[f"{tag}_new_tokens"], flush=True)
    pathlib.Path("reports").mkdir(exist_ok=True)
    json.dump(R, open(f"reports/new_words_{tag}.json", "w"), indent=1)


if __name__ == "__main__":
    main()