rustyml 0.14.0

//! Mean Shift clustering
//!
//! Provides the [`MeanShift`] estimator, which finds clusters by shifting points toward
//! higher-density regions. The number of clusters does not need to be set in advance. Also
//! provides the [`estimate_bandwidth`] helper for choosing a bandwidth from the data

use crate::error::Error;
use crate::machine_learning::parallel::map_collect;
use crate::machine_learning::validation::{
    preliminary_check, validate_max_iterations, validate_predict_input, validate_tolerance,
};
use crate::math::matmul::{cache_resident, gemm_chunk_rows, matvec};
use crate::math::squared_euclidean_distance_row;
use crate::parallel_gates::scan_f64_parallel_min_elems;
use crate::{Deserialize, Serialize};
use ahash::AHashMap;
use gemmkit_ndarray::{Parallelism, dot};
use ndarray::{Array1, Array2, ArrayBase, Axis, Data, Ix2, Zip, s};
use ndarray_rand::rand::seq::SliceRandom;
use rayon::prelude::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator};

/// Mean Shift clustering algorithm
///
/// A centroid-based clustering algorithm that iteratively shifts data points toward areas of
/// higher density. Each point moves in the direction of the mean of points within its current
/// window until convergence. The number of clusters does not need to be set in advance
///
/// # Notes
///
/// - If unsure about a good bandwidth value, use the `estimate_bandwidth` function
/// - The bandwidth value strongly affects the result. Choose it based on the data
/// - For large datasets, set `bin_seeding = true` to reduce fit time
///
/// # Examples
///
/// ```rust
/// use rustyml::machine_learning::MeanShift;
/// use ndarray::Array2;
///
/// // Create a 2D dataset
/// let data = Array2::<f64>::from_shape_vec((10, 2),
///     vec![1.0, 2.0, 1.1, 2.2, 0.9, 1.9, 1.0, 2.1,
///          10.0, 10.0, 10.2, 9.9, 10.1, 10.0, 9.9, 9.8,
///          5.0, 5.0, 5.1, 4.9]).unwrap();
///
/// // Create a MeanShift instance with default parameters
/// let mut ms = MeanShift::default();
///
/// // Fit the model and predict cluster labels
/// let labels = ms.fit_predict(&data).unwrap();
///
/// // Get the cluster centers
/// let centers = ms.get_cluster_centers().clone().unwrap();
/// ```
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct MeanShift {
    /// Kernel bandwidth that sets the search radius. Larger values lead to fewer clusters
    bandwidth: f64,
    /// Maximum number of iterations to prevent infinite loops
    max_iter: usize,
    /// Convergence tolerance. A point has converged when it moves less than this value
    tol: f64,
    /// Whether to use the bin-seeding strategy for faster execution
    bin_seeding: bool,
    /// Whether to assign all points to clusters, including potential noise
    cluster_all: bool,
    /// Number of samples assigned to each cluster center
    n_samples_per_center: Option<Array1<usize>>,
    /// Final cluster centers found by the algorithm
    cluster_centers: Option<Array2<f64>>,
    /// Cluster labels assigned to each input sample. `-1` marks a point left unassigned when
    /// `cluster_all` is off
    labels: Option<Array1<isize>>,
    /// Actual number of iterations performed during fitting
    n_iter: Option<usize>,
}

impl Default for MeanShift {
    /// Creates a new MeanShift instance with default parameter values
    ///
    /// # Default Values
    ///
    /// - `bandwidth` - `1.0`. A larger value results in fewer clusters
    /// - `max_iter` - `300`. Maximum number of iterations, to prevent infinite loops
    /// - `tol` - `1e-3`. Convergence tolerance threshold
    /// - `bin_seeding` - `false`. Bin seeding is disabled by default
    /// - `cluster_all` - `true`. All data points are assigned to clusters by default
    ///
    /// # Returns
    ///
    /// - `Self` - A new MeanShift instance with default parameters
    fn default() -> Self {
        Self::new(1.0).expect("Default parameters should be valid")
    }
}

impl MeanShift {
    /// Creates a new MeanShift instance with the specified bandwidth
    ///
    /// `bandwidth` is the dominant hyperparameter. The remaining settings have sensible
    /// defaults. Tune them afterward with the builder methods listed in the Notes
    ///
    /// # Parameters
    ///
    /// - `bandwidth` - Bandwidth that determines the size of the kernel. Must be positive
    ///   and finite
    ///
    /// # Returns
    ///
    /// - `Result<Self, Error>` - A new MeanShift instance, or an Error
    ///
    /// # Notes
    ///
    /// Configure lower-priority settings after construction. The convergence-related
    /// setters validate their input and return `Result`. The boolean toggles return `Self`:
    ///
    /// - [`with_max_iter`](Self::with_max_iter) - maximum iterations (default: `300`)
    /// - [`with_tolerance`](Self::with_tolerance) - convergence tolerance (default: `1e-3`)
    /// - [`with_bin_seeding`](Self::with_bin_seeding) - bin-seeding for faster init
    ///   (default: `false`)
    /// - [`with_cluster_all`](Self::with_cluster_all) - assign all points to clusters
    ///   (default: `true`)
    ///
    /// # Errors
    ///
    /// - `Error::InvalidParameter` - If `bandwidth` is non-positive or not finite
    pub fn new(bandwidth: f64) -> Result<Self, Error> {
        if bandwidth <= 0.0 || !bandwidth.is_finite() {
            return Err(Error::invalid_parameter(
                "bandwidth",
                format!("must be positive and finite, got {}", bandwidth),
            ));
        }

        Ok(MeanShift {
            bandwidth,
            max_iter: 300,
            tol: 1e-3,
            bin_seeding: false,
            cluster_all: true,
            n_samples_per_center: None,
            cluster_centers: None,
            labels: None,
            n_iter: None,
        })
    }

    /// Sets the maximum number of iterations (default: `300`)
    ///
    /// # Errors
    ///
    /// - `Error::InvalidParameter` - If `max_iter` is 0
    pub fn with_max_iter(mut self, max_iter: usize) -> Result<Self, Error> {
        validate_max_iterations(max_iter)?;
        self.max_iter = max_iter;
        Ok(self)
    }

    /// Sets the convergence tolerance (default: `1e-3`)
    ///
    /// # Errors
    ///
    /// - `Error::InvalidParameter` - If `tol` is non-positive or not finite
    pub fn with_tolerance(mut self, tol: f64) -> Result<Self, Error> {
        validate_tolerance(tol)?;
        self.tol = tol;
        Ok(self)
    }

    /// Enables or disables the bin-seeding initialization strategy (default: `false`)
    ///
    /// Set this to `true` to reduce fit time on large datasets
    pub fn with_bin_seeding(mut self, bin_seeding: bool) -> Self {
        self.bin_seeding = bin_seeding;
        self
    }

    /// Sets whether to assign every point to a cluster, including potential noise
    /// (default: `true`)
    ///
    /// With this off, the model labels a point further than one bandwidth from every cluster
    /// center as `-1`, as in scikit-learn
    pub fn with_cluster_all(mut self, cluster_all: bool) -> Self {
        self.cluster_all = cluster_all;
        self
    }

    // Getters
    get_field!(get_bandwidth, bandwidth, f64);
    get_field_as_ref!(get_cluster_centers, cluster_centers, Option<&Array2<f64>>);
    get_field_as_ref!(get_labels, labels, Option<&Array1<isize>>);
    get_field_as_ref!(
        get_n_samples_per_center,
        n_samples_per_center,
        Option<&Array1<usize>>
    );
    get_field!(get_actual_iterations, n_iter, Option<usize>);
    get_field!(get_max_iterations, max_iter, usize);
    get_field!(get_tolerance, tol, f64);
    get_field!(get_bin_seeding, bin_seeding, bool);
    get_field!(get_cluster_all, cluster_all, bool);

    /// Fits the MeanShift clustering model to the input data
    ///
    /// # Parameters
    ///
    /// - `x` - The input data where each row is a sample
    ///
    /// # Returns
    ///
    /// - `Result<&mut Self, Error>` - A mutable reference to the fitted model, or an Error
    ///
    /// # Errors
    ///
    /// - `Error::EmptyInput` - If the input data is empty
    /// - `Error::NonFinite` - If the input data contains NaN or infinite values
    ///
    /// # Performance
    ///
    /// Parallelizes when the total scan work clears the calibrated scan-class gate (see
    /// `crate::parallel_gates`)
    pub fn fit<S>(&mut self, x: &ArrayBase<S, Ix2>) -> Result<&mut Self, Error>
    where
        S: Data<Elem = f64> + Send + Sync,
    {
        preliminary_check(x, None)?;

        let n_samples = x.shape()[0];
        let n_features = x.shape()[1];

        // Initialize seed points
        let seeds: Vec<usize> = if self.bin_seeding {
            self.get_bin_seeds(x)
        } else {
            (0..n_samples).collect()
        };

        let tol_squared = self.tol * self.tol;
        let bandwidth_squared = self.bandwidth * self.bandwidth;
        // Per-sample squared norms, shared across all seeds and iterations
        let x_sq = x.map_axis(Axis(1), |row| row.dot(&row));

        // Each seed task runs O(iterations * n * d) of work
        let use_parallel = seeds
            .len()
            .saturating_mul(n_samples)
            .saturating_mul(n_features)
            >= scan_f64_parallel_min_elems();

        // When the seed axis alone fills the pool, keep each per-seed matvec serial. This avoids
        // nested rayon forks
        let seed_matvec_par = if use_parallel && seeds.len() >= rayon::current_num_threads() {
            Parallelism::Serial
        } else {
            Parallelism::Rayon(0)
        };

        // Mean shift on a single seed. The third tuple element is how many points the window
        // held at convergence, the mode's intensity. The merge phase below uses it to rank modes
        let process_seed = |seed_idx: usize| -> (Array1<f64>, usize, f64) {
            let mut center = x.row(seed_idx).to_owned();
            let mut completed_iterations = 0;

            let window_weight = loop {
                // Flat kernel: membership in the bandwidth ball, as a 0/1 weight, computed through
                // the squared-norm identity. Writing it as a weight vector keeps the 2 products
                // below as single GEMV sweeps. `weighted_sum / weight_sum` is then the mean of
                // the points inside the ball, and `weight_sum` is exactly how many points that is
                let center_sq = center.dot(&center);
                let projections = matvec(x, &center, seed_matvec_par);
                let weights: Array1<f64> =
                    Zip::from(&projections)
                        .and(&x_sq)
                        .map_collect(|&proj, &x_norm_sq| {
                            let dist_sq = (center_sq + x_norm_sq - 2.0 * proj).max(0.0);
                            if dist_sq <= bandwidth_squared {
                                1.0
                            } else {
                                0.0
                            }
                        });
                // Serial sum
                let weight_sum = weights.sum();

                let weighted_sum = matvec(&x.t(), &weights, seed_matvec_par);
                let new_center = resolve_shifted_center(weighted_sum, weight_sum, &center);

                // Check convergence using squared distance to avoid sqrt
                let shift_squared = squared_euclidean_distance_row(&center, &new_center);
                center = new_center;

                completed_iterations += 1;

                // An empty ball cannot move the seed, so stop rather than spin. scikit-learn
                // takes the same early exit when a seed's neighborhood is empty
                if shift_squared < tol_squared
                    || completed_iterations >= self.max_iter
                    || weight_sum == 0.0
                {
                    break weight_sum;
                }
            };

            (center, completed_iterations, window_weight)
        };

        #[cfg(feature = "show_progress")]
        let progress_bar = {
            let pb = crate::create_progress_bar(
                seeds.len() as u64,
                "[{elapsed_precise}] {bar:40} {pos}/{len} seeds | {msg}",
            );
            pb.set_message("Processing seeds...");
            pb
        };

        let results: Vec<(Array1<f64>, usize, f64)> = if use_parallel {
            seeds
                .par_iter()
                .map(|&seed_idx| {
                    let result = process_seed(seed_idx);
                    #[cfg(feature = "show_progress")]
                    progress_bar.inc(1);
                    result
                })
                .collect()
        } else {
            seeds
                .iter()
                .map(|&seed_idx| {
                    let result = process_seed(seed_idx);
                    #[cfg(feature = "show_progress")]
                    progress_bar.inc(1);
                    result
                })
                .collect()
        };
        #[cfg(feature = "show_progress")]
        progress_bar.finish_with_message("All seeds processed");

        // Extract centers and calculate actual max iterations
        let centers: Vec<Array1<f64>> = results.iter().map(|(c, _, _)| c.clone()).collect();
        let max_actual_iter = results.iter().map(|(_, i, _)| *i).max().unwrap_or(0);
        self.n_iter = Some(max_actual_iter);

        // Merges duplicate modes, scikit-learn's way. Ranks the converged modes by how many
        // points their window held. Then walks that order, keeping each mode and suppressing
        // every other mode within one bandwidth of it. The kept center is the dense mode itself,
        // not an average of the suppressed modes. So the result does not depend on the order in
        // which the pass processes the seeds
        let mut order: Vec<usize> = (0..centers.len()).collect();
        order.sort_by(|&a, &b| {
            // Descending by window weight, then descending by coordinates, so equally dense
            // modes come out in the same order as scikit-learn's `(intensity, center)` sort.
            // Only the cluster numbering depends on this order, not the clustering itself
            results[b]
                .2
                .partial_cmp(&results[a].2)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| {
                    centers[b]
                        .iter()
                        .zip(centers[a].iter())
                        .map(|(x, y)| x.total_cmp(y))
                        .find(|ord| *ord != std::cmp::Ordering::Equal)
                        .unwrap_or(std::cmp::Ordering::Equal)
                })
        });

        let mut unique_centers: Vec<Array1<f64>> = Vec::with_capacity(centers.len());
        for &idx in &order {
            let center = &centers[idx];
            let suppressed = unique_centers
                .iter()
                .any(|kept| squared_euclidean_distance_row(center, kept) <= bandwidth_squared);
            if !suppressed {
                unique_centers.push(center.clone());
            }
        }

        // Create cluster_centers array
        let n_clusters = unique_centers.len();
        let mut cluster_centers = Array2::zeros((n_clusters, n_features));
        for (i, center) in unique_centers.iter().enumerate() {
            cluster_centers.row_mut(i).assign(center);
        }

        // Find the nearest cluster label for one sample
        let find_label = |i: usize| -> isize {
            let point = x.row(i);
            let mut min_dist_squared = f64::INFINITY;
            let mut label = 0;

            for (j, center) in unique_centers.iter().enumerate() {
                let dist_squared = squared_euclidean_distance_row(&point, center);
                if dist_squared < min_dist_squared {
                    min_dist_squared = dist_squared;
                    label = j;
                }
            }

            // With `cluster_all` off, a point further than one bandwidth from every center is
            // noise and gets scikit-learn's `-1`. It does not get an index one past the last
            // real cluster
            if !self.cluster_all && min_dist_squared > bandwidth_squared {
                -1
            } else {
                label as isize
            }
        };

        // Assign cluster labels to each data point (scan-class gate: n tasks, each an
        // O(centers * d) distance scan)
        let label_work = n_samples
            .saturating_mul(n_clusters)
            .saturating_mul(x.ncols());
        let labels = map_collect(
            n_samples,
            label_work >= scan_f64_parallel_min_elems(),
            find_label,
        );

        // Count how many samples belong to each cluster center, skipping the noise label
        let mut samples_per_center = vec![0usize; n_clusters];
        for &label in &labels {
            if label >= 0 {
                samples_per_center[label as usize] += 1;
            }
        }

        self.cluster_centers = Some(cluster_centers);
        self.labels = Some(Array1::from(labels));
        self.n_samples_per_center = Some(Array1::from(samples_per_center));

        Ok(self)
    }

    /// Predicts cluster labels for the input data
    ///
    /// # Parameters
    ///
    /// - `x` - The input data where each row is a sample
    ///
    /// # Returns
    ///
    /// - `Result<Array1<isize>, Error>` - The predicted cluster labels (`-1` for noise), or
    ///   an Error
    ///
    /// # Errors
    ///
    /// - `Error::NotFitted` - If the model has not been fitted yet
    /// - `Error::EmptyInput` - If the input data is empty
    /// - `Error::DimensionMismatch` - If the feature count does not match the training data
    /// - `Error::NonFinite` - If the input data contains NaN or infinite values
    ///
    /// # Performance
    ///
    /// Parallelizes when the total scan work clears the calibrated scan-class gate (see
    /// `crate::parallel_gates`)
    pub fn predict<S>(&self, x: &ArrayBase<S, Ix2>) -> Result<Array1<isize>, Error>
    where
        S: Data<Elem = f64> + Sync,
    {
        let centers = self
            .cluster_centers
            .as_ref()
            .ok_or_else(|| Error::not_fitted("MeanShift"))?;

        // Validate input against the feature count seen during fitting
        validate_predict_input(x, centers.ncols())?;

        let n_samples = x.nrows();
        let n_clusters = centers.nrows();
        let bandwidth_squared = self.bandwidth * self.bandwidth;

        // Nearest cluster center for one sample. Returns the noise label `-1` when `cluster_all`
        // is off and the point lies farther than the bandwidth from every center
        let find_nearest = |i: usize| -> isize {
            let point = x.row(i);
            let mut min_dist_squared = f64::INFINITY;
            let mut label = 0;

            for j in 0..n_clusters {
                let center = centers.row(j);
                let dist_squared = squared_euclidean_distance_row(&point, &center);
                if dist_squared < min_dist_squared {
                    min_dist_squared = dist_squared;
                    label = j;
                }
            }

            if !self.cluster_all && min_dist_squared > bandwidth_squared {
                -1
            } else {
                label as isize
            }
        };

        // Scan-class gate: n tasks, each an O(centers * d) distance scan
        let label_work = n_samples
            .saturating_mul(n_clusters)
            .saturating_mul(x.ncols());
        let labels = map_collect(
            n_samples,
            label_work >= scan_f64_parallel_min_elems(),
            find_nearest,
        );

        Ok(Array1::from(labels))
    }

    /// Fits the model to the input data and predicts cluster labels
    ///
    /// # Parameters
    ///
    /// - `x` - The input data where each row is a sample
    ///
    /// # Returns
    ///
    /// - `Result<Array1<isize>, Error>` - The predicted cluster labels (`-1` for noise), or
    ///   an Error
    ///
    /// # Errors
    ///
    /// - `Error::EmptyInput` - If the input data is empty
    /// - `Error::NonFinite` - If the input data contains NaN or infinite values
    pub fn fit_predict<S>(&mut self, x: &ArrayBase<S, Ix2>) -> Result<Array1<isize>, Error>
    where
        S: Data<Elem = f64> + Sync + Send,
    {
        self.fit(x)?;
        Ok(self.labels.clone().unwrap())
    }

    /// Computes seed points by binning the feature space onto a grid
    ///
    /// Maps each sample to a grid cell of side length `bandwidth` and returns one representative
    /// point per non-empty cell. This reduces the number of seeds the mean-shift iterations must
    /// process on dense datasets
    ///
    /// # Parameters
    ///
    /// - `x` - The input data where each row is a sample
    ///
    /// # Returns
    ///
    /// - `Vec<usize>` - Indices of the selected seed points (one per occupied grid cell)
    fn get_bin_seeds<S>(&self, x: &ArrayBase<S, Ix2>) -> Vec<usize>
    where
        S: Data<Elem = f64> + Sync + Send,
    {
        let n_samples = x.shape()[0];
        let n_features = x.shape()[1];

        // Calculate min for each feature (scan-class gate: d tasks of an O(n) column scan)
        let scan_parallel = n_samples.saturating_mul(n_features) >= scan_f64_parallel_min_elems();
        let col_min = |j: usize| {
            let col = x.column(j);
            col.fold(f64::INFINITY, |a, &b| a.min(b))
        };
        let mins: Vec<f64> = if scan_parallel {
            (0..n_features).into_par_iter().map(col_min).collect()
        } else {
            (0..n_features).map(col_min).collect()
        };

        // Builds the grid in a shared HashMap, which makes this part harder to parallelize
        let bin_size = self.bandwidth;

        let bins_mutex = std::sync::Mutex::new(AHashMap::<Vec<i64>, Vec<usize>>::new());

        // Assign points to bins (same scan-class gate: n tasks of O(d) quantization each)
        let assign_bin = |i: usize| {
            let point = x.row(i);
            let mut bin_index = Vec::with_capacity(n_features);

            for j in 0..n_features {
                let idx = ((point[j] - mins[j]) / bin_size).floor() as i64;
                bin_index.push(idx);
            }

            // Lock the HashMap only when updating
            let mut bins = bins_mutex.lock().unwrap();
            bins.entry(bin_index).or_default().push(i);
        };
        if scan_parallel {
            (0..n_samples).into_par_iter().for_each(assign_bin);
        } else {
            (0..n_samples).for_each(assign_bin);
        }

        let bins = bins_mutex.into_inner().unwrap();

        // One seed per cell
        let mut seeds = Vec::new();
        for (_, indices) in bins {
            if let Some(&min_idx) = indices.iter().min() {
                seeds.push(min_idx);
            }
        }
        seeds.sort_unstable();

        seeds
    }

    model_save_and_load_methods!(MeanShift);
}

/// Estimates a bandwidth to use with the MeanShift algorithm
///
/// Matches scikit-learn's `estimate_bandwidth`. With `k = max(1, floor(n * quantile))`, it takes
/// each point's distance to its `k`-th nearest neighbor and returns the mean of those distances.
/// That is a local-density statistic. It answers how far a typical point is from the edge of
/// its own neighborhood, which is what a bandwidth needs to be. A quantile of the whole
/// pairwise-distance distribution is a global spread statistic instead. It runs much larger on
/// clustered data and collapses everything into one cluster
///
/// # Parameters
///
/// - `x` - The input data where each row is a sample
/// - `quantile` - The quantile of the pairwise distances to use as the bandwidth. Defaults
///   to `0.3`
/// - `n_samples` - The number of samples to use for the distance calculation. Clamped to the
///   dataset size. Defaults to all rows
/// - `random_state` - Seed for random number generation
///
/// # Returns
///
/// - `Result<f64, Error>` - The estimated bandwidth, or an Error
///
/// # Errors
///
/// - `Error::InvalidParameter` - If `quantile` is not in the open range (0, 1)
pub fn estimate_bandwidth<S>(
    x: &ArrayBase<S, Ix2>,
    quantile: Option<f64>,
    n_samples: Option<usize>,
    random_state: Option<u64>,
) -> Result<f64, Error>
where
    S: Data<Elem = f64>,
{
    let quantile = quantile.unwrap_or(0.3);
    if quantile <= 0.0 || quantile >= 1.0 {
        return Err(Error::invalid_parameter(
            "quantile",
            "must be in the open range (0, 1)",
        ));
    }

    let (n_samples_total, _) = x.dim();
    // Clamp the requested sample count to the dataset size
    let n_samples = n_samples.unwrap_or(n_samples_total).min(n_samples_total);

    let mut rng = crate::random::make_rng(random_state);

    // When the requested count covers every row, use all samples. Otherwise, sample at random
    let x_samples = if n_samples >= n_samples_total {
        x.to_owned()
    } else {
        let mut indices: Vec<usize> = (0..n_samples_total).collect();
        indices.shuffle(&mut rng);
        let indices = &indices[..n_samples];

        let mut samples = Array2::zeros((n_samples, x.ncols()));
        for (i, &idx) in indices.iter().enumerate() {
            samples.row_mut(i).assign(&x.row(idx));
        }
        samples
    };

    let x_sq = x_samples.map_axis(Axis(1), |row| row.dot(&row));
    let distances: Vec<f64> = if cache_resident::<f64>(n_samples, x_samples.ncols()) {
        (0..n_samples)
            .into_par_iter()
            .flat_map(|i| {
                // Forced serial: this already runs one task per row, so a matvec that forked
                // again would nest inside its own rayon task
                let proj_row = matvec(&x_samples, &x_samples.row(i), Parallelism::Serial);
                ((i + 1)..n_samples)
                    .map(|j| (x_sq[i] + x_sq[j] - 2.0 * proj_row[j]).max(0.0).sqrt())
                    .collect::<Vec<f64>>()
            })
            .collect()
    } else {
        let chunk_rows = gemm_chunk_rows(n_samples);
        let mut distances: Vec<f64> = Vec::new();
        for chunk_start in (0..n_samples).step_by(chunk_rows) {
            let chunk_end = (chunk_start + chunk_rows).min(n_samples);
            let projections = dot(
                &x_samples.slice(s![chunk_start..chunk_end, ..]),
                &x_samples.t(),
            );
            let chunk: Vec<f64> = (chunk_start..chunk_end)
                .into_par_iter()
                .flat_map(|i| {
                    let proj_row = projections.row(i - chunk_start);
                    ((i + 1)..n_samples)
                        .map(|j| (x_sq[i] + x_sq[j] - 2.0 * proj_row[j]).max(0.0).sqrt())
                        .collect::<Vec<f64>>()
                })
                .collect();
            distances.extend(chunk);
        }
        distances
    };

    if n_samples < 2 {
        // scikit-learn returns 0.0 here and lets `MeanShift` reject it, rather than erroring in
        // the helper
        return Ok(0.0);
    }

    // scikit-learn's neighbor index counts the `k` closest points including the query point
    // itself, then takes the largest of them. So the statistic is really each point's distance
    // to its `(k - 1)`-th nearest neighbor. Reproducing this off-by-one matters: at quantile 0.3
    // on 10 points, reading the k-th instead of the (k-1)-th moves the bandwidth from 1.1229 to
    // 1.2213
    let k = ((n_samples as f64 * quantile) as usize).max(1);
    if k < 2 {
        // Only the query point itself falls inside the window, so its farthest "neighbor" is
        // itself, at distance zero. scikit-learn returns 0.0 here too and lets `MeanShift` reject
        // the bandwidth rather than erroring in this helper
        return Ok(0.0);
    }
    let neighbour_rank = (k - 2).min(n_samples - 2);

    let total: f64 = (0..n_samples)
        .map(|i| {
            let mut row: Vec<f64> = (0..n_samples)
                .filter(|&j| j != i)
                .map(|j| distances[pair_index(i, j, n_samples)])
                .collect();
            row.select_nth_unstable_by(neighbour_rank, f64::total_cmp);
            row[neighbour_rank]
        })
        .sum();

    Ok(total / n_samples as f64)
}

/// Index of the `(i, j)` pair inside the flattened strict upper triangle of a distance matrix
///
/// The layout stores the triangle row by row. Row `i` starts at `i * n - i * (i + 1) / 2`.
/// It holds the distances to `i + 1 ..= n - 1`. The matrix is symmetric, so swapping `i` and
/// `j` answers a query where `i > j`
#[inline]
fn pair_index(i: usize, j: usize, n: usize) -> usize {
    let (lo, hi) = if i < j { (i, j) } else { (j, i) };
    lo * n - lo * (lo + 1) / 2 + (hi - lo - 1)
}

/// Resolves the shifted center for one mean-shift iteration
///
/// Normally the new center is the count-normalized mean `weighted_sum / weight_sum`. When the
/// bandwidth ball is empty, the total weight is exactly zero. The function then keeps the
/// center in place instead of collapsing it to the origin. Otherwise, a spurious cluster
/// center unrelated to the data would appear
///
/// # Parameters
///
/// - `weighted_sum` - Sum of the points inside the bandwidth ball
/// - `weight_sum` - Number of points inside the bandwidth ball
/// - `center` - Current center, returned unchanged when `weight_sum` is zero
///
/// # Returns
///
/// - `Array1<f64>` - The shifted center for the next iteration
fn resolve_shifted_center(
    weighted_sum: Array1<f64>,
    weight_sum: f64,
    center: &Array1<f64>,
) -> Array1<f64> {
    if weight_sum > 0.0 {
        weighted_sum / weight_sum
    } else {
        center.to_owned()
    }
}

/// Unit tests for the `resolve_shifted_center` helper
#[cfg(test)]
mod tests {
    use super::*;
    use ndarray::array;

    /// A degenerate window (all weights underflowed to 0) must keep the current center,
    /// not collapse it to the origin
    #[test]
    fn resolve_shifted_center_keeps_center_when_weights_underflow() {
        let center = array![3.0, -4.0, 5.0];
        // weight_sum == 0 implies the weighted sum is the zero vector
        let weighted_sum = array![0.0, 0.0, 0.0];
        let out = resolve_shifted_center(weighted_sum, 0.0, &center);
        assert_eq!(
            out, center,
            "zero-weight window must keep the current center, got {out:?}"
        );
    }

    /// With positive total weight the center is the weight-normalized mean
    #[test]
    fn resolve_shifted_center_normalizes_when_weights_present() {
        let center = array![100.0, 100.0];
        let weighted_sum = array![6.0, 9.0];
        let out = resolve_shifted_center(weighted_sum, 3.0, &center);
        assert_eq!(out, array![2.0, 3.0]);
    }
}