"""Stage 1: stream corpus -> BGE embeddings (sharded memmap) -> faiss k-means K clusters.

    python scripts/embed_cluster.py --clusters 10000 --n-docs 4000000
Scales to K=1e6 by bumping --clusters (and --n-docs to keep >=~40*K docs).
"""
import argparse
import json
import os

import numpy as np
import torch
from datasets import load_dataset

from mxf.config import ClusterConfig


def embed_stream(cfg, model, tok, device):
    """Yield (texts, np.float32[n,d]) batches from the streamed corpus."""
    ds = load_dataset(cfg.corpus, split="en", streaming=True)  # Ultra-FineWeb splits are en/zh (no "train"); field is "content"
    buf, seen = [], 0
    for row in ds:
        t = (row.get("content") or row.get("text") or "").strip()
        if not (cfg.min_chars <= len(t)):
            continue
        buf.append(t[: cfg.max_chars])
        if len(buf) == cfg.embed_batch:
            yield buf, _encode(buf, model, tok, device)
            seen += len(buf); buf = []
            if seen >= cfg.n_docs:
                return
    if buf:
        yield buf, _encode(buf, model, tok, device)


@torch.no_grad()
def _encode(texts, model, tok, device):
    enc = tok(texts, padding=True, truncation=True, max_length=256, return_tensors="pt").to(device)
    out = model(**enc).last_hidden_state[:, 0]         # BGE: CLS pooling
    return torch.nn.functional.normalize(out, dim=-1).float().cpu().numpy()


def main():
    from transformers import AutoModel, AutoTokenizer

    cfg = ClusterConfig()
    ap = argparse.ArgumentParser()
    for f in ("corpus", "embed_model", "out_dir"):
        ap.add_argument(f"--{f.replace('_','-')}", default=getattr(cfg, f))
    ap.add_argument("--n-docs", type=int, default=cfg.n_docs)
    ap.add_argument("--clusters", type=int, default=cfg.clusters)
    a = ap.parse_args()
    cfg.corpus, cfg.embed_model, cfg.out_dir = a.corpus, a.embed_model, a.out_dir
    cfg.n_docs, cfg.clusters = a.n_docs, a.clusters
    os.makedirs(cfg.out_dir, exist_ok=True)
    device = "cuda"

    tok = AutoTokenizer.from_pretrained(cfg.embed_model)
    model = AutoModel.from_pretrained(cfg.embed_model, torch_dtype=torch.bfloat16).to(device).eval()
    d = model.config.hidden_size

    # sharded memmap of embeddings + a jsonl of the source texts (for centroid targets later)
    emb_path = os.path.join(cfg.out_dir, "emb.f32")
    emb = np.memmap(emb_path, dtype=np.float32, mode="w+", shape=(cfg.n_docs, d))
    txt = open(os.path.join(cfg.out_dir, "texts.jsonl"), "w")
    n = 0
    for texts, vecs in embed_stream(cfg, model, tok, device):
        k = min(len(texts), cfg.n_docs - n)
        emb[n : n + k] = vecs[:k]
        for t in texts[:k]:
            txt.write(json.dumps({"t": t}) + "\n")
        n += k
        if n // 100_000 != (n - k) // 100_000:  # batch-sized steps almost never land on exact multiples
            print(f"embedded {n}/{cfg.n_docs}", flush=True)
        if n >= cfg.n_docs:
            break
    txt.close(); emb.flush()
    print(f"embedded {n} docs, d={d}", flush=True)

    # faiss GPU k-means (falls back to CPU/MiniBatchKMeans if faiss-gpu absent)
    X = np.asarray(emb[:n])
    try:
        import faiss

        faiss.omp_set_num_threads(min(os.cpu_count() or 32, 128))
        # gpu=False: faiss-gpu wheels have no Blackwell sm_103 kernels ("no kernel image"); CPU
        # faiss is multithreaded and fine for k-means. (GPU k-means: use recluster_torch.py.)
        km = faiss.Kmeans(d, cfg.clusters, niter=20, verbose=True, gpu=False, seed=cfg.seed)
        km.train(X)
        _, assign = km.index.search(X, 1)
        assign = assign.ravel()
        centroids = km.centroids.reshape(cfg.clusters, d)
    except Exception as e:
        print("faiss unavailable, MiniBatchKMeans:", e, flush=True)
        from sklearn.cluster import MiniBatchKMeans

        km = MiniBatchKMeans(cfg.clusters, batch_size=10000, n_init=3, random_state=cfg.seed).fit(X)
        assign, centroids = km.labels_, km.cluster_centers_

    np.save(os.path.join(cfg.out_dir, "assign.npy"), assign.astype(np.int32))
    np.save(os.path.join(cfg.out_dir, "centroids.npy"), centroids.astype(np.float32))
    json.dump({"n_docs": int(n), "d": int(d), "clusters": cfg.clusters},
              open(os.path.join(cfg.out_dir, "meta.json"), "w"))
    print(f"CLUSTERED {n} docs into {cfg.clusters} clusters -> {cfg.out_dir}", flush=True)


if __name__ == "__main__":
    main()