"""Validation loss split by how often the target word appeared in TRAINING.

Question: does the frozen table win on rare words, where a learned table is starved of examples, even if the average
loss favours the learned table (common words dominate the average)?

  python freq_eval.py --arm A --ckpt A/step_07820.pt [--tag A] [--seqs 512]
  python freq_eval.py --compare reports/freq_A.json reports/freq_B.json reports/freq_C.json reports/freq_B2.json

Buckets by the target token's count in the 2.05B training tokens (data/unigram_counts.npy):
  never seen (0) | 1-9 | 10-99 | 100-999 | 1k-9.9k | 10k-99k | 100k-999k | 1M+
Saved per validation sequence (loss sum and count per bucket), so --compare gives paired 95% ranges (bootstrap over
sequences, 10,000 re-draws) for every arm minus the first file.
"""
import argparse, json, pathlib, time
import numpy as np
import torch
import torch.nn.functional as F

SEQ = 1024
EDGES = [0, 1, 10, 100, 1_000, 10_000, 100_000, 1_000_000, np.inf]
NAMES = ["never seen", "1-9", "10-99", "100-999", "1k-9.9k", "10k-99k", "100k-999k", "1M+"]


def bucket_of(counts):
    b = np.searchsorted(np.array(EDGES[1:-1]), counts, side="right")       # 0 -> bucket 0 (count 0), 1..9 -> 1, ...
    b[counts == 0] = 0
    b[(counts > 0) & (b == 0)] = 1
    return b


@torch.no_grad()
def run(a):
    import train
    counts = np.load(pathlib.Path(a.data) / "unigram_counts.npy")
    bk = torch.from_numpy(bucket_of(counts)).cuda()
    m = train.Model(a.arm, counts).cuda().eval()
    st = torch.load(a.ckpt, map_location="cuda", weights_only=False)
    m.load_state_dict(st["model"])
    val = np.load(pathlib.Path(a.data) / "val.npy", mmap_mode="r")
    x = torch.from_numpy(val[:a.seqs * (SEQ + 1)].astype(np.int64).reshape(a.seqs, SEQ + 1))
    nb = len(NAMES)
    sums = np.zeros((a.seqs, nb)); cnts = np.zeros((a.seqs, nb))
    for i in range(a.seqs):
        ids = x[i:i + 1].cuda()
        with torch.autocast("cuda", dtype=torch.bfloat16):
            lg = m.logits(m.hidden(ids[:, :-1])).float()
        tgt = ids[0, 1:]
        l = F.cross_entropy(lg[0], tgt, reduction="none")
        b = bk[tgt]
        sums[i] = torch.zeros(nb, device="cuda").scatter_add_(0, b, l).cpu().numpy()
        cnts[i] = torch.bincount(b, minlength=nb).float().cpu().numpy()
    tag = a.tag or st.get("tag", a.arm)
    out = {"tag": tag, "ckpt": a.ckpt, "buckets": NAMES, "sums": sums.tolist(), "counts": cnts.tolist(),
           "vocab_per_bucket": np.bincount(bucket_of(counts), minlength=nb).tolist(), "date": time.strftime("%F %H:%M")}
    pathlib.Path("reports").mkdir(exist_ok=True)
    p = pathlib.Path("reports") / f"freq_{tag}.json"; json.dump(out, open(p, "w"))
    print(f"{tag}:", " | ".join(f"{n} {sums[:, j].sum() / max(cnts[:, j].sum(), 1):.3f}" for j, n in enumerate(NAMES)))
    print("saved", p)


def compare(files):
    rng = np.random.default_rng(0)
    R = [json.load(open(f)) for f in files]
    ref = R[0]; S0 = np.array(ref["sums"]); C0 = np.array(ref["counts"])
    nseq = S0.shape[0]; idx = rng.integers(0, nseq, (10_000, nseq))
    L = ["# Validation loss by training frequency of the target word", "",
         f"{nseq} held-out sequences x 1,024 tokens. Loss in nats/token, lower is better. Differences vs {ref['tag']}, paired,"
         " 95% bootstrap range over sequences.", "",
         "| bucket (count in training) | vocab words | share of val tokens | " + " | ".join(r["tag"] for r in R) + " | "
         + " | ".join(f"{r['tag']} - {ref['tag']}" for r in R[1:]) + " |",
         "|---|---|---|" + "---|" * (len(R) + len(R) - 1)]
    tot = C0.sum()
    for j, name in enumerate(ref["buckets"]):
        if C0[:, j].sum() == 0:
            continue
        cells = []
        for r in R:
            S = np.array(r["sums"]); C = np.array(r["counts"])
            cells.append(f"{S[:, j].sum() / C[:, j].sum():.3f}")
        diffs = []
        for r in R[1:]:
            S = np.array(r["sums"]); C = np.array(r["counts"])
            d = S[:, j].sum() / C[:, j].sum() - S0[:, j].sum() / C0[:, j].sum()
            bs = S[idx, j].sum(1) / np.maximum(C[idx, j].sum(1), 1) - S0[idx, j].sum(1) / np.maximum(C0[idx, j].sum(1), 1)
            lo, hi = np.percentile(bs, [2.5, 97.5])
            verdict = "tied" if lo <= 0 <= hi else ("better" if d < 0 else "worse")
            diffs.append(f"{d:+.3f} ({lo:+.3f} to {hi:+.3f}) {verdict}")
        L.append(f"| {name} | {ref['vocab_per_bucket'][j]:,} | {C0[:, j].sum() / tot * 100:.2f}% | " + " | ".join(cells) + " | "
                 + " | ".join(diffs) + " |")
    open("reports/FREQ_SPLIT.md", "w").write("\n".join(L) + "\n")
    print("\n".join(L))


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", choices=["A", "B", "C", "D"])
    ap.add_argument("--ckpt"); ap.add_argument("--tag", default="")
    ap.add_argument("--data", default="data"); ap.add_argument("--seqs", type=int, default=512)
    ap.add_argument("--compare", nargs="+")
    a = ap.parse_args()
    compare(a.compare) if a.compare else run(a)