"""Cluster cached embeddings (emb.f32) on GPU with a minimal torch k-means — works on Blackwell
where faiss-gpu has no sm_103 kernels. Reuses an existing embed step; writes assign/centroids/meta.

    python scripts/recluster_torch.py --clusters-dir data/pilot_clusters --clusters 10000
"""
import argparse
import json

import numpy as np
import torch


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--clusters-dir", required=True)
    ap.add_argument("--clusters", type=int, required=True)
    ap.add_argument("--iters", type=int, default=20)
    ap.add_argument("--chunk", type=int, default=200_000)
    ap.add_argument("--seed", type=int, default=0)
    a = ap.parse_args()
    dev = "cuda"

    n = sum(1 for _ in open(f"{a.clusters_dir}/texts.jsonl"))
    d = np.memmap(f"{a.clusters_dir}/emb.f32", dtype=np.float32, mode="r").shape[0] // n
    emb = np.memmap(f"{a.clusters_dir}/emb.f32", dtype=np.float32, mode="r", shape=(n, d))
    print(f"kmeans: n={n} d={d} k={a.clusters}", flush=True)

    g = torch.Generator().manual_seed(a.seed)
    init = torch.from_numpy(np.asarray(emb[torch.randperm(n, generator=g)[: a.clusters].numpy()]))
    C = init.to(dev, torch.float32)                       # [k, d]
    Xg = torch.from_numpy(np.ascontiguousarray(emb)).to(dev, torch.float32)  # [n, d] (n*d*4 bytes)
    assign = torch.empty(n, dtype=torch.long, device=dev)
    for it in range(a.iters):
        Cn = (C * C).sum(1)                               # [k]
        for s in range(0, n, a.chunk):
            xb = Xg[s : s + a.chunk]
            d2 = Cn[None] - 2 * xb @ C.T                   # argmin over k (||x||² constant per row)
            assign[s : s + a.chunk] = d2.argmin(1)
        C.zero_()
        C.index_add_(0, assign, Xg)                       # sum members
        cnt = torch.bincount(assign, minlength=a.clusters).clamp(min=1)
        C /= cnt[:, None]
        empty = cnt == 1  # reseed empties to random points (bincount clamp hides true-empties)
        if empty.any():
            C[empty] = Xg[torch.randint(0, n, (int(empty.sum()),), device=dev)]
        print(f"  iter {it}", flush=True)

    np.save(f"{a.clusters_dir}/assign.npy", assign.cpu().numpy().astype(np.int32))
    np.save(f"{a.clusters_dir}/centroids.npy", C.cpu().numpy().astype(np.float32))
    json.dump({"n_docs": int(n), "d": int(d), "clusters": a.clusters},
              open(f"{a.clusters_dir}/meta.json", "w"))
    print(f"CLUSTERED {n} -> {a.clusters} clusters", flush=True)


if __name__ == "__main__":
    main()