docbert-plaid 0.9.0

//! K-means clustering over token embeddings.
//!
//! PLAID treats the entire set of document token embeddings as a cloud of
//! points in ℝᵈ and summarises it with `k` centroids. Every stored token is
//! then quantized relative to its nearest centroid, and queries use the
//! centroids as the first-stage filter. The quality of that summary caps
//! the quality of the whole index, so k-means lives at the base of the
//! stack.
//!
//! Points and centroids are both represented as flat row-major `&[f32]`
//! slices at the public boundary; internally the hot paths
//! ([`assign_points`], [`fit_with_init`]) materialise them into
//! `candle_core::Tensor`s and compute distances as matmul + broadcast.
//! With the `cuda` feature enabled those matmuls run on the GPU;
//! without it, candle's CPU `gemm` is still orders of magnitude faster
//! than the previous scalar nested loops.
//!
//! Training phase follows the PLAID / ColBERTv2 recipe: Lloyd's
//! algorithm on a random `k · MAX_POINTS_PER_CENTROID` subsample of the
//! corpus, seeded from the first `k` rows of that shuffle. Training
//! complexity is thus O(k² · MAX_POINTS_PER_CENTROID · dim · iters)
//! and independent of corpus size — a million-token corpus and a
//! billion-token corpus do the same amount of training work. The
//! full-corpus assign happens later in [`crate::codec`] during
//! encoding.
//!
//! [`nearest_centroid`] (single point) is the only scalar primitive
//! left — it's a pure-CPU helper used outside the training loop.
//! [`update_centroids`] (the M-step) and the assignment step both run
//! through candle, so the whole Lloyd loop dispatches to whichever
//! backend `p_tensor.device()` lives on.

use candle_core::{DType, Tensor};
use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom};

use crate::{Result, device::default_device, distance::squared_l2};

/// Upper bound on training samples per centroid.
///
/// ColBERTv2's paper uses `256 · k` as its training set and notes no
/// measurable recall difference past that point — more samples only
/// add compute, not cluster quality. fast-plaid's reference Python
/// implementation uses the same constant. For a 6.78M-token docbert
/// corpus at k=256 this caps training at 65,536 points, keeping the
/// Lloyd loop's per-iter cost flat as the corpus grows.
const MAX_POINTS_PER_CENTROID: usize = 256;

/// Seed for the training-subsample shuffle.
///
/// Fixed so every run on the same corpus picks the same training set
/// and produces the same centroids — the encoded index is a function
/// of the input pool alone. Callers that want a different sample can
/// shuffle `points` externally before calling [`fit`].
const TRAINING_SHUFFLE_SEED: u64 = 0x7BE0_0CBE_7BE0_0CBE;

/// Index of the centroid nearest to `point` under squared L2 distance.
///
/// Ties are broken by preferring the earlier centroid, which keeps the
/// function deterministic regardless of input length.
///
/// # Panics
///
/// Panics if `centroids.len()` is not a positive multiple of `dim`, or if
/// `point.len() != dim`.
///
/// # Examples
///
/// ```
/// use docbert_plaid::kmeans::nearest_centroid;
///
/// // Two centroids in 2-D at (0, 0) and (10, 10). Point (9, 9) is closer
/// // to the second centroid.
/// let centroids = [0.0, 0.0, 10.0, 10.0];
/// assert_eq!(nearest_centroid(&[9.0, 9.0], &centroids, 2), 1);
/// ```
pub fn nearest_centroid(point: &[f32], centroids: &[f32], dim: usize) -> usize {
    assert_eq!(
        point.len(),
        dim,
        "nearest_centroid: point length {} does not match dim {}",
        point.len(),
        dim,
    );
    assert!(dim > 0, "nearest_centroid: dim must be positive");
    assert!(
        !centroids.is_empty() && centroids.len().is_multiple_of(dim),
        "nearest_centroid: centroids length {} is not a positive multiple of dim {}",
        centroids.len(),
        dim,
    );

    let mut best_idx = 0usize;
    let mut best_dist = f32::INFINITY;
    for (i, centroid) in centroids.chunks_exact(dim).enumerate() {
        let d = squared_l2(point, centroid);
        if d < best_dist {
            best_dist = d;
            best_idx = i;
        }
    }
    best_idx
}

/// Recompute centroids as the mean of the points currently assigned to them.
///
/// This is the M-step of Lloyd's algorithm, expressed as candle ops so a
/// single implementation runs on whichever device the input tensors live
/// on — CUDA when the build has the `cuda` feature and the points were
/// uploaded to the GPU, the candle CPU backend otherwise. There's no
/// separate scalar fallback to keep in sync.
///
/// Builds a `[N, k]` one-hot of the assignments, then computes
/// `sums = oneᵀ × points` (`[k, dim]`) and `counts = oneᵀ.sum(1)`
/// (`[k, 1]`) in one matmul each. Empty clusters keep their `previous`
/// centroid via `where_cond` so the cluster count stays stable across
/// iterations — same behaviour as fast-plaid's reference.
///
/// `p_tensor` is `[N, dim]` f32. `assignments` is `[N]` u32 with values
/// in `0..k` (the contract of [`assign_as_tensor`]). `previous` is
/// `[k, dim]` f32. The returned tensor is `[k, dim]` f32 on the same
/// device.
///
/// # Errors
///
/// Returns [`PlaidError::Tensor`] if any underlying tensor op fails.
///
/// [`PlaidError::Tensor`]: crate::PlaidError::Tensor
pub fn update_centroids(
    p_tensor: &Tensor,
    assignments: &Tensor,
    previous: &Tensor,
) -> Result<Tensor> {
    let (n, _) = p_tensor.dims2()?;
    let (k, _) = previous.dims2()?;
    let device = p_tensor.device();
    let dtype = p_tensor.dtype();

    // No points → every cluster is empty → `previous` survives intact.
    // Skipping avoids calling matmul against an empty operand, which a
    // few candle backends are picky about.
    if n == 0 {
        return Ok(previous.clone());
    }

    // One-hot encode assignments. `assignments[i] == j` is broadcast
    // against an `arange(k)` row; the result lives on-device and is
    // cast to the points' dtype so the matmul below has matching
    // operands.
    let cluster_ids =
        Tensor::arange(0u32, k as u32, device)?.reshape((1, k))?;
    let one_hot = assignments
        .reshape((n, 1))?
        .broadcast_eq(&cluster_ids)?
        .to_dtype(dtype)?;

    // [k, N] × [N, dim] = [k, dim]; [k, N].sum(1) = [k, 1].
    let one_hot_t = one_hot.t()?.contiguous()?;
    let sums = one_hot_t.matmul(p_tensor)?;
    let counts = one_hot_t.sum_keepdim(1)?;

    // Empty clusters would divide by zero; clamp the divisor and then
    // overwrite those rows with `previous` via `where_cond`.
    let safe_counts = counts.clamp(1.0f32, f32::MAX)?;
    let means = sums.broadcast_div(&safe_counts)?;
    let empty_mask = counts.eq(0.0f32)?.broadcast_as(means.shape())?;
    Ok(empty_mask.where_cond(previous, &means)?)
}

/// Run Lloyd's algorithm starting from an explicit set of initial centroids.
///
/// Iterates "assign points to nearest centroid, recompute centroids" up to
/// `max_iters` times, stopping early when no point changes cluster between
/// iterations. Returns the final centroids as a flat `k × dim` buffer.
///
/// Taking the initial centroids as an argument keeps the function
/// deterministic and trivial to test. Random initialization (e.g.
/// sampling `k` distinct points) is layered on top in [`fit`].
///
/// Internally we allocate the points tensor exactly once and reuse it
/// across iterations, so the only per-iteration cost is the centroids
/// upload (`k × dim` floats — negligible) plus the matmul itself.
///
/// # Errors
///
/// Returns [`PlaidError::Tensor`] if any tensor allocation or matmul
/// fails (e.g. CUDA OOM).
///
/// # Panics
///
/// Panics on any shape violation between `points`, `initial`, and `dim`,
/// or if `dim == 0`.
///
/// [`PlaidError::Tensor`]: crate::PlaidError::Tensor
pub fn fit_with_init(
    points: &[f32],
    initial: &[f32],
    dim: usize,
    max_iters: usize,
) -> Result<Vec<f32>> {
    assert!(dim > 0, "fit_with_init: dim must be positive");
    assert!(
        initial.len().is_multiple_of(dim) && !initial.is_empty(),
        "fit_with_init: initial centroids length {} is not a positive multiple of dim {}",
        initial.len(),
        dim,
    );
    assert!(
        points.len().is_multiple_of(dim),
        "fit_with_init: points length {} is not a multiple of dim {}",
        points.len(),
        dim,
    );

    if points.is_empty() || max_iters == 0 {
        return Ok(initial.to_vec());
    }

    let n_points = points.len() / dim;
    let device = default_device();
    let p_tensor = Tensor::from_slice(points, (n_points, dim), device)?;
    fit_with_init_on_tensor(&p_tensor, initial, dim, max_iters)
}

/// Run Lloyd's algorithm against an already-uploaded `p_tensor`.
///
/// The whole loop stays on `p_tensor.device()`: assignments come back
/// from [`assign_as_tensor`] as a device-resident u32 tensor, the
/// M-step folds them into new centroids via candle ops in
/// [`update_centroids`], and the convergence check compares
/// successive assignment tensors via a single scalar host pull. Only
/// the final centroids round-trip back to a host `Vec<f32>` for the
/// caller.
///
/// Factored out so [`fit`] can upload the points tensor once and share
/// it with [`farthest_first_init_on_tensor`]. Duplicating the upload
/// cost 3.47 GB of VRAM on the docbert production corpus, which
/// pushed the build into CUDA OOM on 12 GB cards.
fn fit_with_init_on_tensor(
    p_tensor: &Tensor,
    initial: &[f32],
    dim: usize,
    max_iters: usize,
) -> Result<Vec<f32>> {
    let k = initial.len() / dim;
    let device = p_tensor.device();
    let mut centroids_t = Tensor::from_slice(initial, (k, dim), device)?;

    let mut previous_assignments: Option<Tensor> = None;
    for _ in 0..max_iters {
        let assignments = assign_as_tensor(p_tensor, &centroids_t)?;
        if let Some(prev) = &previous_assignments {
            let mismatches = assignments
                .ne(prev)?
                .to_dtype(DType::F32)?
                .sum_all()?
                .to_scalar::<f32>()?;
            if mismatches == 0.0 {
                break;
            }
        }
        centroids_t = update_centroids(p_tensor, &assignments, &centroids_t)?;
        previous_assignments = Some(assignments);
    }
    Ok(centroids_t.flatten_all()?.to_vec1::<f32>()?)
}

/// Run k-means over a subsample of `points`.
///
/// Follows fast-plaid's recipe (matching ColBERTv2's original paper):
///
/// 1. Sample up to `k · MAX_POINTS_PER_CENTROID` points with a seeded
///    shuffle. Training complexity becomes constant in corpus size.
/// 2. Use the first `k` rows of the shuffle as initial centroids.
///    The shuffle is random so the seeds are effectively a random
///    pick from the subsample, which is itself a random pick from
///    the full corpus.
/// 3. Run Lloyd's algorithm on the subsample with [`fit_with_init`].
///
/// `256 · k` is the sample size the PLAID paper validates as having
/// no measurable recall loss versus full-corpus training; this crate
/// targets that point on the quality / build-time curve.
///
/// # Errors
///
/// Returns [`PlaidError::Tensor`] on underlying tensor failure.
///
/// # Panics
///
/// Panics if `k == 0`, if `dim == 0`, if `points` is empty, or if
/// `points` has fewer than `k` rows.
///
/// [`PlaidError::Tensor`]: crate::PlaidError::Tensor
pub fn fit(
    points: &[f32],
    k: usize,
    dim: usize,
    max_iters: usize,
) -> Result<Vec<f32>> {
    assert!(k > 0, "fit: k must be positive");
    assert!(dim > 0, "fit: dim must be positive");
    assert!(
        points.len().is_multiple_of(dim),
        "fit: points length {} is not a multiple of dim {}",
        points.len(),
        dim,
    );
    let n_points = points.len() / dim;
    assert!(
        n_points >= k,
        "fit: need at least k={k} points, got {n_points}",
    );

    let (training_points, _) = sample_training_points(points, n_points, k, dim);

    let device = default_device();
    let n_train = training_points.len() / dim;
    let training_tensor =
        Tensor::from_slice(&training_points, (n_train, dim), device)?;
    let initial = training_points[..k * dim].to_vec();
    fit_with_init_on_tensor(&training_tensor, &initial, dim, max_iters)
}

/// Same as [`fit`] but accepts an already-uploaded `p_tensor` covering
/// the full corpus, so the caller can reuse that tensor after training.
///
/// The training phase itself still runs on a `k · MAX_POINTS_PER_CENTROID`
/// subsample — we don't train on the full corpus. `p_tensor` is kept
/// around for the post-training encode pass that runs downstream in
/// [`crate::codec::ResidualCodec::batch_encode_tokens_on_tensor`].
///
/// `points` must be the host-side backing buffer for `p_tensor`.
pub fn fit_on_tensor(
    p_tensor: &Tensor,
    points: &[f32],
    k: usize,
    dim: usize,
    max_iters: usize,
) -> Result<Vec<f32>> {
    let n = p_tensor.dim(0)?;
    let (training_points, same_as_pool) =
        sample_training_points(points, n, k, dim);
    let initial = training_points[..k * dim].to_vec();

    if same_as_pool {
        // Tiny corpus — reuse the caller's tensor directly.
        return fit_with_init_on_tensor(p_tensor, &initial, dim, max_iters);
    }

    // Otherwise build the training tensor from the subsample. The
    // subsample holds `k · MAX_POINTS_PER_CENTROID` rows at most so
    // the extra on-device allocation is bounded (≤ ~33 MiB at the
    // typical ColBERT-scale k/dim) and completely dwarfed by the
    // full pool tensor the caller already owns.
    let device = p_tensor.device();
    let n_train = training_points.len() / dim;
    let training_tensor =
        Tensor::from_slice(&training_points, (n_train, dim), device)?;
    fit_with_init_on_tensor(&training_tensor, &initial, dim, max_iters)
}

/// Sample up to `k · MAX_POINTS_PER_CENTROID` rows from `points` with
/// a seeded Fisher–Yates shuffle.
///
/// Returns `(sample, same_as_pool)` — when the corpus already has
/// `≤ k · MAX_POINTS_PER_CENTROID` rows, the caller's slice is cloned
/// as-is and `same_as_pool = true`, letting [`fit_on_tensor`] reuse
/// the pre-uploaded tensor without a second copy.
fn sample_training_points(
    points: &[f32],
    n: usize,
    k: usize,
    dim: usize,
) -> (Vec<f32>, bool) {
    let target = (k * MAX_POINTS_PER_CENTROID).min(n);
    if target == n {
        return (points.to_vec(), true);
    }

    let mut rng = StdRng::seed_from_u64(TRAINING_SHUFFLE_SEED);
    let mut indices: Vec<u32> = (0..n as u32).collect();
    // Partial Fisher–Yates: shuffle just the first `target` positions.
    // Equivalent to a full shuffle followed by `truncate(target)`,
    // minus the unused tail work.
    let (head, _) = indices.partial_shuffle(&mut rng, target);
    let head: Vec<u32> = head.to_vec();

    let mut sample = Vec::with_capacity(target * dim);
    for &i in &head {
        let start = i as usize * dim;
        sample.extend_from_slice(&points[start..start + dim]);
    }
    (sample, false)
}

/// Assign every point in `points` to the index of its nearest centroid.
///
/// `points` is a flat row-major `n_points × dim` buffer; `centroids` is a
/// flat row-major `k × dim` buffer. The returned vector has length
/// `n_points` and each entry is in `0..k`.
///
/// # Errors
///
/// Returns [`PlaidError::Tensor`] if any tensor allocation or matmul
/// fails (e.g. CUDA OOM).
///
/// # Panics
///
/// Panics if `points.len() % dim != 0`, if `centroids.len() % dim != 0`,
/// if `centroids.len() == 0`, or if `dim == 0`.
///
/// # Examples
///
/// ```
/// use docbert_plaid::kmeans::assign_points;
///
/// let centroids = [0.0, 0.0, 10.0, 10.0];
/// let points = [0.1, 0.1, 9.5, 9.5, -1.0, 0.0];
/// assert_eq!(assign_points(&points, &centroids, 2).unwrap(), vec![0, 1, 0]);
/// ```
///
/// [`PlaidError::Tensor`]: crate::PlaidError::Tensor
pub fn assign_points(
    points: &[f32],
    centroids: &[f32],
    dim: usize,
) -> Result<Vec<usize>> {
    assert!(dim > 0, "assign_points: dim must be positive");
    assert!(
        points.len().is_multiple_of(dim),
        "assign_points: points length {} is not a multiple of dim {}",
        points.len(),
        dim,
    );
    let n_points = points.len() / dim;
    if n_points == 0 {
        return Ok(Vec::new());
    }

    let device = default_device();
    let p = Tensor::from_slice(points, (n_points, dim), device)?;
    assign_on_tensor(&p, centroids, dim)
}

/// Same as [`assign_points`] but reuses a pre-uploaded points tensor.
///
/// Callers that run many assign passes against the same corpus (the
/// PLAID builder does k-means plus a final full encoding) can upload
/// the `[n, dim]` tensor once and hand it to each pass, saving both
/// the host→device copy and the risk of cudarc's caching allocator
/// holding two 3.47 GB blocks side by side.
///
/// # Errors
///
/// Returns [`PlaidError::Tensor`] if any tensor allocation or matmul
/// fails.
///
/// # Panics
///
/// Panics if `centroids.len() % dim != 0`, if `centroids.len() == 0`,
/// or if `dim == 0`.
///
/// [`PlaidError::Tensor`]: crate::PlaidError::Tensor
pub fn assign_on_tensor(
    p_tensor: &Tensor,
    centroids: &[f32],
    dim: usize,
) -> Result<Vec<usize>> {
    assert!(dim > 0, "assign_on_tensor: dim must be positive");
    assert!(
        !centroids.is_empty() && centroids.len().is_multiple_of(dim),
        "assign_on_tensor: centroids length {} is not a positive multiple of dim {}",
        centroids.len(),
        dim,
    );
    let k = centroids.len() / dim;
    let device = p_tensor.device();
    let c = Tensor::from_slice(centroids, (k, dim), device)?;
    assign_tensor(p_tensor, &c)
}

/// Compute nearest-centroid assignments for every row of `p_tensor`,
/// returning them as a device tensor of u32 indices.
///
/// Same E-step matmul as [`assign_on_tensor`] / [`assign_points`], but
/// keeps the argmin output on the device instead of pulling it back to
/// host. Callers that feed the result into another GPU op (e.g. the
/// `index_select` + residual + pack chain in
/// [`crate::codec::ResidualCodec::batch_encode_tokens_on_tensor`]) save
/// the host round-trip and any downstream upload.
///
/// The returned tensor has dtype `u32` and shape `[p_tensor.dim(0)]`.
///
/// # Errors
///
/// Returns [`PlaidError::Tensor`] on underlying tensor failure.
///
/// [`PlaidError::Tensor`]: crate::PlaidError::Tensor
pub fn assign_as_tensor(
    p_tensor: &Tensor,
    centroids_tensor: &Tensor,
) -> Result<Tensor> {
    let n = p_tensor.dim(0)?;
    if n == 0 {
        // Empty points → empty u32 tensor. Use zeros (safe dtype default).
        return Ok(Tensor::zeros(
            0,
            candle_core::DType::U32,
            p_tensor.device(),
        )?);
    }
    let k = centroids_tensor.dim(0)?.max(1);
    let bytes_per_row = k * std::mem::size_of::<f32>();
    let chunk_rows = (ASSIGN_CHUNK_BYTES / bytes_per_row).max(1).min(n);

    let c_t = centroids_tensor.t()?;
    let c_sq = centroids_tensor.sqr()?.sum_keepdim(1)?.t()?; // [1, k]

    let mut chunks: Vec<Tensor> = Vec::new();
    let mut start = 0usize;
    while start < n {
        let len = chunk_rows.min(n - start);
        let p_chunk = p_tensor.narrow(0, start, len)?;
        let dot = p_chunk.matmul(&c_t)?;
        let scores = dot.affine(-2.0, 0.0)?.broadcast_add(&c_sq)?;
        chunks.push(scores.argmin(1)?); // [len] u32
        start += len;
    }
    if chunks.len() == 1 {
        Ok(chunks.into_iter().next().unwrap())
    } else {
        Ok(Tensor::cat(&chunks, 0)?)
    }
}

/// Cap on the transient bytes materialised by a single assignment
/// matmul.
///
/// Each chunk produces three tensors of shape `[chunk_rows, k]`: the
/// matmul output, its affine-scaled twin, and the broadcast-added
/// score matrix. Sizing each of those to ~128 MiB keeps the working
/// set near ~384 MiB beside the resident points tensor — small enough
/// to fit on a 6 GiB GPU after the 3 GB pool upload, big enough that
/// cuBLAS kernels still amortise their launch overhead.
const ASSIGN_CHUNK_BYTES: usize = 128 * 1024 * 1024;

/// Tensor-based Lloyd E-step: returns the index of the nearest centroid
/// for every row of `p` against every row of `c`.
///
/// Delegates to [`assign_tensor_chunked`] with a chunk size derived
/// from [`ASSIGN_CHUNK_BYTES`]. Splitting this into a pair lets tests
/// drive the chunked path with tiny chunks without having to
/// synthesise gigabytes of input.
fn assign_tensor(p: &Tensor, c: &Tensor) -> Result<Vec<usize>> {
    let n = p.dim(0)?;
    if n == 0 {
        return Ok(Vec::new());
    }
    let k = c.dim(0)?.max(1);
    let bytes_per_row = k * std::mem::size_of::<f32>();
    let chunk_rows = (ASSIGN_CHUNK_BYTES / bytes_per_row).max(1).min(n);
    assign_tensor_chunked(p, c, chunk_rows)
}

/// Chunked implementation of the Lloyd E-step.
///
/// Argmin uses the formula `||c||² − 2·p·c` (the `||p||²` term is
/// constant for argmin per row, so we drop it). For typical ColBERT
/// centroids of ~unit norm, the dominant cost is the `[chunk, dim] ×
/// [dim, k]` matmul, which is exactly the GEMM candle is best at.
///
/// Point rows are sliced via [`Tensor::narrow`] — a zero-copy view —
/// so the caller's points tensor stays resident across chunks and
/// across Lloyd iterations. The centroid-side quantities (`c_t` and
/// `c_sq`) only depend on `c`, so they're computed once outside the
/// loop and reused.
fn assign_tensor_chunked(
    p: &Tensor,
    c: &Tensor,
    chunk_rows: usize,
) -> Result<Vec<usize>> {
    assert!(
        chunk_rows > 0,
        "assign_tensor_chunked: chunk_rows must be positive",
    );
    let n = p.dim(0)?;
    if n == 0 {
        return Ok(Vec::new());
    }
    let c_t = c.t()?;
    let c_sq = c.sqr()?.sum_keepdim(1)?.t()?; // [1, k]

    let mut out = Vec::with_capacity(n);
    let mut start = 0usize;
    while start < n {
        let len = chunk_rows.min(n - start);
        let p_chunk = p.narrow(0, start, len)?;
        let dot = p_chunk.matmul(&c_t)?; // [len, k]
        // scores[i, j] = ||c[j]||² - 2·p[i]·c[j], argmin gives nearest.
        let scores = dot.affine(-2.0, 0.0)?.broadcast_add(&c_sq)?;
        let argmin_u32: Vec<u32> = scores.argmin(1)?.to_vec1::<u32>()?;
        out.extend(argmin_u32.into_iter().map(|x| x as usize));
        start += len;
    }
    Ok(out)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn nearest_centroid_returns_zero_for_single_centroid() {
        let centroids = [1.0, 2.0, 3.0];
        assert_eq!(nearest_centroid(&[0.0, 0.0, 0.0], &centroids, 3), 0);
    }

    #[test]
    fn nearest_centroid_picks_the_closest_of_many() {
        // Centroids at (0,0), (5,5), (10,10). Point (4,4) is closest to
        // the middle one at index 1.
        let centroids = [0.0, 0.0, 5.0, 5.0, 10.0, 10.0];
        assert_eq!(nearest_centroid(&[4.0, 4.0], &centroids, 2), 1);
    }

    #[test]
    fn nearest_centroid_breaks_ties_toward_earlier_index() {
        // Point (5,0) is exactly the same distance from (0,0) and (10,0).
        // The earlier centroid wins.
        let centroids = [0.0, 0.0, 10.0, 0.0];
        assert_eq!(nearest_centroid(&[5.0, 0.0], &centroids, 2), 0);
    }

    #[test]
    #[should_panic(expected = "point length")]
    fn nearest_centroid_panics_on_dim_mismatch() {
        let centroids = [0.0, 0.0];
        let _ = nearest_centroid(&[1.0], &centroids, 2);
    }

    #[test]
    #[should_panic(expected = "centroids length")]
    fn nearest_centroid_panics_on_ragged_centroid_block() {
        // 3 f32s with dim=2 can't split cleanly into centroids.
        let centroids = [0.0, 0.0, 1.0];
        let _ = nearest_centroid(&[0.0, 0.0], &centroids, 2);
    }

    #[test]
    fn assign_points_maps_each_point_to_its_cluster() {
        let centroids = [0.0, 0.0, 10.0, 10.0];
        let points = [0.1, 0.2, 9.0, 9.5, 0.0, -0.1, 10.1, 10.1];
        assert_eq!(
            assign_points(&points, &centroids, 2).unwrap(),
            vec![0, 1, 0, 1],
        );
    }

    #[test]
    fn assign_points_handles_empty_input() {
        let centroids = [0.0, 1.0];
        assert_eq!(
            assign_points(&[], &centroids, 2).unwrap(),
            Vec::<usize>::new(),
        );
    }

    #[test]
    #[should_panic(expected = "points length")]
    fn assign_points_panics_on_ragged_points() {
        let centroids = [0.0, 0.0];
        let _ = assign_points(&[1.0, 2.0, 3.0], &centroids, 2);
    }

    #[test]
    fn assign_tensor_chunked_agrees_with_unchunked() {
        // Build a dataset with >1 row and exercise chunk boundaries
        // that are 1, prime, round, off-by-one, and the full length.
        // The chunked and unchunked paths must agree row-for-row
        // regardless of chunk size — if they ever diverge, the
        // per-chunk argmin has drifted from the per-row argmin.
        let dim = 3;
        let k = 4;
        let centroids: Vec<f32> = vec![
            0.0, 0.0, 0.0, //
            10.0, 0.0, 0.0, //
            0.0, 10.0, 0.0, //
            10.0, 10.0, 0.0, //
        ];
        let n = 173;
        let mut points: Vec<f32> = Vec::with_capacity(n * dim);
        for i in 0..n {
            let a = ((i * 37) % 11) as f32;
            let b = ((i * 53) % 13) as f32;
            points.extend_from_slice(&[a, b, 0.0]);
        }

        let device = default_device();
        let p = Tensor::from_slice(&points, (n, dim), device).unwrap();
        let c = Tensor::from_slice(&centroids, (k, dim), device).unwrap();

        let baseline = assign_tensor_chunked(&p, &c, n).unwrap();
        assert_eq!(baseline.len(), n);
        for chunk in [1usize, 7, 64, n - 1, n, n + 5] {
            let chunked = assign_tensor_chunked(&p, &c, chunk).unwrap();
            assert_eq!(
                chunked, baseline,
                "chunk_rows={chunk} must agree with the unchunked result",
            );
        }
    }

    // -- Lloyd M-step --

    /// Materialise a `[k, dim]` host slice from a `[k, dim]` tensor for
    /// equality assertions. Only used in M-step tests.
    fn centroids_to_vec(t: &Tensor) -> Vec<f32> {
        t.flatten_all().unwrap().to_vec1::<f32>().unwrap()
    }

    #[test]
    fn update_centroids_averages_assigned_points() {
        // Two clusters. Cluster 0 gets (0,0) and (2,0); cluster 1 gets (10,0).
        let device = default_device();
        let points = Tensor::from_vec(
            vec![0.0_f32, 0.0, 2.0, 0.0, 10.0, 0.0],
            (3, 2),
            device,
        )
        .unwrap();
        let assignments =
            Tensor::from_vec(vec![0u32, 0, 1], (3,), device).unwrap();
        let previous =
            Tensor::from_vec(vec![0.0_f32; 4], (2, 2), device).unwrap();

        let updated =
            update_centroids(&points, &assignments, &previous).unwrap();

        assert_eq!(centroids_to_vec(&updated), vec![1.0, 0.0, 10.0, 0.0]);
    }

    #[test]
    fn update_centroids_keeps_previous_for_empty_clusters() {
        // Cluster 1 receives no points; its centroid must survive.
        let device = default_device();
        let points =
            Tensor::from_vec(vec![0.0_f32, 0.0, 2.0, 0.0], (2, 2), device)
                .unwrap();
        let assignments =
            Tensor::from_vec(vec![0u32, 0], (2,), device).unwrap();
        let previous =
            Tensor::from_vec(vec![5.0_f32, 5.0, 99.0, -99.0], (2, 2), device)
                .unwrap();

        let updated =
            update_centroids(&points, &assignments, &previous).unwrap();
        let updated = centroids_to_vec(&updated);

        assert_eq!(updated[..2], [1.0, 0.0]);
        assert_eq!(updated[2..], [99.0, -99.0]);
    }

    // -- Lloyd driver --

    #[test]
    fn fit_with_init_converges_on_well_separated_clusters() {
        // Two tight clusters at (0,0) and (10,10). Starting from slightly
        // off-centre initial centroids, one round of Lloyd's is enough to
        // snap them to the true cluster means.
        let points = vec![
            0.0, 0.0, 0.2, -0.1, -0.1, 0.1, // near (0,0)
            10.0, 10.0, 10.1, 9.9, 9.9, 10.1, // near (10,10)
        ];
        let initial = [0.5, 0.5, 9.5, 9.5];

        let fitted = fit_with_init(&points, &initial, 2, 20).unwrap();

        assert_eq!(fitted.len(), 4);
        // Cluster 0 mean
        assert!((fitted[0] - 0.0333).abs() < 1e-3);
        assert!((fitted[1] - 0.0).abs() < 1e-3);
        // Cluster 1 mean
        assert!((fitted[2] - 10.0).abs() < 1e-3);
        assert!((fitted[3] - 10.0).abs() < 1e-3);
    }

    #[test]
    fn fit_with_init_is_idempotent_after_convergence() {
        // Running Lloyd again on an already-converged centroid set must
        // produce the same centroids (no further movement).
        let points = vec![0.0, 0.0, 1.0, 0.0, 10.0, 0.0, 11.0, 0.0];
        let initial = [0.5, 0.0, 10.5, 0.0];

        let once = fit_with_init(&points, &initial, 2, 50).unwrap();
        let twice = fit_with_init(&points, &once, 2, 50).unwrap();

        assert_eq!(once, twice);
    }

    #[test]
    fn fit_with_init_returns_initial_when_no_iterations_allowed() {
        let points = vec![0.0, 0.0, 10.0, 10.0];
        let initial = [1.0, 1.0, 9.0, 9.0];
        assert_eq!(
            fit_with_init(&points, &initial, 2, 0).unwrap(),
            initial.to_vec(),
        );
    }

    #[test]
    fn fit_converges_to_well_separated_centroids_on_bimodal_corpus() {
        // Two tight clusters, one near origin and one near (10, 10).
        // Random-subsample seeding can (and does) sometimes pick two
        // points from the same cluster, but Lloyd's algorithm pulls
        // them apart within a few iterations. 20 iters is the same
        // cap `PlaidBuildParams::default()` uses in production.
        let points = vec![
            0.0, 0.0, 0.1, 0.1, -0.1, 0.1, // cluster near origin
            10.0, 10.0, 10.1, 9.9, 9.9, 10.1, // cluster near (10, 10)
        ];
        let fitted = fit(&points, 2, 2, 20).unwrap();
        let c0 = &fitted[0..2];
        let c1 = &fitted[2..4];
        let dist = squared_l2(c0, c1);
        assert!(
            dist > 100.0,
            "expected well-separated centroids after Lloyd, got {c0:?} and {c1:?} (dist²={dist})",
        );
    }

    #[test]
    fn fit_is_deterministic_across_identical_calls() {
        // Farthest-first seeding must still be reproducible — no RNG.
        let points = vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8];
        let a = fit(&points, 3, 2, 10).unwrap();
        let b = fit(&points, 3, 2, 10).unwrap();
        assert_eq!(a, b);
    }

    #[test]
    #[should_panic(expected = "need at least")]
    fn fit_panics_when_asked_for_more_clusters_than_points() {
        let points = vec![0.0, 0.0];
        let _ = fit(&points, 5, 2, 1);
    }
}