"""Brand-new control tokens, usable from context with ZERO training (the paper's adding-tokens claim).

Five families of tokens that exist in no vocabulary: tool calls, thinking markers, SQL blocks, citations, JSON fields.
Each family has 12 hand-written examples. For every rotation r (24) and shot count k (1, 3, 5), the prompt shows k
examples and a new query; we measure whether the model puts the OPEN token at the right place (top-1, and its
probability) and the CLOSE token right after the content (top-1). B/C: new row = random +-1 code + prior at the 75th
percentile of the seal (a frozen table's normal way to take a row). A: new row drawn at the mean norm of its learned
rows (a learned table has no construction rule; this is the fair untrained equivalent). 95% bootstrap over trials.

  python control_tokens.py --arm B --ckpt hf/B/step_07820.pt [--tag B]
"""
import argparse, json, pathlib
import numpy as np, torch, torch.nn.functional as F
import fastc

FAMILIES = {
 "tool_call": (["<tool_call>", "</tool_call>", "<tool_result>"], "User: {q}\nAssistant:", " {a}", " {r}\n{f}\n", [
  ("What is 17 times 23?", "calc(17*23)", "391", "17 times 23 is 391."),
  ("What is the weather in Paris?", "weather(\"Paris\")", "18C, cloudy", "It is 18C and cloudy in Paris."),
  ("How far is Sydney from Melbourne?", "distance(\"Sydney\",\"Melbourne\")", "713 km", "About 713 km."),
  ("What is 144 divided by 12?", "calc(144/12)", "12", "144 divided by 12 is 12."),
  ("What time is it in Tokyo?", "time(\"Tokyo\")", "14:05", "It is 14:05 in Tokyo."),
  ("What is the square root of 81?", "calc(sqrt(81))", "9", "The square root of 81 is 9."),
  ("Who wrote Hamlet?", "search(\"author of Hamlet\")", "William Shakespeare", "Hamlet was written by William Shakespeare."),
  ("Convert 5 miles to km.", "convert(5, \"mi\", \"km\")", "8.05 km", "5 miles is about 8.05 km."),
  ("What is the capital of Peru?", "search(\"capital of Peru\")", "Lima", "The capital of Peru is Lima."),
  ("What is 2 to the power 10?", "calc(2**10)", "1024", "2 to the power 10 is 1024."),
  ("Is 97 a prime number?", "is_prime(97)", "true", "Yes, 97 is prime."),
  ("What day is 2026-03-01?", "weekday(\"2026-03-01\")", "Sunday", "2026-03-01 is a Sunday.")]),
 "think": (["<think>", "</think>"], "Question: {q}\nAnswer:", " {a}", " {f}\n", [
  ("If a train leaves at 3pm and arrives at 5pm, how long is the trip?", "5pm minus 3pm is two hours.", "", "Two hours."),
  ("Tom has 3 apples and buys 4 more. How many?", "3 plus 4 is 7.", "", "7 apples."),
  ("Which is heavier, a kilo of iron or a kilo of feathers?", "Both are one kilo.", "", "They weigh the same."),
  ("What is half of 50?", "50 divided by 2 is 25.", "", "25."),
  ("A week has how many hours?", "7 days times 24 hours is 168.", "", "168 hours."),
  ("Is 15 odd or even?", "15 ends in 5, so it is odd.", "", "Odd."),
  ("Sara is older than Ben. Ben is older than Li. Who is youngest?", "Li is younger than Ben, who is younger than Sara.", "", "Li."),
  ("How many legs do three cats have?", "3 cats times 4 legs is 12.", "", "12 legs."),
  ("What comes after Thursday?", "The day after Thursday is Friday.", "", "Friday."),
  ("If I have 10 dollars and spend 3, what is left?", "10 minus 3 is 7.", "", "7 dollars."),
  ("Double 21.", "21 times 2 is 42.", "", "42."),
  ("How many minutes in 2 hours?", "2 times 60 is 120.", "", "120 minutes.")]),
 "sql": (["<sql>", "</sql>", "<rows>"], "Request: {q}\nQuery:", " {a}", " {r}\n{f}\n", [
  ("How many users signed up today?", "SELECT COUNT(*) FROM users WHERE signup_date = CURRENT_DATE;", "42", "42 users signed up today."),
  ("List the oldest customer.", "SELECT name FROM customers ORDER BY birth_date LIMIT 1;", "Maria Lopez", "The oldest customer is Maria Lopez."),
  ("Total sales in March?", "SELECT SUM(amount) FROM sales WHERE month = 3;", "18250.50", "March sales were 18,250.50."),
  ("Which product sold most?", "SELECT product FROM sales GROUP BY product ORDER BY SUM(qty) DESC LIMIT 1;", "Widget A", "Widget A sold the most."),
  ("How many orders are pending?", "SELECT COUNT(*) FROM orders WHERE status = 'pending';", "7", "7 orders are pending."),
  ("Average order value?", "SELECT AVG(total) FROM orders;", "63.20", "The average order is 63.20."),
  ("Emails of users in Spain?", "SELECT email FROM users WHERE country = 'ES';", "ana@x.es, luis@y.es", "Two users: ana@x.es and luis@y.es."),
  ("Cheapest item in stock?", "SELECT name FROM items WHERE stock > 0 ORDER BY price LIMIT 1;", "Pencil", "The cheapest item in stock is Pencil."),
  ("How many distinct cities?", "SELECT COUNT(DISTINCT city) FROM customers;", "19", "19 distinct cities."),
  ("Latest invoice date?", "SELECT MAX(date) FROM invoices;", "2026-10-13", "The latest invoice is from 2026-10-13."),
  ("Staff in the sales team?", "SELECT COUNT(*) FROM staff WHERE team = 'sales';", "11", "11 staff are in sales."),
  ("Refunds over 100?", "SELECT COUNT(*) FROM refunds WHERE amount > 100;", "3", "3 refunds were over 100.")]),
 "cite": (["<cite>", "</cite>"], "Claim: {q}\nSource:", " {a}", "\n{f}\n", [
  ("Water boils at 100 C at sea level.", "Physics Handbook, ch. 4", "", "Confirmed."),
  ("The Nile is about 6,650 km long.", "World Atlas 2024, p. 88", "", "Confirmed."),
  ("Light travels at about 300,000 km/s.", "Smith, Optics (2019), p. 12", "", "Confirmed."),
  ("Canberra is the capital of Australia.", "Australian Government Directory", "", "Confirmed."),
  ("The human heart has four chambers.", "Gray's Anatomy, 41st ed.", "", "Confirmed."),
  ("Mount Everest is 8,849 m tall.", "Nepal Survey Department, 2020", "", "Confirmed."),
  ("Honey never spoils.", "Food Science Review, vol. 3", "", "Confirmed."),
  ("Pluto was reclassified in 2006.", "IAU Resolution B5 (2006)", "", "Confirmed."),
  ("A marathon is 42.195 km.", "IAAF Competition Rules", "", "Confirmed."),
  ("The Great Wall is over 20,000 km long.", "China Heritage Survey, 2012", "", "Confirmed."),
  ("Sound travels faster in water than in air.", "Acoustics Primer, ch. 2", "", "Confirmed."),
  ("Venus is hotter than Mercury.", "NASA Planetary Fact Sheet", "", "Confirmed.")]),
 "field": (["<field>", "</field>", "<value>"], "Record: {q}\nExtract:", " {a}", " {r}\n{f}\n", [
  ("Name: Ada Lovelace; Born: 1815; City: London", "born", "1815", "Done."),
  ("Order 551: 3 x lamp, total 90.00, ship to Perth", "total", "90.00", "Done."),
  ("Patient 77, age 54, blood type O+", "blood_type", "O+", "Done."),
  ("Flight QF1, gate 12, departs 21:40", "gate", "12", "Done."),
  ("Book: Dune, author Frank Herbert, 1965", "author", "Frank Herbert", "Done."),
  ("Car: Toyota Corolla, 2018, 64,000 km", "year", "2018", "Done."),
  ("Invoice 9981, due 2026-11-01, amount 420", "due", "2026-11-01", "Done."),
  ("Employee: Li Wei, dept: Finance, ext 3321", "dept", "Finance", "Done."),
  ("Package 4A, weight 2.3 kg, fragile: yes", "weight", "2.3 kg", "Done."),
  ("Room 204, floor 2, capacity 12", "capacity", "12", "Done."),
  ("Song: Imagine, artist John Lennon, 1971", "artist", "John Lennon", "Done."),
  ("Plant: basil, water twice a week, sun: full", "water", "twice a week", "Done.")]),
}
ROT, SHOTS = 24, (1, 3, 5)


@torch.no_grad()
def next_lp(m, ids):
    x = torch.tensor([ids], 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 add_rows(m, arm, n, seed):
    if arm == "A":
        W = m.body.embed_tokens.weight; V = W.shape[0]
        g = torch.Generator(device="cpu").manual_seed(seed)
        r = torch.randn(n, W.shape[1], generator=g).cuda(); r = r / r.norm(dim=1, keepdim=True) * W.norm(dim=1).mean()
        m.body.embed_tokens.weight = torch.nn.Parameter(torch.cat([W.data, r.to(W.dtype)]), requires_grad=False)
        m.body.embed_tokens.num_embeddings = V + n
    else:
        V = m.table.shape[0]; g = torch.Generator(device="cpu").manual_seed(seed)
        codes = (torch.randint(0, 2, (n, m.table.shape[1]), generator=g) * 2 - 1).float().cuda()
        prior = torch.quantile(m.seal.detach().float(), 0.75)
        m.table = torch.cat([m.table, codes]); seal = torch.cat([m.seal.detach(), prior.repeat(n)])
        if isinstance(m.seal, torch.nn.Parameter):
            del m.seal
        m.seal = seal
    return list(range(V, V + n))


def restore(m, arm, V0):
    if arm == "A":
        W = m.body.embed_tokens.weight
        m.body.embed_tokens.weight = torch.nn.Parameter(W.data[:V0].clone(), requires_grad=False); m.body.embed_tokens.num_embeddings = V0
    else:
        m.table = m.table[:V0]; m.seal = m.seal[:V0]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", required=True); ap.add_argument("--ckpt", required=True); ap.add_argument("--tag", default="")
    ap.add_argument("--data", default="data")
    a = ap.parse_args()
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained("google/gemma-3-270m"); sep = tok.eos_token_id
    enc = lambda s: tok(s, add_special_tokens=False).input_ids
    counts = np.load(pathlib.Path(a.data) / "unigram_counts.npy")
    m, tag = fastc.load(a.arm, a.ckpt, counts); tag = a.tag or tag
    V0 = m.table.shape[0] if a.arm != "A" else m.body.embed_tokens.weight.shape[0]
    R = {"tag": tag, "families": {}}
    rng = np.random.default_rng(0)
    for fi, (fam, (names, head, body, tail, ex)) in enumerate(FAMILIES.items()):
        ids_new = add_rows(m, a.arm, len(names), seed=100 + fi); T = dict(zip(names, ids_new))
        op, cl, res = names[0], names[1], (names[2] if len(names) > 2 else None)
        def example(q, ans, r, f):
            s = enc(head.format(q=q)) + [T[op]] + enc(body.format(a=ans)) + [T[cl]]
            if res: s += [T[res]]
            return s + enc(tail.format(r=r, f=f))
        R["families"][fam] = {}
        for k in SHOTS:
            open_hit, open_p, close_hit = [], [], []
            for r in range(ROT):
                idx = [(r + j) % len(ex) for j in range(k)]; qi = (r + k) % len(ex)
                ids = [sep]
                for j in idx:
                    ids += example(*ex[j])
                q, ans, rr, f = ex[qi]
                ids += enc(head.format(q=q))
                lp = next_lp(m, ids)
                open_hit.append(int(lp.argmax()) == T[op]); open_p.append(float(lp[T[op]].exp()))
                lp2 = next_lp(m, ids + [T[op]] + enc(body.format(a=ans)))
                close_hit.append(int(lp2.argmax()) == T[cl])
            oh, ch = np.array(open_hit, float), np.array(close_hit, float)
            b = rng.integers(0, ROT, (5000, ROT))
            R["families"][fam][k] = {"open_top1": float(oh.mean()), "open_ci": [float(np.percentile(oh[b].mean(1), 2.5)), float(np.percentile(oh[b].mean(1), 97.5))],
                                     "open_prob": float(np.mean(open_p)), "close_top1": float(ch.mean()), "n": ROT}
            print(f"{tag} {fam:9s} {k}-shot: open {oh.mean()*100:5.1f}% (p={np.mean(open_p):.2f}) close {ch.mean()*100:5.1f}%", flush=True)
        restore(m, a.arm, V0)
    allo = [v["open_top1"] for f in R["families"].values() for v in f.values()]
    allc = [v["close_top1"] for f in R["families"].values() for v in f.values()]
    R["overall"] = {"open_top1": float(np.mean(allo)), "close_top1": float(np.mean(allc))}
    print(f"{tag} OVERALL open {np.mean(allo)*100:.1f}% close {np.mean(allc)*100:.1f}%", flush=True)
    pathlib.Path("reports").mkdir(exist_ok=True)
    json.dump(R, open(f"reports/control_tokens_{tag}.json", "w"), indent=1)


if __name__ == "__main__":
    main()