"""Stage 1 (activation version): cluster the corpus in Qwen3's OWN space — mean-pooled layer-27
residuals — not a separate embedding model. One pass produces both the clustering features AND the
per-doc probe residuals (build_data reads emb.f32 rows straight from here), and clusters live in
the same space as the probe/injection/reward.

OUT-OF-CORE at 200M docs / k=1e6: each rank streams its corpus shard and writes residuals straight
into its contiguous row-block of a shared disk memmap (emb.f32 — 3.3TB at 200M, never in RAM);
texts stream to per-rank jsonl (rank 0 concatenates in rank order, so rows stay aligned across
emb.f32 / assign.npy / texts.jsonl). K-means holds only the [k, d] centroids on GPU (16GB at k=1M)
and streams point-chunks from the memmap for assign+accumulate, data-parallel across ranks;
centroids train on a --kmeans-sample random subset, then ALL docs get assigned by streaming.

    python scripts/embed_cluster_acts.py --n-docs 50000 --clusters 300 --inspect   # coherence probe
    torchrun --standalone --nproc_per_node=8 scripts/embed_cluster_acts.py \
        --n-docs 200000000 --clusters 1000000 --kmeans-sample 20000000
"""
import argparse
import json
import os
import shutil

import numpy as np
import torch
import torch.distributed as dist
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

from mxf.config import CORPUS, D_MODEL, MODEL, READ_LAYER
from mxf.inject import read_resid

DIST_BYTES = 15e9   # [chunk, k] fp32 distance buffer budget → chunk*k*4 < 15GB (spec'd for k=1e6)
X_BYTES = 4e9       # cap on the streamed point-chunk itself (binds when k is small)


def _argmin_chunk(X, Cb, Cn):
    """Nearest centroid per row. bf16 matmul (fp32 tensor-core accum) at ~2x fp32 speed; the fp32
    ||c||² correction keeps the argmin stable for cluster assignment."""
    return (Cn[None] - 2 * (X.to(torch.bfloat16) @ Cb.T).float()).argmin(1)


def _chunk_rows(k, d):
    return max(256, min(int(DIST_BYTES / (k * 4)), int(X_BYTES / (d * 4))))


def stream_kmeans(emb, k, iters, seed, sample, device, rank, world):
    """Lloyd's with only the centroids resident on GPU; points stream from the disk memmap.
    Centroids train on a random `sample`-doc subsample (0 = all docs), sharded across ranks
    (all_reduce of sums/counts). Every rank runs the same rng → identical init and updates, no
    broadcasts. Empty clusters keep their previous centroid. Returns [k, d] fp32 on GPU."""
    n, d = emb.shape
    chunk = _chunk_rows(k, d)
    rng = np.random.default_rng(seed)
    samp = np.sort(rng.choice(n, min(sample or n, n), replace=False))
    C = torch.from_numpy(np.asarray(emb[np.sort(rng.choice(samp, k, replace=False))])).to(device)
    mine = samp[rank::world]
    ones = torch.ones(chunk, device=device)
    for it in range(iters):
        Cb, Cn = C.to(torch.bfloat16), (C * C).sum(1)
        Csum = torch.zeros_like(C)
        cnt = torch.zeros(k, device=device)
        for s in range(0, len(mine), chunk):
            X = torch.from_numpy(np.asarray(emb[mine[s : s + chunk]])).to(device)
            a = _argmin_chunk(X, Cb, Cn)
            Csum.index_add_(0, a, X)
            cnt.index_add_(0, a, ones[: len(X)])
        if world > 1:
            dist.all_reduce(Csum); dist.all_reduce(cnt)
        live = cnt > 0
        C = torch.where(live[:, None], Csum / cnt.clamp(min=1)[:, None], C)
        if rank == 0:
            print(f"  kmeans iter {it}: {int(live.sum())}/{k} live", flush=True)
    return C


def stream_assign(emb, C, out, lo, hi, rank):
    """Assign rows [lo, hi) to their nearest centroid, streaming memmap→GPU→memmap."""
    chunk = _chunk_rows(*C.shape)
    Cb, Cn = C.to(torch.bfloat16), (C * C).sum(1)
    for s in range(lo, hi, chunk):
        X = torch.from_numpy(np.asarray(emb[s : min(s + chunk, hi)])).to(C.device)
        out[s : s + len(X)] = _argmin_chunk(X, Cb, Cn).to(torch.int32).cpu().numpy()
        if rank == 0 and (s - lo) % (50 * chunk) == 0:
            print(f"  assign {s - lo}/{hi - lo}", flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--corpus", default=CORPUS)
    ap.add_argument("--n-docs", type=int, default=50000)
    ap.add_argument("--clusters", type=int, default=300)
    ap.add_argument("--batch", type=int, default=256)
    ap.add_argument("--min-chars", type=int, default=200)
    ap.add_argument("--max-tok", type=int, default=64)
    ap.add_argument("--iters", type=int, default=20)
    ap.add_argument("--kmeans-sample", type=int, default=0,
                    help="train centroids on N random docs (0 = all); ALL docs still get assigned")
    ap.add_argument("--out-dir", default="data/actclusters")
    ap.add_argument("--inspect", action="store_true", help="print sample docs/cluster (coherence check)")
    ap.add_argument("--seed", type=int, default=0)
    a = ap.parse_args()
    rank = int(os.environ.get("RANK", 0)); world = int(os.environ.get("WORLD_SIZE", 1))
    local = int(os.environ.get("LOCAL_RANK", 0))
    if world > 1:
        dist.init_process_group("nccl"); torch.cuda.set_device(local)
    device = f"cuda:{local}"
    os.makedirs(a.out_dir, exist_ok=True)
    per_rank = a.n_docs // world; n_total = per_rank * world; lo = rank * per_rank

    if rank == 0:  # sparse files, instant to create; ranks write disjoint contiguous row-blocks
        np.memmap(f"{a.out_dir}/emb.f32", dtype=np.float32, mode="w+", shape=(n_total, D_MODEL)).flush()
        np.lib.format.open_memmap(f"{a.out_dir}/assign.npy", mode="w+", dtype=np.int32, shape=(n_total,)).flush()
    if world > 1:
        dist.barrier()
    emb = np.memmap(f"{a.out_dir}/emb.f32", dtype=np.float32, mode="r+", shape=(n_total, D_MODEL))
    assign = np.lib.format.open_memmap(f"{a.out_dir}/assign.npy", mode="r+")

    tok = AutoTokenizer.from_pretrained(MODEL)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    tok.padding_side = "right"
    model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16,
                                                 attn_implementation="sdpa", device_map={"": device}).eval()

    # each rank streams its own file-shard of the corpus; residuals go straight into its emb.f32
    # block as batches complete (never accumulated in RAM), texts to a per-rank jsonl
    ds = load_dataset(a.corpus, split="en", streaming=True)
    if world > 1:
        ds = ds.shard(num_shards=world, index=rank)   # file-level: each rank reads 1/world of shards
    tf = open(f"{a.out_dir}/texts_rank{rank}.jsonl", "w")
    buf, done = [], 0

    @torch.no_grad()
    def flush(batch):
        nonlocal done
        enc = tok(batch, padding=True, truncation=True, max_length=a.max_tok,
                  return_tensors="pt", add_special_tokens=True).to(device)
        emb[lo + done : lo + done + len(batch)] = read_resid(model, READ_LAYER, dict(enc), pool="mean").cpu().numpy()
        for t in batch:
            tf.write(json.dumps({"t": t}) + "\n")
        done += len(batch)

    for row in ds:
        t = (row.get("content") or row.get("text") or "").strip()
        if len(t) < a.min_chars:
            continue
        buf.append(t[:2000])
        if len(buf) == a.batch or done + len(buf) == per_rank:
            flush(buf); buf = []
            if rank == 0 and done % 25600 == 0:
                print(f"resid {done}/{per_rank}", flush=True)
            if done >= per_rank:
                break
    assert done == per_rank, f"corpus exhausted: rank {rank} got {done}/{per_rank} docs"
    tf.close(); emb.flush()
    del model
    torch.cuda.empty_cache()
    if world > 1:
        dist.barrier()

    if rank == 0:
        print(f"clustering {n_total} docs (d={D_MODEL}) into {a.clusters} "
              f"(centroid sample {a.kmeans_sample or n_total})", flush=True)
    C = stream_kmeans(emb, a.clusters, a.iters, a.seed, a.kmeans_sample, device, rank, world)
    stream_assign(emb, C, assign, lo, lo + per_rank, rank)
    assign.flush()
    if world > 1:
        dist.barrier()

    if rank == 0:
        np.save(f"{a.out_dir}/centroids.npy", C.cpu().numpy().astype(np.float32))
        with open(f"{a.out_dir}/texts.jsonl", "wb") as out:  # rank order == emb/assign row order
            for r in range(world):
                with open(f"{a.out_dir}/texts_rank{r}.jsonl", "rb") as f:
                    shutil.copyfileobj(f, out)
                os.remove(f"{a.out_dir}/texts_rank{r}.jsonl")
        json.dump({"n_docs": n_total, "d": D_MODEL, "clusters": a.clusters, "space": "qwen3_l27_mean"},
                  open(f"{a.out_dir}/meta.json", "w"))
        print(f"CLUSTERED {n_total} -> {a.clusters} (qwen3 layer-{READ_LAYER} mean-pool)", flush=True)
        if a.inspect:
            A = np.asarray(assign)
            rng = np.random.default_rng(1)
            sizes = np.bincount(A, minlength=a.clusters)
            big = np.where(sizes >= 4)[0]
            picks = {int(c): np.where(A == c)[0][:4] for c in rng.choice(big, min(8, len(big)), replace=False)}
            want = {int(r) for rows in picks.values() for r in rows}
            txt = {i: json.loads(l)["t"] for i, l in enumerate(open(f"{a.out_dir}/texts.jsonl")) if i in want}
            for c, rows in picks.items():
                print(f"\n=== cluster {c} ({sizes[c]} docs) ===", flush=True)
                for r in rows:
                    print("  -", txt[int(r)][:100].replace("\n", " "), flush=True)
    if world > 1:
        dist.barrier(); dist.destroy_process_group()


if __name__ == "__main__":
    main()