//! Density Classifier: Discover structure in embedding spaces
//!
//! Implements Krapivin-style PCA-based LSH for density-aware classification.
//! See: Krapivin et al. (2025) "Optimal Bounds for Open Addressing Without Reordering"
//! <https://arxiv.org/abs/2501.02305>
//!
//! ## Algorithm
//!
//! Two-stage hashing ensures hyperplanes align with data structure:
//! 1. Random hash → initial buckets → compute bucket centroids
//! 2. PCA on centroids → data-adapted hyperplanes
//! 3. Re-hash with PCA hyperplanes → better bucket distribution
//!
//! ## Outputs per item
//!
//! - bucket_id: LSH bucket assignment
//! - bucket_size: Number of items in the bucket
//! - centroid_similarity: Cosine similarity to bucket centroid (0-1)
//! - isolation_score: How isolated the item is (top_k_sim - median_sim)
//! - stability_score: How stable bucket assignment is across multiple seeds (0-1)

use nalgebra::DMatrix;
use ndarray::{Array2, ArrayView2};
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;
use rayon::prelude::*;
use std::collections::HashMap;

use crate::compute::ComputeBackend;

// Link BLAS
extern crate blas_src;

/// Multi-resolution analysis results - finds structure among outliers at coarser resolutions
///
/// Since PCA-LSH bit 0 (LSB) = highest variance component, masking the low K bits
/// gives the K most informative bits. This enables "recovery" - finding structure
/// among outliers at coarser resolutions - essentially for free.
#[derive(Debug, Clone)]
pub struct MultiResolutionAnalysis {
    /// Number of items analyzed
    pub n_items: usize,
    /// Number of bits in the full-resolution bucket IDs
    pub num_bits: usize,
    /// Per-item: smallest depth K (1..num_bits) where coarse bucket_size >= dense_threshold.
    /// 0 means already dense at full resolution. num_bits+1 means never recovered.
    pub recovery_depth: Vec<u32>,
    /// Per-item: bucket_size / expected_size at recovery depth.
    /// expected_size = n_items / 2^K. Ratio > 1 means denser than expected.
    pub recovery_ratio: Vec<f32>,
    /// Per depth K (index 0 = K=1, index num_bits-1 = K=num_bits):
    /// number of unique coarse buckets at that depth
    pub buckets_per_depth: Vec<usize>,
    /// Per depth K: mean coarse bucket size
    pub mean_size_per_depth: Vec<f32>,
    /// Dense threshold used for analysis
    pub dense_threshold: u32,
}

/// Bridge analysis results - reveals region connectivity structure
#[derive(Debug, Clone)]
pub struct BridgeAnalysis {
    /// Indices of bridge items (low centroid similarity)
    pub bridge_indices: Vec<usize>,
    /// For each bridge: (item_idx, primary_bucket, connected_buckets)
    pub bridge_connections: Vec<(usize, u64, Vec<u64>)>,
    /// Unique bucket IDs in order
    pub bucket_ids: Vec<u64>,
    /// Adjacency matrix: bucket_i -> bucket_j -> bridge count
    /// Flattened as Vec of length n_buckets * n_buckets
    pub adjacency_matrix: Vec<u32>,
    /// Number of unique buckets
    pub n_buckets: usize,
    /// Threshold used for bridge detection
    pub bridge_threshold: f32,
}

/// Bridge persistence analysis — multi-resolution connector detection
///
/// # Why relative thresholds?
///
/// At each LSH depth K, items are grouped into 2^K coarse buckets. A "connector"
/// is an item that's similar to a bucket centroid other than its own. The naive
/// approach uses an absolute similarity threshold (e.g., sim >= 0.3), but this
/// fails across depths:
///
/// - At K=1 (2 buckets), centroids are broad averages. Median max-other-sim ≈ 0.0.
///   A threshold of 0.3 catches almost nothing.
/// - At K=12 (thousands of buckets), centroids are tight. Median max-other-sim ≈ 0.5.
///   A threshold of 0.3 catches 99.5% of items — meaningless.
///
/// The **relative threshold** normalizes this: `max_other_sim / own_sim >= T`.
/// This ratio is stable across depths (median ≈ 0.8-0.98) because both numerator
/// and denominator scale together as centroids get tighter. A ratio ≥ 0.8 means
/// "this item is at least 80% as similar to another bucket's centroid as to its own."
///
/// # Algorithm
///
/// For each depth K = 1..num_bits:
/// 1. Mask bucket IDs to K low bits → coarse buckets
/// 2. Compute coarse centroids (O(n × d))
/// 3. BLAS matmul: all items × all centroids → similarity matrix (n × n_coarse)
/// 4. Per-item: own_sim = sim to own centroid. For each other centroid,
///    if other_sim / own_sim >= relative_threshold → connector to that bucket.
/// 5. Map connected coarse pairs to full-res pairs via u64 depth bitmask.
///
/// # Output
///
/// - `bridge_persistence[i]`: number of depths where item i is a connector (0..num_bits)
/// - `max_bridge_depth[i]`: coarsest depth where item i connects (0 = never)
/// - `bridge_ratio[i]`: peak max_other/own ratio across all depths (strength of connection)
/// - `persistent_pairs`: full-res bucket pairs ranked by how many depths maintain the link
#[derive(Debug, Clone)]
pub struct BridgePersistence {
    /// Number of items analyzed
    pub n_items: usize,
    /// Number of bits in the full-resolution bucket IDs
    pub num_bits: usize,
    /// Per-item: count of depths where this item is a connector (0..num_bits)
    pub bridge_persistence: Vec<u32>,
    /// Per-item: coarsest depth where item is a connector (0 = never)
    pub max_bridge_depth: Vec<u32>,
    /// Per-item: finest depth where item is a connector (0 = never)
    pub min_bridge_depth: Vec<u32>,
    /// Per-item: peak relative similarity (max_other_sim / own_sim) across all depths.
    /// Higher = stronger cross-bucket affinity. 0.0 = never a connector.
    pub bridge_ratio: Vec<f32>,
    /// Per depth K (index 0 = K=1): number of connector items at that depth
    pub bridges_per_depth: Vec<usize>,
    /// Per depth K: number of connected bucket pairs at that depth
    pub pairs_per_depth: Vec<usize>,
    /// Per depth K: number of unique coarse buckets at that depth
    pub buckets_per_depth: Vec<usize>,
    /// Relative threshold used: items with max_other_sim/own_sim >= this are connectors
    pub relative_threshold: f32,
    /// Persistent bucket pairs: (bucket_a, bucket_b, persistence_count, min_depth, max_depth)
    pub persistent_pairs: Vec<(u64, u64, u32, u32, u32)>,
}

/// Report from density classification
#[derive(Debug, Clone)]
pub struct DensityReport {
    pub corpus_size: usize,
    pub num_buckets: usize,
    pub stage1_variance_explained: f32,
    pub mean_bucket_size: f32,
    pub median_bucket_size: usize,
    pub max_bucket_size: usize,
    pub mean_centroid_similarity: f32,
    pub mean_isolation_score: f32,
    pub mean_stability_score: f32,
}

/// Statistics for a single bit depth analysis
#[derive(Debug, Clone)]
pub struct BitDepthStats {
    pub bits: usize,
    pub num_buckets: usize,
    pub max_bucket_size: usize,
    pub median_bucket_size: usize,
    pub min_bucket_size: usize,
    pub mean_bucket_size: f32,
    pub buckets_over_1000: usize,
    pub buckets_under_10: usize,
}

/// Analyze multiple bit depths to find optimal LSH configuration
pub fn analyze_bit_depths(
    data: &[f32],
    n_samples: usize,
    embedding_dim: usize,
    bit_range: std::ops::Range<usize>,
    seed: u64,
) -> Vec<BitDepthStats> {
    bit_range
        .map(|bits| {
            let mut classifier = DensityClassifier::new(embedding_dim, bits, seed, 3, true);
            classifier.fit_flat(data, n_samples, embedding_dim);

            let bucket_sizes = classifier.get_bucket_sizes();
            let mut unique_sizes: HashMap<u64, usize> = HashMap::new();
            for (&bid, &size) in classifier.bucket_ids.iter().zip(bucket_sizes.iter()) {
                unique_sizes.insert(bid, size as usize);
            }

            let mut sizes: Vec<usize> = unique_sizes.values().copied().collect();
            sizes.sort();

            let (max_sz, median_sz, min_sz) = if sizes.is_empty() {
                (0, 0, 0)
            } else {
                (sizes[sizes.len() - 1], sizes[sizes.len() / 2], sizes[0])
            };

            let mean_sz = if sizes.is_empty() {
                0.0
            } else {
                sizes.iter().sum::<usize>() as f32 / sizes.len() as f32
            };

            BitDepthStats {
                bits,
                num_buckets: unique_sizes.len(),
                max_bucket_size: max_sz,
                median_bucket_size: median_sz,
                min_bucket_size: min_sz,
                mean_bucket_size: mean_sz,
                buckets_over_1000: sizes.iter().filter(|&&s| s > 1000).count(),
                buckets_under_10: sizes.iter().filter(|&&s| s < 10).count(),
            }
        })
        .collect()
}

/// Select optimal bit depth based on heuristics
pub fn select_optimal_bit_depth(
    stats: &[BitDepthStats],
    target_mean_bucket: f32,
    max_bucket: usize,
) -> Option<usize> {
    // Find one that has reasonable mean bucket size and max bucket
    for s in stats {
        if s.mean_bucket_size >= target_mean_bucket && s.max_bucket_size <= max_bucket {
            return Some(s.bits);
        }
    }

    // Fallback: best trade-off
    stats
        .iter()
        .filter(|s| s.max_bucket_size <= max_bucket * 2)
        .min_by_key(|s| (s.mean_bucket_size - target_mean_bucket).abs() as usize)
        .map(|s| s.bits)
        .or_else(|| stats.first().map(|s| s.bits))
}

/// Density Classifier using PCA-based LSH
///
/// Returns raw metrics per item - no categorical classification.
pub struct DensityClassifier {
    // Configuration
    embedding_dim: usize,
    num_bits: usize,
    seed: u64,
    num_stability_seeds: usize,
    skip_isolation: bool,

    // Compute backend (CPU BLAS or GPU)
    backend: ComputeBackend,

    // PCA hyperplanes
    hyperplanes: Option<DMatrix<f32>>,

    // Per-item results
    bucket_ids: Vec<u64>,
    bucket_sizes: Vec<u32>,
    centroid_similarities: Vec<f32>,
    isolation_scores: Vec<f32>,
    stability_scores: Vec<f32>,

    // Stats
    stage1_variance: f32,
    num_buckets: usize,
    eigenvalues: Vec<f32>,

    fitted: bool,
}

impl DensityClassifier {
    /// Create new density classifier
    ///
    /// # Arguments
    /// * `embedding_dim` - Dimensionality of embeddings
    /// * `num_bits` - Bits for PCA LSH (default: 14)
    /// * `seed` - Random seed
    /// * `num_stability_seeds` - Number of seeds for stability scoring (default: 3)
    /// * `skip_isolation` - Skip isolation/stability score computation (default: false)
    pub fn new(embedding_dim: usize, num_bits: usize, seed: u64, num_stability_seeds: usize, skip_isolation: bool) -> Self {
        DensityClassifier {
            embedding_dim,
            num_bits,
            seed,
            num_stability_seeds,
            skip_isolation,
            backend: ComputeBackend::auto(),
            hyperplanes: None,
            bucket_ids: Vec::new(),
            bucket_sizes: Vec::new(),
            centroid_similarities: Vec::new(),
            isolation_scores: Vec::new(),
            stability_scores: Vec::new(),
            stage1_variance: 0.0,
            num_buckets: 0,
            eigenvalues: Vec::new(),
            fitted: false,
        }
    }

    /// Create with default parameters
    pub fn with_defaults(embedding_dim: usize) -> Self {
        Self::new(embedding_dim, 14, 31, 3, false)
    }

    /// Create with an explicit compute backend.
    pub fn with_backend(
        embedding_dim: usize,
        num_bits: usize,
        seed: u64,
        num_stability_seeds: usize,
        skip_isolation: bool,
        backend: ComputeBackend,
    ) -> Self {
        DensityClassifier {
            embedding_dim,
            num_bits,
            seed,
            num_stability_seeds,
            skip_isolation,
            backend,
            hyperplanes: None,
            bucket_ids: Vec::new(),
            bucket_sizes: Vec::new(),
            centroid_similarities: Vec::new(),
            isolation_scores: Vec::new(),
            stability_scores: Vec::new(),
            stage1_variance: 0.0,
            num_buckets: 0,
            eigenvalues: Vec::new(),
            fitted: false,
        }
    }

    /// Fit from a flat slice (zero-copy from numpy)
    pub fn fit_flat(&mut self, data: &[f32], n_samples: usize, dim: usize) -> &Self {
        if n_samples == 0 || dim != self.embedding_dim {
            self.fitted = true;
            return self;
        }

        // Normalize data into a flat buffer
        let mut normalized = vec![0.0f32; n_samples * dim];
        normalized
            .par_chunks_mut(dim)
            .enumerate()
            .for_each(|(i, out_row)| {
                let start = i * dim;
                let in_row = &data[start..start + dim];
                let norm: f32 = in_row.iter().map(|x| x * x).sum::<f32>().sqrt();
                if norm > 0.0 {
                    for (j, &val) in in_row.iter().enumerate() {
                        out_row[j] = val / norm;
                    }
                } else {
                    out_row.copy_from_slice(in_row);
                }
            });

        self.fit_internal_flat(&normalized, n_samples, dim)
    }

    /// Fit using pre-computed hyperplanes (e.g., from sklearn PCA)
    ///
    /// Skips the random-hash → centroid → PCA pipeline. Only does:
    /// normalization, hashing with provided hyperplanes, bucket metrics,
    /// and isolation/stability (if not skipped).
    pub fn fit_with_hyperplanes_flat(
        &mut self, data: &[f32], n: usize, d: usize,
        hp_data: &[f32], hp_rows: usize,
    ) -> &Self {
        if n == 0 || d != self.embedding_dim {
            self.fitted = true;
            return self;
        }

        // Normalize data
        let mut normalized = vec![0.0f32; n * d];
        normalized
            .par_chunks_mut(d)
            .enumerate()
            .for_each(|(i, out_row)| {
                let start = i * d;
                let in_row = &data[start..start + d];
                let norm: f32 = in_row.iter().map(|x| x * x).sum::<f32>().sqrt();
                if norm > 0.0 {
                    for (j, &val) in in_row.iter().enumerate() {
                        out_row[j] = val / norm;
                    }
                } else {
                    out_row.copy_from_slice(in_row);
                }
            });

        // Build DMatrix from flat hyperplane data (hp_rows × d)
        let hp = DMatrix::from_row_slice(hp_rows, d, hp_data);
        self.hyperplanes = Some(hp.clone());
        self.num_bits = hp_rows;
        self.stage1_variance = 0.0;

        // Hash with provided hyperplanes
        let hashes = self.hash_embeddings_flat(&normalized, n, d, &hp);
        self.bucket_ids = hashes.clone();

        // Bucket mapping + per-item metrics
        let mut bucket_to_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (idx, &hash) in hashes.iter().enumerate() {
            bucket_to_indices.entry(hash).or_default().push(idx);
        }
        self.num_buckets = bucket_to_indices.len();

        self.bucket_sizes = vec![0u32; n];
        self.centroid_similarities = vec![0.0f32; n];
        for (&_bucket_id, indices) in &bucket_to_indices {
            let bucket_size = indices.len() as u32;
            let centroid = self.compute_centroid_flat(&normalized, d, indices);
            for &idx in indices {
                self.bucket_sizes[idx] = bucket_size;
                let row = &normalized[idx * d..(idx + 1) * d];
                self.centroid_similarities[idx] = self.cosine_similarity(row, &centroid);
            }
        }

        // Isolation/stability
        if !self.skip_isolation {
            self.compute_isolation_scores_flat(&normalized, n, d, 10);
            self.compute_stability_scores_flat(&normalized, n, d);
        } else {
            self.isolation_scores = vec![0.0f32; n];
            self.stability_scores = vec![1.0f32; n];
        }

        self.fitted = true;
        self
    }

    /// Fit using PCA computed directly on the full dataset
    ///
    /// Skips the random-hash → centroid approximation and runs PCA on all
    /// normalized data points. Gives better hyperplane quality than
    /// centroid-PCA, especially in moderate dimensions (e.g., 128d).
    pub fn fit_raw_pca_flat(&mut self, data: &[f32], n: usize, d: usize) -> &Self {
        if n == 0 || d != self.embedding_dim {
            self.fitted = true;
            return self;
        }

        // Normalize data
        let mut normalized = vec![0.0f32; n * d];
        normalized
            .par_chunks_mut(d)
            .enumerate()
            .for_each(|(i, out_row)| {
                let start = i * d;
                let in_row = &data[start..start + d];
                let norm: f32 = in_row.iter().map(|x| x * x).sum::<f32>().sqrt();
                if norm > 0.0 {
                    for (j, &val) in in_row.iter().enumerate() {
                        out_row[j] = val / norm;
                    }
                } else {
                    out_row.copy_from_slice(in_row);
                }
            });

        // PCA directly on full normalized data
        let (pca_hp, variance, ev) = if n > self.num_bits {
            self.compute_pca_flat(&normalized, n, d, self.num_bits)
        } else {
            let mut rng = ChaCha8Rng::seed_from_u64(self.seed);
            (self.generate_random_hyperplanes(&mut rng, self.num_bits), 0.0, Vec::new())
        };
        self.hyperplanes = Some(pca_hp.clone());
        self.stage1_variance = variance;
        self.eigenvalues = ev;

        // Hash with PCA hyperplanes
        let hashes = self.hash_embeddings_flat(&normalized, n, d, &pca_hp);
        self.bucket_ids = hashes.clone();

        // Bucket mapping + per-item metrics
        let mut bucket_to_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (idx, &hash) in hashes.iter().enumerate() {
            bucket_to_indices.entry(hash).or_default().push(idx);
        }
        self.num_buckets = bucket_to_indices.len();

        self.bucket_sizes = vec![0u32; n];
        self.centroid_similarities = vec![0.0f32; n];
        for (&_bucket_id, indices) in &bucket_to_indices {
            let bucket_size = indices.len() as u32;
            let centroid = self.compute_centroid_flat(&normalized, d, indices);
            for &idx in indices {
                self.bucket_sizes[idx] = bucket_size;
                let row = &normalized[idx * d..(idx + 1) * d];
                self.centroid_similarities[idx] = self.cosine_similarity(row, &centroid);
            }
        }

        // Isolation/stability
        if !self.skip_isolation {
            self.compute_isolation_scores_flat(&normalized, n, d, 10);
            self.compute_stability_scores_flat(&normalized, n, d);
        } else {
            self.isolation_scores = vec![0.0f32; n];
            self.stability_scores = vec![1.0f32; n];
        }

        self.fitted = true;
        self
    }

    /// Fit the density classifier on embeddings
    pub fn fit(&mut self, embeddings: &[Vec<f32>]) -> &Self {
        let n = embeddings.len();
        if n == 0 {
            self.fitted = true;
            return self;
        }

        // Normalize embeddings (parallel)
        let embeddings: Vec<Vec<f32>> = embeddings
            .par_iter()
            .map(|e| {
                let norm: f32 = e.iter().map(|x| x * x).sum::<f32>().sqrt();
                if norm > 0.0 {
                    e.iter().map(|x| x / norm).collect()
                } else {
                    e.clone()
                }
            })
            .collect();

        // Convert to flat
        let dim = embeddings[0].len();
        let mut flat = vec![0.0f32; n * dim];
        for (i, emb) in embeddings.iter().enumerate() {
            flat[i * dim..(i + 1) * dim].copy_from_slice(emb);
        }

        self.fit_internal_flat(&flat, n, dim)
    }

    /// Ensemble fit: use consensus voting across multiple seeds
    ///
    /// Runs PCA-LSH with multiple seeds and assigns each item to
    /// its most common bucket (majority vote). Reduces boundary noise.
    ///
    /// Returns number of items where consensus differed from primary seed.
    pub fn fit_ensemble_flat(&mut self, data: &[f32], n: usize, d: usize, num_seeds: usize) -> usize {
        if n == 0 || d != self.embedding_dim || num_seeds == 0 {
            self.fitted = true;
            return 0;
        }

        // Normalize data
        let mut normalized = vec![0.0f32; n * d];
        normalized
            .par_chunks_mut(d)
            .enumerate()
            .for_each(|(i, out_row)| {
                let start = i * d;
                let in_row = &data[start..start + d];
                let norm: f32 = in_row.iter().map(|x| x * x).sum::<f32>().sqrt();
                if norm > 0.0 {
                    for (j, &val) in in_row.iter().enumerate() {
                        out_row[j] = val / norm;
                    }
                } else {
                    out_row.copy_from_slice(in_row);
                }
            });

        // First, do standard PCA-LSH to get base hyperplanes
        let mut rng = ChaCha8Rng::seed_from_u64(self.seed);
        let random_hp = self.generate_random_hyperplanes(&mut rng, self.num_bits);
        let random_hashes = self.hash_embeddings_flat(&normalized, n, d, &random_hp);
        let centroids = self.compute_all_centroids_flat(&normalized, n, d, &random_hashes, 2);
        let num_centroids = centroids.len() / d;

        let base_hp = if num_centroids > self.num_bits {
            let (hp, variance, _) = self.compute_pca_flat(&centroids, num_centroids, d, self.num_bits);
            self.stage1_variance = variance;
            hp
        } else {
            random_hp
        };

        // Run with base hyperplanes + perturbations
        let mut all_hashes: Vec<Vec<u64>> = Vec::with_capacity(num_seeds);

        for seed_idx in 0..num_seeds {
            let perturbed_hp = if seed_idx == 0 {
                // First seed uses exact base hyperplanes
                base_hp.clone()
            } else {
                // Perturb hyperplanes slightly (1% noise)
                let seed = self.seed + (seed_idx as u64 * 1000);
                let mut rng = ChaCha8Rng::seed_from_u64(seed);
                let mut hp = base_hp.clone();

                for i in 0..self.num_bits {
                    for j in 0..d {
                        hp[(i, j)] += rng.gen_range(-0.01f32..0.01f32);
                    }
                    // Re-normalize
                    let norm: f32 = (0..d).map(|j| hp[(i, j)].powi(2)).sum::<f32>().sqrt();
                    if norm > 0.0 {
                        for j in 0..d {
                            hp[(i, j)] /= norm;
                        }
                    }
                }
                hp
            };

            let hashes = self.hash_embeddings_flat(&normalized, n, d, &perturbed_hp);
            all_hashes.push(hashes);
        }

        let primary_hyperplanes = Some(base_hp);

        // Majority vote for each item
        let consensus_hashes: Vec<u64> = (0..n)
            .into_par_iter()
            .map(|i| {
                // Count votes for each bucket
                let mut votes: HashMap<u64, usize> = HashMap::new();
                for seed_hashes in &all_hashes {
                    *votes.entry(seed_hashes[i]).or_default() += 1;
                }

                // Find bucket with most votes
                votes
                    .into_iter()
                    .max_by_key(|(_, count)| *count)
                    .map(|(bucket, _)| bucket)
                    .unwrap_or(all_hashes[0][i])
            })
            .collect();

        // Count how many items differ from primary seed
        let changed = consensus_hashes
            .iter()
            .zip(all_hashes[0].iter())
            .filter(|(consensus, primary)| consensus != primary)
            .count();

        // Store results
        self.hyperplanes = primary_hyperplanes;
        self.bucket_ids = consensus_hashes;

        // Build bucket mapping and compute metrics
        let mut bucket_to_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (idx, &hash) in self.bucket_ids.iter().enumerate() {
            bucket_to_indices.entry(hash).or_default().push(idx);
        }
        self.num_buckets = bucket_to_indices.len();

        // Compute per-item metrics
        self.bucket_sizes = vec![0u32; n];
        self.centroid_similarities = vec![0.0f32; n];

        for (&_bucket_id, indices) in &bucket_to_indices {
            let bucket_size = indices.len() as u32;
            let centroid = self.compute_centroid_flat(&normalized, d, indices);

            for &idx in indices {
                self.bucket_sizes[idx] = bucket_size;
                let row = &normalized[idx * d..(idx + 1) * d];
                self.centroid_similarities[idx] = self.cosine_similarity(row, &centroid);
            }
        }

        // Compute stability using the ensemble hashes we already have
        self.stability_scores = (0..n)
            .into_par_iter()
            .map(|i| {
                let mut buckets: Vec<u64> = all_hashes.iter().map(|h| h[i]).collect();
                buckets.sort();
                buckets.dedup();
                let unique_count = buckets.len();

                if num_seeds <= 1 {
                    1.0
                } else {
                    1.0 - (unique_count as f32 - 1.0) / (num_seeds as f32 - 1.0)
                }
            })
            .collect();

        // Compute isolation scores
        if !self.skip_isolation {
            self.compute_isolation_scores_flat(&normalized, n, d, 10);
        } else {
            self.isolation_scores = vec![0.0f32; n];
        }

        self.fitted = true;
        changed
    }

    /// Iterative fit: refine hyperplanes by repeating PCA on centroids
    ///
    /// Each iteration:
    /// 1. Hash items with current hyperplanes
    /// 2. Compute bucket centroids
    /// 3. PCA on centroids → new hyperplanes
    /// 4. Repeat until convergence or max_iterations
    ///
    /// Returns (self, iterations_run, items_changed_last_iter)
    pub fn fit_iterative_flat(&mut self, data: &[f32], n: usize, d: usize, max_iterations: usize) -> (usize, usize) {
        if n == 0 || d != self.embedding_dim {
            self.fitted = true;
            return (0, 0);
        }

        // Normalize data
        let mut normalized = vec![0.0f32; n * d];
        normalized
            .par_chunks_mut(d)
            .enumerate()
            .for_each(|(i, out_row)| {
                let start = i * d;
                let in_row = &data[start..start + d];
                let norm: f32 = in_row.iter().map(|x| x * x).sum::<f32>().sqrt();
                if norm > 0.0 {
                    for (j, &val) in in_row.iter().enumerate() {
                        out_row[j] = val / norm;
                    }
                } else {
                    out_row.copy_from_slice(in_row);
                }
            });

        // Initial random hash
        let mut rng = ChaCha8Rng::seed_from_u64(self.seed);
        let mut current_hp = self.generate_random_hyperplanes(&mut rng, self.num_bits);
        let mut current_hashes = self.hash_embeddings_flat(&normalized, n, d, &current_hp);

        let mut iterations = 0;
        let mut last_changed = n; // First iteration "changes" everything

        for iter in 0..max_iterations {
            // Compute centroids from current buckets
            let centroids = self.compute_all_centroids_flat(&normalized, n, d, &current_hashes, 2);
            let num_centroids = centroids.len() / d;

            if num_centroids <= self.num_bits {
                // Not enough centroids for PCA, stop iterating
                break;
            }

            // PCA on centroids → new hyperplanes
            let (pca_hp, variance, _) = self.compute_pca_flat(&centroids, num_centroids, d, self.num_bits);
            self.stage1_variance = variance;

            // Damped update: blend new hyperplanes with old (reduces oscillation)
            // First align signs (eigenvectors have arbitrary sign)
            let alpha = 0.5f32;
            let mut new_hp = DMatrix::<f32>::zeros(self.num_bits, d);
            for i in 0..self.num_bits {
                // Check if PCA hyperplane points same direction as current
                let dot: f32 = (0..d).map(|j| pca_hp[(i, j)] * current_hp[(i, j)]).sum();
                let sign = if dot >= 0.0 { 1.0f32 } else { -1.0f32 };

                // Blend with aligned sign
                for j in 0..d {
                    new_hp[(i, j)] = alpha * sign * pca_hp[(i, j)] + (1.0 - alpha) * current_hp[(i, j)];
                }
                // Re-normalize row
                let norm: f32 = (0..d).map(|j| new_hp[(i, j)].powi(2)).sum::<f32>().sqrt();
                if norm > 0.0 {
                    for j in 0..d {
                        new_hp[(i, j)] /= norm;
                    }
                }
            }

            // Re-hash with blended hyperplanes
            let new_hashes = self.hash_embeddings_flat(&normalized, n, d, &new_hp);

            // Count how many items changed buckets
            let changed: usize = current_hashes
                .iter()
                .zip(new_hashes.iter())
                .filter(|(old, new)| old != new)
                .count();

            current_hp = new_hp;
            current_hashes = new_hashes;
            iterations = iter + 1;
            last_changed = changed;

            // Convergence check: if very few items moved, stop
            if changed < n / 100 {
                // Less than 1% changed
                break;
            }
        }

        // Store final hyperplanes and hashes
        self.hyperplanes = Some(current_hp);
        self.bucket_ids = current_hashes;

        // Build bucket mapping and compute metrics
        let mut bucket_to_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (idx, &hash) in self.bucket_ids.iter().enumerate() {
            bucket_to_indices.entry(hash).or_default().push(idx);
        }
        self.num_buckets = bucket_to_indices.len();

        // Compute per-item metrics
        self.bucket_sizes = vec![0u32; n];
        self.centroid_similarities = vec![0.0f32; n];

        for (&_bucket_id, indices) in &bucket_to_indices {
            let bucket_size = indices.len() as u32;
            let centroid = self.compute_centroid_flat(&normalized, d, indices);

            for &idx in indices {
                self.bucket_sizes[idx] = bucket_size;
                let row = &normalized[idx * d..(idx + 1) * d];
                self.centroid_similarities[idx] = self.cosine_similarity(row, &centroid);
            }
        }

        // Compute isolation and stability scores
        if !self.skip_isolation {
            self.compute_isolation_scores_flat(&normalized, n, d, 10);
            self.compute_stability_scores_flat(&normalized, n, d);
        } else {
            self.isolation_scores = vec![0.0f32; n];
            self.stability_scores = vec![1.0f32; n];
        }

        self.fitted = true;
        (iterations, last_changed)
    }

    /// Internal fit implementation using flat slices
    fn fit_internal_flat(&mut self, data: &[f32], n: usize, d: usize) -> &Self {
        // Step 1: Random hash to get initial buckets for centroid computation
        let mut rng = ChaCha8Rng::seed_from_u64(self.seed);
        let random_hp = self.generate_random_hyperplanes(&mut rng, self.num_bits);
        let random_hashes = self.hash_embeddings_flat(data, n, d, &random_hp);

        // Step 2: Compute bucket centroids
        let centroids = self.compute_all_centroids_flat(data, n, d, &random_hashes, 2);
        let num_centroids = centroids.len() / d;

        // Step 3: PCA on centroids
        let (pca_hp, variance, ev) = if num_centroids > self.num_bits {
            self.compute_pca_flat(&centroids, num_centroids, d, self.num_bits)
        } else {
            (random_hp.clone(), 0.0, Vec::new())
        };
        self.hyperplanes = Some(pca_hp.clone());
        self.stage1_variance = variance;
        self.eigenvalues = ev;

        // Step 4: Re-hash with PCA hyperplanes
        let hashes = self.hash_embeddings_flat(data, n, d, &pca_hp);
        self.bucket_ids = hashes.clone();

        // Build bucket mapping
        let mut bucket_to_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (idx, &hash) in hashes.iter().enumerate() {
            bucket_to_indices.entry(hash).or_default().push(idx);
        }
        self.num_buckets = bucket_to_indices.len();

        // Step 5: Compute per-item metrics
        self.bucket_sizes = vec![0u32; n];
        self.centroid_similarities = vec![0.0f32; n];

        for (&_bucket_id, indices) in &bucket_to_indices {
            let bucket_size = indices.len() as u32;

            // Compute centroid
            let centroid = self.compute_centroid_flat(data, d, indices);

            // Set bucket size and centroid similarity for each item
            for &idx in indices {
                self.bucket_sizes[idx] = bucket_size;
                let row = &data[idx * d..(idx + 1) * d];
                self.centroid_similarities[idx] = self.cosine_similarity(row, &centroid);
            }
        }

        // Step 6: Compute isolation scores
        // Step 7: Compute stability scores across multiple seeds
        if !self.skip_isolation {
            self.compute_isolation_scores_flat(data, n, d, 10);
            self.compute_stability_scores_flat(data, n, d);
        } else {
            self.isolation_scores = vec![0.0f32; n];
            self.stability_scores = vec![1.0f32; n];
        }

        self.fitted = true;
        self
    }

    /// Compute isolation score for each item
    /// isolation = mean(top_k similarities) - median(all similarities)
    fn compute_isolation_scores_flat(&mut self, data: &[f32], n: usize, d: usize, k: usize) {
        self.isolation_scores = vec![0.0f32; n];

        // For efficiency, sample the corpus for median computation
        let sample_size = 1000.min(n);
        let mut rng = ChaCha8Rng::seed_from_u64(self.seed + 12345);
        let sample_indices: Vec<usize> = if n <= sample_size {
            (0..n).collect()
        } else {
            let mut indices: Vec<usize> = (0..n).collect();
            for i in 0..sample_size {
                let j = rng.gen_range(i..n);
                indices.swap(i, j);
            }
            indices[..sample_size].to_vec()
        };

        // Compute isolation scores in parallel
        self.isolation_scores = (0..n)
            .into_par_iter()
            .map(|i| {
                let row_i = &data[i * d..(i + 1) * d];

                // Compute similarities to sample
                let mut sims: Vec<f32> = sample_indices
                    .iter()
                    .filter(|&&j| j != i)
                    .map(|&j| {
                        let row_j = &data[j * d..(j + 1) * d];
                        self.cosine_similarity(row_i, row_j)
                    })
                    .collect();

                if sims.is_empty() {
                    return 0.0;
                }

                // Sort to get top-k and median
                sims.sort_by(|a, b| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal));

                let top_k_mean: f32 = sims.iter().take(k).sum::<f32>() / k.min(sims.len()) as f32;
                let median = sims[sims.len() / 2];

                top_k_mean - median
            })
            .collect();
    }

    /// Compute stability score for each item across multiple seeds
    /// stability = 1 - (unique_buckets - 1) / (num_seeds - 1)
    /// 1.0 = same bucket in all seeds, 0.0 = different bucket in each seed
    fn compute_stability_scores_flat(&mut self, data: &[f32], n: usize, d: usize) {
        if self.num_stability_seeds <= 1 {
            // No stability computation with 0 or 1 seed
            self.stability_scores = vec![1.0f32; n];
            return;
        }

        // Compute bucket assignments for each seed
        // Start offset at 1 to avoid correlation with main seed's hyperplanes
        let mut all_hashes: Vec<Vec<u64>> = Vec::with_capacity(self.num_stability_seeds);

        for seed_idx in 0..self.num_stability_seeds {
            let seed = self.seed + ((seed_idx + 1) as u64 * 1000);

            // Generate random hyperplanes for this seed
            let mut rng = ChaCha8Rng::seed_from_u64(seed);
            let random_hp = self.generate_random_hyperplanes(&mut rng, self.num_bits);
            let random_hashes = self.hash_embeddings_flat(data, n, d, &random_hp);

            // Compute centroids for PCA
            let centroids = self.compute_all_centroids_flat(data, n, d, &random_hashes, 2);
            let num_centroids = centroids.len() / d;

            // PCA on centroids
            let pca_hp = if num_centroids > self.num_bits {
                let (hp, _, _) = self.compute_pca_flat(&centroids, num_centroids, d, self.num_bits);
                hp
            } else {
                random_hp
            };

            // Hash with PCA hyperplanes
            let hashes = self.hash_embeddings_flat(data, n, d, &pca_hp);
            all_hashes.push(hashes);
        }

        // Compute stability for each item
        self.stability_scores = (0..n)
            .into_par_iter()
            .map(|i| {
                // Count unique buckets for this item
                let mut buckets: Vec<u64> = all_hashes.iter().map(|h| h[i]).collect();
                buckets.sort();
                buckets.dedup();
                let unique_count = buckets.len();

                // stability = 1 - (unique - 1) / (num_seeds - 1)
                // 1 unique -> 1.0, num_seeds unique -> 0.0
                if self.num_stability_seeds == 1 {
                    1.0
                } else {
                    1.0 - (unique_count as f32 - 1.0) / (self.num_stability_seeds as f32 - 1.0)
                }
            })
            .collect();
    }

    /// Generate random hyperplanes
    fn generate_random_hyperplanes(&self, rng: &mut ChaCha8Rng, num_bits: usize) -> DMatrix<f32> {
        let mut data = Vec::with_capacity(num_bits * self.embedding_dim);
        for _ in 0..(num_bits * self.embedding_dim) {
            data.push(rng.gen_range(-1.0f32..1.0f32));
        }

        let mut hp = DMatrix::from_vec(num_bits, self.embedding_dim, data);

        // Normalize each row
        for i in 0..num_bits {
            let row = hp.row(i);
            let norm: f32 = row.iter().map(|x| x * x).sum::<f32>().sqrt();
            if norm > 0.0 {
                for j in 0..self.embedding_dim {
                    hp[(i, j)] /= norm;
                }
            }
        }

        hp
    }

    /// Hash embeddings from flat slice — delegates to ComputeBackend.
    fn hash_embeddings_flat(&self, data: &[f32], n: usize, d: usize, hyperplanes: &DMatrix<f32>) -> Vec<u64> {
        let num_bits = hyperplanes.nrows();

        // Flatten DMatrix to row-major (num_bits × d) for the backend
        let mut hp_flat = vec![0.0f32; num_bits * d];
        for i in 0..num_bits {
            for j in 0..d {
                hp_flat[i * d + j] = hyperplanes[(i, j)];
            }
        }

        self.backend.hash_embeddings(data, n, d, &hp_flat, num_bits)
    }

    /// Compute all centroids from flat data
    fn compute_all_centroids_flat(&self, data: &[f32], _n: usize, d: usize, hashes: &[u64], min_size: usize) -> Vec<f32> {
        let mut bucket_sizes: HashMap<u64, usize> = HashMap::new();
        for &hash in hashes {
            *bucket_sizes.entry(hash).or_default() += 1;
        }

        let valid_buckets: Vec<u64> = bucket_sizes
            .iter()
            .filter(|(_, &size)| size >= min_size)
            .map(|(&hash, _)| hash)
            .collect();

        let bucket_to_idx: HashMap<u64, usize> = valid_buckets
            .iter()
            .enumerate()
            .map(|(idx, &hash)| (hash, idx))
            .collect();

        let num_centroids = valid_buckets.len();
        if num_centroids == 0 {
            return Vec::new();
        }

        let mut sums = vec![0.0f32; num_centroids * d];
        let mut counts = vec![0usize; num_centroids];

        for (i, &hash) in hashes.iter().enumerate() {
            if let Some(&centroid_idx) = bucket_to_idx.get(&hash) {
                let row_start = i * d;
                let sum_start = centroid_idx * d;
                for j in 0..d {
                    sums[sum_start + j] += data[row_start + j];
                }
                counts[centroid_idx] += 1;
            }
        }

        // Normalize centroids
        sums.par_chunks_mut(d)
            .zip(counts.par_iter())
            .for_each(|(sum_chunk, &count)| {
                if count > 0 {
                    let n_f = count as f32;
                    for val in sum_chunk.iter_mut() {
                        *val /= n_f;
                    }
                    let norm: f32 = sum_chunk.iter().map(|x| x * x).sum::<f32>().sqrt();
                    if norm > 0.0 {
                        for val in sum_chunk.iter_mut() {
                            *val /= norm;
                        }
                    }
                }
            });

        sums.chunks(d)
            .zip(counts.iter())
            .filter(|(_, &count)| count > 0)
            .flat_map(|(chunk, _)| chunk.iter().copied())
            .collect()
    }

    /// Compute centroid from flat data for specific indices
    fn compute_centroid_flat(&self, data: &[f32], d: usize, indices: &[usize]) -> Vec<f32> {
        if indices.is_empty() {
            return vec![0.0f32; d];
        }

        let mut centroid = vec![0.0f32; d];
        for &idx in indices {
            let row_start = idx * d;
            for j in 0..d {
                centroid[j] += data[row_start + j];
            }
        }

        let n = indices.len() as f32;
        for val in &mut centroid {
            *val /= n;
        }

        let norm: f32 = centroid.iter().map(|x| x * x).sum::<f32>().sqrt();
        if norm > 0.0 {
            for val in &mut centroid {
                *val /= norm;
            }
        }
        centroid
    }

    /// Compute PCA hyperplanes from flat data.
    /// Returns (hyperplanes, variance_ratio, eigenvalues).
    fn compute_pca_flat(&self, data: &[f32], n: usize, d: usize, num_components: usize) -> (DMatrix<f32>, f32, Vec<f32>) {
        if n < 2 || d == 0 {
            return (
                self.generate_random_hyperplanes(&mut ChaCha8Rng::seed_from_u64(self.seed), num_components),
                0.0,
                Vec::new(),
            );
        }

        let data_view = ArrayView2::from_shape((n, d), data).unwrap();
        let mean = data_view.mean_axis(ndarray::Axis(0)).unwrap();
        let centered = &data_view - &mean;
        let xtx = centered.t().dot(&centered);

        let k = num_components.min(d).min(n - 1);
        let (components, eigenvalues) = self.power_iteration_pca_ndarray(&xtx, k, d, 20);

        let total_variance = if !eigenvalues.is_empty() {
            let sum: f32 = eigenvalues.iter().sum();
            if sum > 0.0 {
                eigenvalues.iter().sum::<f32>() / (sum + 0.001)
            } else {
                0.0
            }
        } else {
            0.0
        };

        (components, total_variance.min(1.0), eigenvalues)
    }

    /// Power iteration PCA — delegates to ComputeBackend.
    fn power_iteration_pca_ndarray(
        &self,
        xtx: &Array2<f32>,
        k: usize,
        d: usize,
        max_iter: usize,
    ) -> (DMatrix<f32>, Vec<f32>) {
        // Flatten ndarray to row-major slice for the backend
        let xtx_flat = xtx.as_slice().expect("xtx must be contiguous row-major");

        let (comp_flat, eigenvalues) = self.backend.power_iteration(xtx_flat, d, k, max_iter, self.seed);

        // Pack flat components (k × d) into DMatrix
        let components = DMatrix::from_row_slice(k, d, &comp_flat);

        (components, eigenvalues)
    }

    /// Cosine similarity between two vectors
    fn cosine_similarity(&self, a: &[f32], b: &[f32]) -> f32 {
        a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
    }

    // =========================================================================
    // Public Accessors
    // =========================================================================

    /// Get bucket IDs for all items
    pub fn get_bucket_ids(&self) -> &[u64] {
        &self.bucket_ids
    }

    /// Get bucket ID for a single item
    pub fn get_bucket_id(&self, idx: usize) -> Option<u64> {
        self.bucket_ids.get(idx).copied()
    }

    /// Get bucket sizes for all items
    pub fn get_bucket_sizes(&self) -> &[u32] {
        &self.bucket_sizes
    }

    /// Get centroid similarities for all items
    pub fn get_centroid_similarities(&self) -> &[f32] {
        &self.centroid_similarities
    }

    /// Get isolation scores for all items
    pub fn get_isolation_scores(&self) -> &[f32] {
        &self.isolation_scores
    }

    /// Get stability scores for all items (0-1, higher = more stable)
    pub fn get_stability_scores(&self) -> &[f32] {
        &self.stability_scores
    }

    /// Get classification report
    pub fn report(&self) -> DensityReport {
        let mean_bucket = if self.bucket_sizes.is_empty() {
            0.0
        } else {
            self.bucket_sizes.iter().map(|&x| x as f32).sum::<f32>() / self.bucket_sizes.len() as f32
        };

        let mut sorted_sizes: Vec<u32> = self.bucket_sizes.clone();
        sorted_sizes.sort();
        let median_bucket = if sorted_sizes.is_empty() {
            0
        } else {
            sorted_sizes[sorted_sizes.len() / 2] as usize
        };
        let max_bucket = sorted_sizes.last().copied().unwrap_or(0) as usize;

        let mean_centroid_sim = if self.centroid_similarities.is_empty() {
            0.0
        } else {
            self.centroid_similarities.iter().sum::<f32>() / self.centroid_similarities.len() as f32
        };

        let mean_isolation = if self.isolation_scores.is_empty() {
            0.0
        } else {
            self.isolation_scores.iter().sum::<f32>() / self.isolation_scores.len() as f32
        };

        let mean_stability = if self.stability_scores.is_empty() {
            0.0
        } else {
            self.stability_scores.iter().sum::<f32>() / self.stability_scores.len() as f32
        };

        DensityReport {
            corpus_size: self.bucket_ids.len(),
            num_buckets: self.num_buckets,
            stage1_variance_explained: self.stage1_variance,
            mean_bucket_size: mean_bucket,
            median_bucket_size: median_bucket,
            max_bucket_size: max_bucket,
            mean_centroid_similarity: mean_centroid_sim,
            mean_isolation_score: mean_isolation,
            mean_stability_score: mean_stability,
        }
    }

    /// Get PCA hyperplanes
    pub fn get_hyperplanes(&self) -> Option<&DMatrix<f32>> {
        self.hyperplanes.as_ref()
    }

    /// Get PCA eigenvalues (per-component variance from power iteration)
    pub fn get_eigenvalues(&self) -> &[f32] {
        &self.eigenvalues
    }

    /// Check if fitted
    pub fn is_fitted(&self) -> bool {
        self.fitted
    }

    /// Get embedding dimension
    pub fn embedding_dim(&self) -> usize {
        self.embedding_dim
    }

    /// Adaptive Voronoi refinement: search radius based on item confidence
    ///
    /// Low-confidence items (low centroid_similarity) get wider search.
    /// High-confidence items get narrow search or skip entirely.
    ///
    /// Returns the number of items reassigned.
    pub fn voronoi_refine_adaptive_flat(
        &mut self,
        data: &[f32],
        n: usize,
        d: usize,
    ) -> usize {
        if !self.fitted || n == 0 {
            return 0;
        }

        // Determine Hamming radius per item based on centroid similarity
        // sim < 0.5 → 3, sim 0.5-0.7 → 2, sim > 0.7 → 1
        let item_hamming: Vec<u32> = self.centroid_similarities
            .iter()
            .map(|&sim| {
                if sim < 0.5 { 3 }
                else if sim < 0.7 { 2 }
                else { 1 }
            })
            .collect();

        // Get unique buckets
        let unique_buckets: Vec<u64> = {
            let mut buckets: Vec<u64> = self.bucket_ids.clone();
            buckets.sort();
            buckets.dedup();
            buckets
        };
        let num_buckets = unique_buckets.len();
        if num_buckets <= 1 {
            return 0;
        }

        let bucket_to_idx: HashMap<u64, usize> = unique_buckets
            .iter()
            .enumerate()
            .map(|(i, &b)| (b, i))
            .collect();

        // Build bucket -> item indices
        let mut bucket_item_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (i, &bid) in self.bucket_ids.iter().enumerate() {
            bucket_item_indices.entry(bid).or_default().push(i);
        }

        // Compute centroids
        let mut centroids = Array2::<f32>::zeros((num_buckets, d));
        for (bucket_idx, &bucket_id) in unique_buckets.iter().enumerate() {
            if let Some(indices) = bucket_item_indices.get(&bucket_id) {
                let centroid = self.compute_centroid_flat(data, d, indices);
                for (j, &val) in centroid.iter().enumerate() {
                    centroids[[bucket_idx, j]] = val;
                }
            }
        }

        // Precompute neighbor sets for each Hamming distance
        let neighbors_by_hamming: Vec<Vec<Vec<usize>>> = (0..=3)
            .map(|max_h| {
                unique_buckets
                    .iter()
                    .map(|&bid| {
                        unique_buckets
                            .iter()
                            .enumerate()
                            .filter_map(|(idx, &other)| {
                                if (bid ^ other).count_ones() <= max_h {
                                    Some(idx)
                                } else {
                                    None
                                }
                            })
                            .collect()
                    })
                    .collect()
            })
            .collect();

        // Process items grouped by (bucket, hamming_radius) for batched BLAS
        use std::sync::atomic::{AtomicU64, Ordering};
        let new_ids: Vec<AtomicU64> = (0..n).map(|_| AtomicU64::new(0)).collect();

        // Group items by bucket and max hamming radius
        let mut bucket_hamming_groups: HashMap<(u64, u32), Vec<usize>> = HashMap::new();
        for (i, (&bid, &h)) in self.bucket_ids.iter().zip(item_hamming.iter()).enumerate() {
            bucket_hamming_groups.entry((bid, h)).or_default().push(i);
        }

        // Process groups in parallel
        let groups: Vec<_> = bucket_hamming_groups.into_iter().collect();
        groups.par_iter().for_each(|((bucket_id, hamming_radius), item_indices)| {
            let bucket_idx = bucket_to_idx[bucket_id];
            let neighbors = &neighbors_by_hamming[*hamming_radius as usize][bucket_idx];
            let num_neighbors = neighbors.len();

            // Build neighbor centroid matrix
            let mut neighbor_centroids = Array2::<f32>::zeros((num_neighbors, d));
            for (local_idx, &global_idx) in neighbors.iter().enumerate() {
                for j in 0..d {
                    neighbor_centroids[[local_idx, j]] = centroids[[global_idx, j]];
                }
            }

            // Build item matrix
            let num_items = item_indices.len();
            let mut item_matrix = Array2::<f32>::zeros((num_items, d));
            for (local_i, &global_i) in item_indices.iter().enumerate() {
                for j in 0..d {
                    item_matrix[[local_i, j]] = data[global_i * d + j];
                }
            }

            // BLAS matmul
            let sims = item_matrix.dot(&neighbor_centroids.t());

            // Argmax
            for (local_i, &global_i) in item_indices.iter().enumerate() {
                let row = sims.row(local_i);
                let best_local = row.iter()
                    .enumerate()
                    .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
                    .map(|(idx, _)| idx)
                    .unwrap_or(0);
                new_ids[global_i].store(unique_buckets[neighbors[best_local]], Ordering::Relaxed);
            }
        });

        let new_bucket_ids: Vec<u64> = new_ids.into_iter().map(|a| a.into_inner()).collect();

        // Count reassignments and update
        let reassigned = self.bucket_ids
            .iter()
            .zip(new_bucket_ids.iter())
            .filter(|(old, new)| old != new)
            .count();

        self.bucket_ids = new_bucket_ids;

        // Recompute metrics
        let mut new_bucket_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (i, &bid) in self.bucket_ids.iter().enumerate() {
            new_bucket_indices.entry(bid).or_default().push(i);
        }
        self.num_buckets = new_bucket_indices.len();

        self.bucket_sizes = vec![0u32; n];
        self.centroid_similarities = vec![0.0f32; n];

        for (&_bucket_id, indices) in &new_bucket_indices {
            let bucket_size = indices.len() as u32;
            let centroid = self.compute_centroid_flat(data, d, indices);
            for &idx in indices {
                self.bucket_sizes[idx] = bucket_size;
                let row = &data[idx * d..(idx + 1) * d];
                self.centroid_similarities[idx] = self.cosine_similarity(row, &centroid);
            }
        }

        reassigned
    }

    /// Hierarchical LSH: coarse global hash + fine region-specific hash
    ///
    /// Stage 1 (Coarse): Global PCA on all items → coarse_bits hyperplanes
    /// Stage 2 (Fine): Per-region PCA on items in each coarse bucket → fine_bits hyperplanes
    ///
    /// Final bucket ID combines both: (coarse_id << fine_bits) | fine_id
    ///
    /// This captures both global structure (coarse) and local structure (fine),
    /// improving centroid similarity and NN recall over flat LSH.
    ///
    /// # Arguments
    /// * `data` - Flat embedding data (n * d elements)
    /// * `n` - Number of samples
    /// * `d` - Embedding dimension
    /// * `coarse_bits` - Bits for global hash (default: 6)
    /// * `fine_bits` - Bits for per-region hash (default: 4)
    ///
    /// # Returns
    /// Self for chaining
    pub fn fit_hierarchical_flat(
        &mut self,
        data: &[f32],
        n: usize,
        d: usize,
        coarse_bits: usize,
        fine_bits: usize,
    ) -> &Self {
        if n == 0 || d != self.embedding_dim {
            self.fitted = true;
            return self;
        }

        // Normalize data into a flat buffer
        let mut normalized = vec![0.0f32; n * d];
        normalized
            .par_chunks_mut(d)
            .enumerate()
            .for_each(|(i, out_row)| {
                let start = i * d;
                let in_row = &data[start..start + d];
                let norm: f32 = in_row.iter().map(|x| x * x).sum::<f32>().sqrt();
                if norm > 0.0 {
                    for (j, &val) in in_row.iter().enumerate() {
                        out_row[j] = val / norm;
                    }
                } else {
                    out_row.copy_from_slice(in_row);
                }
            });

        // =========================================================================
        // Stage 1: Coarse hash with global PCA on RAW DATA (not centroids)
        // =========================================================================

        // PCA directly on embeddings - captures global variance directions
        let (coarse_hp, variance, _) = self.compute_pca_flat(&normalized, n, d, coarse_bits);
        self.stage1_variance = variance;

        // Hash all items with coarse hyperplanes
        let coarse_hashes = self.hash_embeddings_flat(&normalized, n, d, &coarse_hp);

        // Build coarse bucket -> item indices mapping
        let mut coarse_bucket_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (i, &hash) in coarse_hashes.iter().enumerate() {
            coarse_bucket_indices.entry(hash).or_default().push(i);
        }

        // =========================================================================
        // Stage 2: Fine hash with per-region PCA on RAW DATA
        // =========================================================================

        // For each coarse bucket, compute region-specific PCA and fine hash
        let mut final_bucket_ids = vec![0u64; n];

        // Process coarse buckets in parallel
        let coarse_buckets: Vec<(u64, Vec<usize>)> = coarse_bucket_indices.into_iter().collect();

        let fine_results: Vec<(Vec<usize>, Vec<u64>)> = coarse_buckets
            .par_iter()
            .map(|(coarse_id, indices)| {
                let num_items = indices.len();

                // Extract region items
                let mut region_data = vec![0.0f32; num_items * d];
                for (local_i, &global_i) in indices.iter().enumerate() {
                    region_data[local_i * d..(local_i + 1) * d]
                        .copy_from_slice(&normalized[global_i * d..(global_i + 1) * d]);
                }

                // Compute fine hyperplanes for this region - PCA directly on data
                let fine_hp = if num_items > fine_bits + 1 {
                    let (hp, _, _) = self.compute_pca_flat(&region_data, num_items, d, fine_bits);
                    hp
                } else {
                    // Not enough items - use random hyperplanes
                    let mut rng = ChaCha8Rng::seed_from_u64(self.seed + *coarse_id);
                    self.generate_random_hyperplanes(&mut rng, fine_bits)
                };

                // Hash region items with fine hyperplanes
                let fine_hashes = self.hash_embeddings_flat(&region_data, num_items, d, &fine_hp);

                // Combine coarse and fine into final bucket IDs
                let combined: Vec<u64> = fine_hashes
                    .iter()
                    .map(|&fine_id| (*coarse_id << fine_bits) | fine_id)
                    .collect();

                (indices.clone(), combined)
            })
            .collect();

        // Scatter results back to original indices
        for (indices, combined_ids) in fine_results {
            for (local_i, &global_i) in indices.iter().enumerate() {
                final_bucket_ids[global_i] = combined_ids[local_i];
            }
        }

        // Store hyperplanes (coarse only for now - fine are region-specific)
        self.hyperplanes = Some(coarse_hp);
        self.bucket_ids = final_bucket_ids;
        self.num_bits = coarse_bits + fine_bits; // Update total bits

        // Build bucket mapping and compute metrics
        let mut bucket_to_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (idx, &hash) in self.bucket_ids.iter().enumerate() {
            bucket_to_indices.entry(hash).or_default().push(idx);
        }
        self.num_buckets = bucket_to_indices.len();

        // Compute per-item metrics
        self.bucket_sizes = vec![0u32; n];
        self.centroid_similarities = vec![0.0f32; n];

        for (&_bucket_id, indices) in &bucket_to_indices {
            let bucket_size = indices.len() as u32;
            let centroid = self.compute_centroid_flat(&normalized, d, indices);

            for &idx in indices {
                self.bucket_sizes[idx] = bucket_size;
                let row = &normalized[idx * d..(idx + 1) * d];
                self.centroid_similarities[idx] = self.cosine_similarity(row, &centroid);
            }
        }

        // Compute isolation and stability scores
        if !self.skip_isolation {
            self.compute_isolation_scores_flat(&normalized, n, d, 10);
            self.compute_stability_scores_flat(&normalized, n, d);
        } else {
            self.isolation_scores = vec![0.0f32; n];
            self.stability_scores = vec![1.0f32; n];
        }

        self.fitted = true;
        self
    }

    /// Voronoi refinement: reassign items to nearest centroid among Hamming neighbors
    ///
    /// This refines LSH bucket assignments by checking if a nearby bucket's centroid
    /// is actually closer. Only checks buckets within `max_hamming_distance` bits
    /// of the current bucket (cheap Voronoi).
    ///
    /// Uses BLAS-accelerated matrix multiply for similarity computation.
    ///
    /// Returns the number of items reassigned.
    pub fn voronoi_refine_flat(
        &mut self,
        data: &[f32],
        n: usize,
        d: usize,
        max_hamming_distance: u32,
    ) -> usize {
        if !self.fitted || n == 0 {
            return 0;
        }

        // Step 1: Get unique buckets and compute centroids
        let unique_buckets: Vec<u64> = {
            let mut buckets: Vec<u64> = self.bucket_ids.clone();
            buckets.sort();
            buckets.dedup();
            buckets
        };
        let num_buckets = unique_buckets.len();
        if num_buckets <= 1 {
            return 0;
        }

        let bucket_to_idx: HashMap<u64, usize> = unique_buckets
            .iter()
            .enumerate()
            .map(|(i, &b)| (b, i))
            .collect();

        // Build bucket -> item indices mapping
        let mut bucket_item_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (i, &bid) in self.bucket_ids.iter().enumerate() {
            bucket_item_indices.entry(bid).or_default().push(i);
        }

        // Compute centroids for all buckets as ndarray (num_buckets x d)
        let mut centroids = Array2::<f32>::zeros((num_buckets, d));
        for (bucket_idx, &bucket_id) in unique_buckets.iter().enumerate() {
            if let Some(indices) = bucket_item_indices.get(&bucket_id) {
                let centroid = self.compute_centroid_flat(data, d, indices);
                for (j, &val) in centroid.iter().enumerate() {
                    centroids[[bucket_idx, j]] = val;
                }
            }
        }

        // Step 2: Choose algorithm based on restriction level
        let new_bucket_ids: Vec<u64> = if max_hamming_distance >= self.num_bits as u32 {
            // Full Voronoi: single BLAS call for all similarities
            let data_view = ArrayView2::from_shape((n, d), data).unwrap();
            let similarities = data_view.dot(&centroids.t());

            (0..n)
                .into_par_iter()
                .map(|i| {
                    let row = similarities.row(i);
                    let best_idx = row.iter()
                        .enumerate()
                        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
                        .map(|(idx, _)| idx)
                        .unwrap_or(0);
                    unique_buckets[best_idx]
                })
                .collect()
        } else {
            // Cheap Voronoi: per-bucket BLAS calls (less total FLOPS)
            // Generate Hamming neighbors by bit-flipping (O(buckets * C(bits,k)) vs O(buckets²))
            let neighbor_indices: Vec<Vec<usize>> = unique_buckets
                .par_iter()
                .map(|&bid| {
                    let mut neighbors = Vec::new();

                    // Generate all Hamming neighbors by flipping bits
                    // Hamming 0: self
                    if let Some(&idx) = bucket_to_idx.get(&bid) {
                        neighbors.push(idx);
                    }

                    // Hamming 1: flip single bits
                    if max_hamming_distance >= 1 {
                        for i in 0..self.num_bits {
                            let neighbor = bid ^ (1u64 << i);
                            if let Some(&idx) = bucket_to_idx.get(&neighbor) {
                                neighbors.push(idx);
                            }
                        }
                    }

                    // Hamming 2: flip pairs of bits
                    if max_hamming_distance >= 2 {
                        for i in 0..self.num_bits {
                            for j in (i + 1)..self.num_bits {
                                let neighbor = bid ^ (1u64 << i) ^ (1u64 << j);
                                if let Some(&idx) = bucket_to_idx.get(&neighbor) {
                                    neighbors.push(idx);
                                }
                            }
                        }
                    }

                    // Hamming 3: flip triplets of bits
                    if max_hamming_distance >= 3 {
                        for i in 0..self.num_bits {
                            for j in (i + 1)..self.num_bits {
                                for k in (j + 1)..self.num_bits {
                                    let neighbor = bid ^ (1u64 << i) ^ (1u64 << j) ^ (1u64 << k);
                                    if let Some(&idx) = bucket_to_idx.get(&neighbor) {
                                        neighbors.push(idx);
                                    }
                                }
                            }
                        }
                    }

                    neighbors
                })
                .collect();

            // Process buckets in parallel, each with its own small BLAS call
            use std::sync::atomic::{AtomicU64, Ordering};
            let new_ids: Vec<AtomicU64> = (0..n).map(|_| AtomicU64::new(0)).collect();

            unique_buckets.par_iter().enumerate().for_each(|(bucket_idx, &bucket_id)| {
                let item_indices = match bucket_item_indices.get(&bucket_id) {
                    Some(indices) => indices,
                    None => return,
                };

                let neighbors = &neighbor_indices[bucket_idx];
                let num_neighbors = neighbors.len();

                // Build neighbor centroid matrix (num_neighbors x d)
                let mut neighbor_centroids = Array2::<f32>::zeros((num_neighbors, d));
                for (local_idx, &global_idx) in neighbors.iter().enumerate() {
                    for j in 0..d {
                        neighbor_centroids[[local_idx, j]] = centroids[[global_idx, j]];
                    }
                }

                // Build item matrix for this bucket (num_items x d)
                let num_items = item_indices.len();
                let mut item_matrix = Array2::<f32>::zeros((num_items, d));
                for (local_i, &global_i) in item_indices.iter().enumerate() {
                    for j in 0..d {
                        item_matrix[[local_i, j]] = data[global_i * d + j];
                    }
                }

                // Small BLAS call: (num_items x d) @ (d x num_neighbors)
                let sims = item_matrix.dot(&neighbor_centroids.t());

                // Argmax for each item
                for (local_i, &global_i) in item_indices.iter().enumerate() {
                    let row = sims.row(local_i);
                    let best_local = row.iter()
                        .enumerate()
                        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
                        .map(|(idx, _)| idx)
                        .unwrap_or(0);
                    new_ids[global_i].store(unique_buckets[neighbors[best_local]], Ordering::Relaxed);
                }
            });

            new_ids.into_iter().map(|a| a.into_inner()).collect()
        };

        // Step 3: Count reassignments and update
        let reassigned = self.bucket_ids
            .iter()
            .zip(new_bucket_ids.iter())
            .filter(|(old, new)| old != new)
            .count();

        self.bucket_ids = new_bucket_ids;

        // Recompute bucket sizes and centroid similarities
        let mut new_bucket_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (i, &bid) in self.bucket_ids.iter().enumerate() {
            new_bucket_indices.entry(bid).or_default().push(i);
        }
        self.num_buckets = new_bucket_indices.len();

        self.bucket_sizes = vec![0u32; n];
        self.centroid_similarities = vec![0.0f32; n];

        for (&_bucket_id, indices) in &new_bucket_indices {
            let bucket_size = indices.len() as u32;
            let centroid = self.compute_centroid_flat(data, d, indices);

            for &idx in indices {
                self.bucket_sizes[idx] = bucket_size;
                let row = &data[idx * d..(idx + 1) * d];
                self.centroid_similarities[idx] = self.cosine_similarity(row, &centroid);
            }
        }

        reassigned
    }

    /// Recursive density-based splitting: subdivide dense buckets
    ///
    /// Algorithm:
    /// 1. LSH/PCA on data → buckets
    /// 2. For each bucket:
    ///    - If dense (size > threshold): recurse on bucket contents
    ///    - If sparse (size <= threshold): done, this is a leaf
    ///
    /// Dense regions get more splits, sparse regions stop early.
    /// Adaptive resolution based on local density.
    ///
    /// # Arguments
    /// * `data` - Flat embedding data
    /// * `n` - Number of samples
    /// * `d` - Embedding dimension
    /// * `bits_per_level` - Bits to use at each recursion level (default: 4)
    /// * `max_depth` - Maximum recursion depth (default: 4)
    /// * `min_bucket_size` - Stop splitting below this size (default: 20)
    ///
    /// # Returns
    /// Self for chaining
    pub fn fit_recursive_flat(
        &mut self,
        data: &[f32],
        n: usize,
        d: usize,
        bits_per_level: usize,
        max_depth: usize,
        min_bucket_size: usize,
    ) -> &Self {
        if n == 0 || d != self.embedding_dim {
            self.fitted = true;
            return self;
        }

        // Normalize data
        let mut normalized = vec![0.0f32; n * d];
        normalized
            .par_chunks_mut(d)
            .enumerate()
            .for_each(|(i, out_row)| {
                let start = i * d;
                let in_row = &data[start..start + d];
                let norm: f32 = in_row.iter().map(|x| x * x).sum::<f32>().sqrt();
                if norm > 0.0 {
                    for (j, &val) in in_row.iter().enumerate() {
                        out_row[j] = val / norm;
                    }
                } else {
                    out_row.copy_from_slice(in_row);
                }
            });

        // Initialize bucket IDs to 0
        let mut bucket_ids = vec![0u64; n];
        let mut current_bit_offset = 0usize;

        // Track which items are still "active" (in buckets that need splitting)
        // Initially all items are active with bucket prefix 0
        let mut active_groups: Vec<(u64, Vec<usize>)> = vec![(0u64, (0..n).collect())];

        for depth in 0..max_depth {
            if active_groups.is_empty() {
                break;
            }

            let mut next_active_groups: Vec<(u64, Vec<usize>)> = Vec::new();

            // Process each active group
            for (prefix, indices) in active_groups {
                let group_size = indices.len();

                // If small enough, this group is done (leaf)
                if group_size <= min_bucket_size {
                    // Keep current bucket ID (prefix is already set)
                    continue;
                }

                // Extract data for this group
                let mut group_data = vec![0.0f32; group_size * d];
                for (local_i, &global_i) in indices.iter().enumerate() {
                    group_data[local_i * d..(local_i + 1) * d]
                        .copy_from_slice(&normalized[global_i * d..(global_i + 1) * d]);
                }

                // LSH/PCA on this group
                let group_hyperplanes = self.compute_group_hyperplanes(
                    &group_data, group_size, d, bits_per_level
                );

                // Hash group items
                let group_hashes = self.hash_embeddings_flat(
                    &group_data, group_size, d, &group_hyperplanes
                );

                // Update bucket IDs with new bits
                for (local_i, &global_i) in indices.iter().enumerate() {
                    let new_bits = group_hashes[local_i];
                    bucket_ids[global_i] = prefix | (new_bits << current_bit_offset);
                }

                // Group items by their new hash within this group
                let mut sub_buckets: HashMap<u64, Vec<usize>> = HashMap::new();
                for (local_i, &global_i) in indices.iter().enumerate() {
                    let full_id = bucket_ids[global_i];
                    sub_buckets.entry(full_id).or_default().push(global_i);
                }

                // Add dense sub-buckets to next iteration
                for (full_id, sub_indices) in sub_buckets {
                    if sub_indices.len() > min_bucket_size {
                        next_active_groups.push((full_id, sub_indices));
                    }
                    // Sparse buckets are done - their IDs are already set
                }
            }

            active_groups = next_active_groups;
            current_bit_offset += bits_per_level;
        }

        // Store results
        self.bucket_ids = bucket_ids;
        self.num_bits = current_bit_offset; // Total bits used

        // Build bucket mapping and compute metrics
        let mut bucket_to_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (idx, &hash) in self.bucket_ids.iter().enumerate() {
            bucket_to_indices.entry(hash).or_default().push(idx);
        }
        self.num_buckets = bucket_to_indices.len();

        // Compute per-item metrics
        self.bucket_sizes = vec![0u32; n];
        self.centroid_similarities = vec![0.0f32; n];

        for (&_bucket_id, indices) in &bucket_to_indices {
            let bucket_size = indices.len() as u32;
            let centroid = self.compute_centroid_flat(&normalized, d, indices);

            for &idx in indices {
                self.bucket_sizes[idx] = bucket_size;
                let row = &normalized[idx * d..(idx + 1) * d];
                self.centroid_similarities[idx] = self.cosine_similarity(row, &centroid);
            }
        }

        // Compute isolation and stability scores
        if !self.skip_isolation {
            self.compute_isolation_scores_flat(&normalized, n, d, 10);
        } else {
            self.isolation_scores = vec![0.0f32; n];
        }
        // Skip stability for recursive (different hyperplanes per group)
        self.stability_scores = vec![1.0f32; n];

        self.fitted = true;
        self
    }

    /// Compute hyperplanes for a group using LSH → centroids → PCA
    fn compute_group_hyperplanes(
        &self,
        data: &[f32],
        n: usize,
        d: usize,
        num_bits: usize,
    ) -> DMatrix<f32> {
        if n <= num_bits + 1 {
            // Not enough data for PCA, use random
            let mut rng = ChaCha8Rng::seed_from_u64(self.seed + n as u64);
            return self.generate_random_hyperplanes(&mut rng, num_bits);
        }

        // Random hash to get initial buckets
        let mut rng = ChaCha8Rng::seed_from_u64(self.seed + n as u64);
        let random_hp = self.generate_random_hyperplanes(&mut rng, num_bits);
        let random_hashes = self.hash_embeddings_flat(data, n, d, &random_hp);

        // Compute centroids
        let centroids = self.compute_all_centroids_flat(data, n, d, &random_hashes, 1);
        let num_centroids = centroids.len() / d;

        if num_centroids > num_bits {
            let (pca_hp, _, _) = self.compute_pca_flat(&centroids, num_centroids, d, num_bits);
            pca_hp
        } else {
            random_hp
        }
    }

    /// ITQ (Iterative Quantization): learn rotation to minimize quantization error
    ///
    /// Standard PCA-LSH uses sign(X @ PCA.T) which loses information near zero.
    /// ITQ finds a rotation R such that sign(X @ PCA.T @ R) has minimal quantization loss.
    ///
    /// Algorithm:
    /// 1. Compute PCA projections V = X @ PCA.T
    /// 2. Iterate: B = sign(V @ R), then solve Procrustes for R
    /// 3. Final hyperplanes = PCA @ R
    ///
    /// Typically gives 5-15% improvement over raw PCA.
    ///
    /// # Arguments
    /// * `data` - Flat embedding data
    /// * `n` - Number of samples
    /// * `d` - Embedding dimension
    /// * `max_iterations` - ITQ iterations (default: 50)
    ///
    /// # Returns
    /// Self for chaining
    pub fn fit_itq_flat(
        &mut self,
        data: &[f32],
        n: usize,
        d: usize,
        max_iterations: usize,
    ) -> &Self {
        if n == 0 || d != self.embedding_dim {
            self.fitted = true;
            return self;
        }

        // Normalize data
        let mut normalized = vec![0.0f32; n * d];
        normalized
            .par_chunks_mut(d)
            .enumerate()
            .for_each(|(i, out_row)| {
                let start = i * d;
                let in_row = &data[start..start + d];
                let norm: f32 = in_row.iter().map(|x| x * x).sum::<f32>().sqrt();
                if norm > 0.0 {
                    for (j, &val) in in_row.iter().enumerate() {
                        out_row[j] = val / norm;
                    }
                } else {
                    out_row.copy_from_slice(in_row);
                }
            });

        // Step 1: Standard PCA to get initial hyperplanes
        let mut rng = ChaCha8Rng::seed_from_u64(self.seed);
        let random_hp = self.generate_random_hyperplanes(&mut rng, self.num_bits);
        let random_hashes = self.hash_embeddings_flat(&normalized, n, d, &random_hp);

        let centroids = self.compute_all_centroids_flat(&normalized, n, d, &random_hashes, 2);
        let num_centroids = centroids.len() / d;

        let (pca_hp, variance, ev) = if num_centroids > self.num_bits {
            self.compute_pca_flat(&centroids, num_centroids, d, self.num_bits)
        } else {
            (random_hp.clone(), 0.0, Vec::new())
        };
        self.stage1_variance = variance;
        self.eigenvalues = ev;

        let k = self.num_bits;

        // Step 2: Compute PCA projections V = X @ PCA.T (n × k)
        // PCA is k × d, so we need X @ PCA.T = (n × d) @ (d × k)
        let data_view = ArrayView2::from_shape((n, d), &normalized).unwrap();
        let mut pca_t = Array2::<f32>::zeros((d, k));
        for i in 0..k {
            for j in 0..d {
                pca_t[[j, i]] = pca_hp[(i, j)];
            }
        }
        let v = data_view.dot(&pca_t); // n × k

        // Step 3: Initialize rotation R = Identity (k × k)
        let mut r = DMatrix::<f32>::identity(k, k);

        // Step 4: ITQ iterations
        for _iter in 0..max_iterations {
            // Compute V @ R (n × k)
            let mut vr = Array2::<f32>::zeros((n, k));
            for i in 0..n {
                for j in 0..k {
                    let mut sum = 0.0f32;
                    for l in 0..k {
                        sum += v[[i, l]] * r[(l, j)];
                    }
                    vr[[i, j]] = sum;
                }
            }

            // B = sign(V @ R)
            let b: Array2<f32> = vr.mapv(|x| if x >= 0.0 { 1.0 } else { -1.0 });

            // Solve Procrustes: minimize ||V @ R - B||² subject to R orthogonal
            // C = V.T @ B (k × k)
            let mut c = DMatrix::<f32>::zeros(k, k);
            for i in 0..k {
                for j in 0..k {
                    let mut sum = 0.0f32;
                    for l in 0..n {
                        sum += v[[l, i]] * b[[l, j]];
                    }
                    c[(i, j)] = sum;
                }
            }

            // SVD of C: C = U @ S @ V.T
            // R_new = U @ V.T (closest orthogonal matrix to C)
            let svd = c.svd(true, true);
            if let (Some(u), Some(vt)) = (svd.u, svd.v_t) {
                // R = U @ V.T
                r = &u * &vt;
            }
        }

        // Step 5: Final hyperplanes = PCA @ R
        // PCA is k × d, R is k × k
        // rotated_hp[i] = sum_j R[j,i] * PCA[j]
        let mut rotated_hp = DMatrix::<f32>::zeros(k, d);
        for i in 0..k {
            for j in 0..d {
                let mut sum = 0.0f32;
                for l in 0..k {
                    sum += r[(l, i)] * pca_hp[(l, j)];
                }
                rotated_hp[(i, j)] = sum;
            }
        }

        // Normalize rows of rotated hyperplanes
        for i in 0..k {
            let norm: f32 = (0..d).map(|j| rotated_hp[(i, j)].powi(2)).sum::<f32>().sqrt();
            if norm > 0.0 {
                for j in 0..d {
                    rotated_hp[(i, j)] /= norm;
                }
            }
        }

        self.hyperplanes = Some(rotated_hp.clone());

        // Step 6: Hash with rotated hyperplanes
        let hashes = self.hash_embeddings_flat(&normalized, n, d, &rotated_hp);
        self.bucket_ids = hashes.clone();

        // Build bucket mapping
        let mut bucket_to_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (idx, &hash) in hashes.iter().enumerate() {
            bucket_to_indices.entry(hash).or_default().push(idx);
        }
        self.num_buckets = bucket_to_indices.len();

        // Compute per-item metrics
        self.bucket_sizes = vec![0u32; n];
        self.centroid_similarities = vec![0.0f32; n];

        for (&_bucket_id, indices) in &bucket_to_indices {
            let bucket_size = indices.len() as u32;
            let centroid = self.compute_centroid_flat(&normalized, d, indices);

            for &idx in indices {
                self.bucket_sizes[idx] = bucket_size;
                let row = &normalized[idx * d..(idx + 1) * d];
                self.centroid_similarities[idx] = self.cosine_similarity(row, &centroid);
            }
        }

        // Compute isolation and stability scores
        if !self.skip_isolation {
            self.compute_isolation_scores_flat(&normalized, n, d, 10);
            self.compute_stability_scores_flat(&normalized, n, d);
        } else {
            self.isolation_scores = vec![0.0f32; n];
            self.stability_scores = vec![1.0f32; n];
        }

        self.fitted = true;
        self
    }

    /// Margin-based selective refinement: only refine low-confidence items
    ///
    /// Items near hyperplane boundaries (low margin) are candidates for misassignment.
    /// This method only refines those items, skipping high-confidence ones.
    ///
    /// Margin = minimum |dot product| across all hyperplanes.
    /// Low margin = item is close to at least one decision boundary.
    ///
    /// # Arguments
    /// * `data` - Flat embedding data (same as fit)
    /// * `n` - Number of samples
    /// * `d` - Embedding dimension
    /// * `margin_threshold` - Items with margin below this are refined (default: 0.1)
    /// * `max_hamming_distance` - Hamming radius for neighbor search (default: 2)
    ///
    /// # Returns
    /// (items_refined, items_reassigned) - how many were checked vs actually moved
    pub fn voronoi_refine_margin_flat(
        &mut self,
        data: &[f32],
        n: usize,
        d: usize,
        margin_threshold: f32,
        max_hamming_distance: u32,
    ) -> (usize, usize) {
        if !self.fitted || n == 0 {
            return (0, 0);
        }

        let hyperplanes = match &self.hyperplanes {
            Some(hp) => hp,
            None => return (0, 0),
        };

        // Step 1: Compute margins for all items
        // Margin = min(|dot product with hyperplane|) across all hyperplanes
        let num_hp = hyperplanes.nrows();

        // Convert hyperplanes to ndarray for BLAS
        let mut hp_array = Array2::<f32>::zeros((num_hp, d));
        for i in 0..num_hp {
            for j in 0..d {
                hp_array[[i, j]] = hyperplanes[(i, j)];
            }
        }

        // Compute all dot products: (n x d) @ (d x num_hp) = (n x num_hp)
        let data_view = ArrayView2::from_shape((n, d), data).unwrap();
        let dots = data_view.dot(&hp_array.t());

        // Compute margin for each item (min absolute dot product)
        let margins: Vec<f32> = (0..n)
            .into_par_iter()
            .map(|i| {
                let row = dots.row(i);
                row.iter()
                    .map(|&x| x.abs())
                    .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
                    .unwrap_or(0.0)
            })
            .collect();

        // Step 2: Identify low-margin items to refine
        let low_margin_indices: Vec<usize> = margins
            .iter()
            .enumerate()
            .filter(|(_, &m)| m < margin_threshold)
            .map(|(i, _)| i)
            .collect();

        let items_to_refine = low_margin_indices.len();
        if items_to_refine == 0 {
            return (0, 0);
        }

        // Step 3: Get unique buckets and compute centroids (same as voronoi_refine)
        let unique_buckets: Vec<u64> = {
            let mut buckets: Vec<u64> = self.bucket_ids.clone();
            buckets.sort();
            buckets.dedup();
            buckets
        };
        let num_buckets = unique_buckets.len();
        if num_buckets <= 1 {
            return (items_to_refine, 0);
        }

        let bucket_to_idx: HashMap<u64, usize> = unique_buckets
            .iter()
            .enumerate()
            .map(|(i, &b)| (b, i))
            .collect();

        // Build bucket -> item indices mapping
        let mut bucket_item_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (i, &bid) in self.bucket_ids.iter().enumerate() {
            bucket_item_indices.entry(bid).or_default().push(i);
        }

        // Compute centroids
        let mut centroids = Array2::<f32>::zeros((num_buckets, d));
        for (bucket_idx, &bucket_id) in unique_buckets.iter().enumerate() {
            if let Some(indices) = bucket_item_indices.get(&bucket_id) {
                let centroid = self.compute_centroid_flat(data, d, indices);
                for (j, &val) in centroid.iter().enumerate() {
                    centroids[[bucket_idx, j]] = val;
                }
            }
        }

        // Step 4: Generate Hamming neighbors for each bucket (same bit-flip approach)
        let neighbor_indices: Vec<Vec<usize>> = unique_buckets
            .par_iter()
            .map(|&bid| {
                let mut neighbors = Vec::new();

                if let Some(&idx) = bucket_to_idx.get(&bid) {
                    neighbors.push(idx);
                }

                if max_hamming_distance >= 1 {
                    for i in 0..self.num_bits {
                        let neighbor = bid ^ (1u64 << i);
                        if let Some(&idx) = bucket_to_idx.get(&neighbor) {
                            neighbors.push(idx);
                        }
                    }
                }

                if max_hamming_distance >= 2 {
                    for i in 0..self.num_bits {
                        for j in (i + 1)..self.num_bits {
                            let neighbor = bid ^ (1u64 << i) ^ (1u64 << j);
                            if let Some(&idx) = bucket_to_idx.get(&neighbor) {
                                neighbors.push(idx);
                            }
                        }
                    }
                }

                neighbors
            })
            .collect();

        // Step 5: Refine only low-margin items
        use std::sync::atomic::{AtomicUsize, Ordering};
        let reassigned_count = AtomicUsize::new(0);

        let new_bucket_ids: Vec<u64> = (0..n)
            .into_par_iter()
            .map(|i| {
                // If high margin, keep original bucket
                if margins[i] >= margin_threshold {
                    return self.bucket_ids[i];
                }

                // Low margin - check neighbor centroids
                let current_bucket = self.bucket_ids[i];
                let bucket_idx = match bucket_to_idx.get(&current_bucket) {
                    Some(&idx) => idx,
                    None => return current_bucket,
                };

                let neighbors = &neighbor_indices[bucket_idx];
                let item_row = &data[i * d..(i + 1) * d];

                // Find best centroid among neighbors
                let mut best_bucket = current_bucket;
                let mut best_sim = f32::NEG_INFINITY;

                for &neighbor_idx in neighbors {
                    let mut sim = 0.0f32;
                    for j in 0..d {
                        sim += item_row[j] * centroids[[neighbor_idx, j]];
                    }
                    if sim > best_sim {
                        best_sim = sim;
                        best_bucket = unique_buckets[neighbor_idx];
                    }
                }

                if best_bucket != current_bucket {
                    reassigned_count.fetch_add(1, Ordering::Relaxed);
                }

                best_bucket
            })
            .collect();

        let items_reassigned = reassigned_count.load(Ordering::Relaxed);

        // Step 6: Update bucket assignments and metrics
        self.bucket_ids = new_bucket_ids;

        let mut new_bucket_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (i, &bid) in self.bucket_ids.iter().enumerate() {
            new_bucket_indices.entry(bid).or_default().push(i);
        }
        self.num_buckets = new_bucket_indices.len();

        self.bucket_sizes = vec![0u32; n];
        self.centroid_similarities = vec![0.0f32; n];

        for (&_bucket_id, indices) in &new_bucket_indices {
            let bucket_size = indices.len() as u32;
            let centroid = self.compute_centroid_flat(data, d, indices);

            for &idx in indices {
                self.bucket_sizes[idx] = bucket_size;
                let row = &data[idx * d..(idx + 1) * d];
                self.centroid_similarities[idx] = self.cosine_similarity(row, &centroid);
            }
        }

        (items_to_refine, items_reassigned)
    }

    /// Multi-resolution analysis: find structure among outliers at coarser resolutions
    ///
    /// Since PCA-LSH bit 0 (LSB) = highest variance PCA component, masking the
    /// low K bits gives the K most informative bits. At each depth K=1..num_bits,
    /// we compute coarse bucket IDs and sizes. For each item, we find the smallest
    /// depth where its coarse bucket exceeds dense_threshold ("recovery depth").
    ///
    /// # Arguments
    /// * `dense_threshold` - Bucket size threshold to consider "dense" (default: 10)
    ///
    /// # Returns
    /// MultiResolutionAnalysis with per-item recovery depth/ratio and per-depth stats
    pub fn multi_resolution_analysis(&self, dense_threshold: u32) -> MultiResolutionAnalysis {
        let n = self.bucket_ids.len();
        let num_bits = self.num_bits;

        if !self.fitted || n == 0 {
            return MultiResolutionAnalysis {
                n_items: 0,
                num_bits,
                recovery_depth: Vec::new(),
                recovery_ratio: Vec::new(),
                buckets_per_depth: Vec::new(),
                mean_size_per_depth: Vec::new(),
                dense_threshold,
            };
        }

        // Check if already dense at full resolution
        let full_res_dense: Vec<bool> = self.bucket_sizes.iter()
            .map(|&sz| sz >= dense_threshold)
            .collect();

        // Per-item recovery depth: start at num_bits+1 (never recovered)
        let mut recovery_depth = vec![(num_bits as u32) + 1; n];
        let mut recovery_ratio = vec![0.0f32; n];

        // Items already dense at full resolution get depth=0
        for i in 0..n {
            if full_res_dense[i] {
                recovery_depth[i] = 0;
                let expected = n as f32 / (1u64 << num_bits) as f32;
                recovery_ratio[i] = if expected > 0.0 {
                    self.bucket_sizes[i] as f32 / expected
                } else {
                    self.bucket_sizes[i] as f32
                };
            }
        }

        let mut buckets_per_depth = Vec::with_capacity(num_bits);
        let mut mean_size_per_depth = Vec::with_capacity(num_bits);

        // For each depth K=1..num_bits, mask low K bits to get coarse bucket IDs
        for k in 1..=num_bits {
            let mask = (1u64 << k) - 1; // e.g., k=3 -> mask=0b111

            // Compute coarse bucket sizes
            let mut coarse_sizes: HashMap<u64, u32> = HashMap::new();
            for &bid in &self.bucket_ids {
                let coarse_id = bid & mask;
                *coarse_sizes.entry(coarse_id).or_default() += 1;
            }

            let n_coarse_buckets = coarse_sizes.len();
            let total: u32 = coarse_sizes.values().sum();
            let mean_size = if n_coarse_buckets > 0 {
                total as f32 / n_coarse_buckets as f32
            } else {
                0.0
            };

            buckets_per_depth.push(n_coarse_buckets);
            mean_size_per_depth.push(mean_size);

            let expected_size = n as f32 / (1u64 << k) as f32;

            // Check each item: if not yet recovered and coarse bucket is dense
            for i in 0..n {
                if recovery_depth[i] <= num_bits as u32 {
                    continue; // Already recovered at a finer depth
                }
                let coarse_id = self.bucket_ids[i] & mask;
                let coarse_size = coarse_sizes[&coarse_id];
                if coarse_size >= dense_threshold {
                    recovery_depth[i] = k as u32;
                    recovery_ratio[i] = if expected_size > 0.0 {
                        coarse_size as f32 / expected_size
                    } else {
                        coarse_size as f32
                    };
                }
            }
        }

        MultiResolutionAnalysis {
            n_items: n,
            num_bits,
            recovery_depth,
            recovery_ratio,
            buckets_per_depth,
            mean_size_per_depth,
            dense_threshold,
        }
    }

    /// Multi-resolution bridge persistence analysis using relative thresholds
    ///
    /// Detects cross-bucket connectors at every LSH depth K=1..num_bits using a
    /// **relative similarity threshold**: an item is a connector if its similarity
    /// to some other bucket's centroid is at least `relative_threshold` × its
    /// similarity to its own centroid.
    ///
    /// ## Why relative thresholds?
    ///
    /// Absolute thresholds (e.g., "sim to other centroid >= 0.3") are
    /// depth-dependent: at coarse depths (few large buckets), centroid similarities
    /// are low and 0.3 catches nothing. At fine depths (many small buckets),
    /// similarities are high and 0.3 catches everything. The relative ratio
    /// `other_sim / own_sim` is stable across depths because both scale together
    /// as centroids get tighter.
    ///
    /// Empirically on wiki-10k with 12-bit LSH:
    /// - Absolute threshold 0.3: catches 0% at K=1, 99.5% at K=12 (useless)
    /// - Relative threshold 0.8: catches 7% at K=1, 50% at K=12 (discriminating)
    ///
    /// ## Algorithm
    ///
    /// For each depth K = 1..num_bits:
    /// 1. Mask bucket IDs to K low bits → 2^K coarse buckets
    /// 2. Compute L2-normalized coarse centroids
    /// 3. BLAS matmul: all items × all centroids → (n × n_coarse) similarity matrix
    /// 4. Per item: own_sim = sim[i, own_bucket]. For each other centroid j,
    ///    if sim[i,j] / own_sim >= relative_threshold → item connects own ↔ j
    /// 5. Map coarse connections to full-resolution bucket pairs via u64 bitmask
    ///
    /// ## Cost
    ///
    /// Dominated by the BLAS matmul at each depth: O(n × n_coarse × d).
    /// Total across all depths: O(n × d × Σ(2^k for k=1..num_bits)).
    /// For n=10k, d=384, 12 bits: ~0.2s. Scales linearly with n.
    ///
    /// # Arguments
    /// * `data` - Flat embedding data (normalized, same as fit)
    /// * `n` - Number of samples
    /// * `d` - Embedding dimension
    /// * `relative_threshold` - Ratio threshold: other_sim/own_sim >= this → connector.
    ///   Recommended: 0.8 (item is 80% as similar to another centroid as its own).
    ///   Lower values catch more connectors; higher values are more selective.
    pub fn bridge_persistence_analysis(
        &self,
        data: &[f32],
        n: usize,
        d: usize,
        relative_threshold: f32,
    ) -> BridgePersistence {
        let num_bits = self.num_bits;

        if !self.fitted || n == 0 {
            return BridgePersistence {
                n_items: 0,
                num_bits,
                bridge_persistence: Vec::new(),
                max_bridge_depth: Vec::new(),
                min_bridge_depth: Vec::new(),
                bridge_ratio: Vec::new(),
                bridges_per_depth: Vec::new(),
                pairs_per_depth: Vec::new(),
                buckets_per_depth: Vec::new(),
                relative_threshold,
                persistent_pairs: Vec::new(),
            };
        }

        let data_view = ArrayView2::from_shape((n, d), data).unwrap();

        // Per-item tracking
        let mut bridge_persistence = vec![0u32; n];
        let mut max_bridge_depth = vec![0u32; n];
        let mut min_bridge_depth = vec![0u32; n];
        let mut bridge_ratio = vec![0.0f32; n]; // peak max_other/own across depths

        // Per-depth stats
        let mut bridges_per_depth = Vec::with_capacity(num_bits);
        let mut pairs_per_depth = Vec::with_capacity(num_bits);
        let mut buckets_per_depth = Vec::with_capacity(num_bits);

        // Track full-res pair connections: (min_bucket, max_bucket) -> depth bitmask
        let mut pair_depth_mask: HashMap<(u64, u64), u64> = HashMap::new();

        for k in 1..=num_bits {
            let mask = (1u64 << k) - 1;

            // Build coarse bucket -> item indices mapping
            let mut coarse_to_items: HashMap<u64, Vec<usize>> = HashMap::new();
            for (i, &bid) in self.bucket_ids.iter().enumerate() {
                let coarse_id = bid & mask;
                coarse_to_items.entry(coarse_id).or_default().push(i);
            }

            let coarse_bucket_ids: Vec<u64> = {
                let mut ids: Vec<u64> = coarse_to_items.keys().copied().collect();
                ids.sort();
                ids
            };
            let n_coarse = coarse_bucket_ids.len();
            buckets_per_depth.push(n_coarse);

            if n_coarse < 2 {
                // Only one coarse bucket — no connections possible
                bridges_per_depth.push(0);
                pairs_per_depth.push(0);
                continue;
            }

            // Compute coarse centroids as Array2 for BLAS matmul
            let mut centroids = Array2::<f32>::zeros((n_coarse, d));
            for (ci, &cid) in coarse_bucket_ids.iter().enumerate() {
                if let Some(indices) = coarse_to_items.get(&cid) {
                    let centroid = self.compute_centroid_flat(data, d, indices);
                    for (j, &val) in centroid.iter().enumerate() {
                        centroids[[ci, j]] = val;
                    }
                }
            }

            let coarse_id_to_idx: HashMap<u64, usize> = coarse_bucket_ids
                .iter()
                .enumerate()
                .map(|(i, &id)| (id, i))
                .collect();

            // BLAS matmul: all items × all centroids → (n, n_coarse) similarity matrix
            let sims = data_view.dot(&centroids.t());

            // Per-item: compute relative threshold, identify connectors
            let mut depth_connectors: usize = 0;
            let mut depth_pairs: usize = 0;

            for i in 0..n {
                let coarse_id = self.bucket_ids[i] & mask;
                let own_idx = coarse_id_to_idx[&coarse_id];
                let own_sim = sims[[i, own_idx]];

                // Skip items with non-positive own similarity (degenerate)
                if own_sim <= 0.0 {
                    continue;
                }

                // Check all other centroids against relative threshold
                let mut is_connector = false;
                let mut connected_coarse: Vec<u64> = Vec::new();

                for (ci, &cid) in coarse_bucket_ids.iter().enumerate() {
                    if ci == own_idx {
                        continue;
                    }
                    let other_sim = sims[[i, ci]];
                    let ratio = other_sim / own_sim;

                    if ratio >= relative_threshold {
                        is_connector = true;
                        connected_coarse.push(cid);

                        // Track peak ratio for this item across all depths
                        if ratio > bridge_ratio[i] {
                            bridge_ratio[i] = ratio;
                        }
                    }
                }

                if is_connector {
                    depth_connectors += 1;
                    bridge_persistence[i] += 1;

                    // Update max_bridge_depth (coarsest = largest K)
                    max_bridge_depth[i] = k as u32;

                    // Update min_bridge_depth (finest = smallest K, first time seen)
                    if min_bridge_depth[i] == 0 {
                        min_bridge_depth[i] = k as u32;
                    }

                    // Map coarse connections to full-res bucket pairs
                    let own_coarse_id = coarse_id;
                    if let Some(own_full_items) = coarse_to_items.get(&own_coarse_id) {
                        let own_full_buckets: Vec<u64> = {
                            let mut b: Vec<u64> = own_full_items.iter().map(|&idx| self.bucket_ids[idx]).collect();
                            b.sort();
                            b.dedup();
                            b
                        };

                        for &other_coarse_id in &connected_coarse {
                            depth_pairs += 1;
                            if let Some(other_full_items) = coarse_to_items.get(&other_coarse_id) {
                                let other_full_buckets: Vec<u64> = {
                                    let mut b: Vec<u64> = other_full_items.iter().map(|&idx| self.bucket_ids[idx]).collect();
                                    b.sort();
                                    b.dedup();
                                    b
                                };

                                // Set bit K for all full-res pairs between these coarse buckets
                                for &fa in &own_full_buckets {
                                    for &fb in &other_full_buckets {
                                        let key = if fa <= fb { (fa, fb) } else { (fb, fa) };
                                        let entry = pair_depth_mask.entry(key).or_insert(0);
                                        *entry |= 1u64 << (k - 1);
                                    }
                                }
                            }
                        }
                    }
                }
            }

            bridges_per_depth.push(depth_connectors);
            pairs_per_depth.push(depth_pairs);
        }

        // Build persistent_pairs from pair_depth_mask
        let mut persistent_pairs: Vec<(u64, u64, u32, u32, u32)> = pair_depth_mask
            .into_iter()
            .map(|((a, b), bitmask)| {
                let persistence = bitmask.count_ones();
                let min_depth = bitmask.trailing_zeros() + 1;
                let max_depth = 64 - bitmask.leading_zeros();
                (a, b, persistence, min_depth, max_depth)
            })
            .collect();
        persistent_pairs.sort_by(|a, b| b.2.cmp(&a.2));

        BridgePersistence {
            n_items: n,
            num_bits,
            bridge_persistence,
            max_bridge_depth,
            min_bridge_depth,
            bridge_ratio,
            bridges_per_depth,
            pairs_per_depth,
            buckets_per_depth,
            relative_threshold,
            persistent_pairs,
        }
    }

    /// Analyze bridge items to discover region connectivity
    ///
    /// Bridges are items with low centroid similarity - they sit between regions.
    /// This analysis finds:
    /// 1. Which items are bridges
    /// 2. Which buckets each bridge connects
    /// 3. A region adjacency matrix showing how regions connect
    ///
    /// # Arguments
    /// * `data` - Flat embedding data (same as fit)
    /// * `n` - Number of samples
    /// * `d` - Embedding dimension
    /// * `bridge_threshold` - Items with centroid_sim below this are bridges (default: 0.5)
    /// * `connection_threshold` - Min similarity to consider a bridge connected to a bucket (default: 0.3)
    ///
    /// # Returns
    /// BridgeAnalysis struct with connectivity information
    pub fn analyze_bridges_flat(
        &self,
        data: &[f32],
        n: usize,
        d: usize,
        bridge_threshold: f32,
        connection_threshold: f32,
    ) -> BridgeAnalysis {
        if !self.fitted || n == 0 {
            return BridgeAnalysis {
                bridge_indices: Vec::new(),
                bridge_connections: Vec::new(),
                bucket_ids: Vec::new(),
                adjacency_matrix: Vec::new(),
                n_buckets: 0,
                bridge_threshold,
            };
        }

        // Get unique buckets and build centroids
        let unique_buckets: Vec<u64> = {
            let mut buckets: Vec<u64> = self.bucket_ids.clone();
            buckets.sort();
            buckets.dedup();
            buckets
        };
        let n_buckets = unique_buckets.len();

        let bucket_to_idx: HashMap<u64, usize> = unique_buckets
            .iter()
            .enumerate()
            .map(|(i, &b)| (b, i))
            .collect();

        // Build bucket -> item indices
        let mut bucket_item_indices: HashMap<u64, Vec<usize>> = HashMap::new();
        for (i, &bid) in self.bucket_ids.iter().enumerate() {
            bucket_item_indices.entry(bid).or_default().push(i);
        }

        // Compute centroids for all buckets
        let mut centroids = Array2::<f32>::zeros((n_buckets, d));
        for (bucket_idx, &bucket_id) in unique_buckets.iter().enumerate() {
            if let Some(indices) = bucket_item_indices.get(&bucket_id) {
                let centroid = self.compute_centroid_flat(data, d, indices);
                for (j, &val) in centroid.iter().enumerate() {
                    centroids[[bucket_idx, j]] = val;
                }
            }
        }

        // Find bridge items (low centroid similarity)
        let bridge_indices: Vec<usize> = self.centroid_similarities
            .iter()
            .enumerate()
            .filter(|(_, &sim)| sim < bridge_threshold)
            .map(|(i, _)| i)
            .collect();

        // For each bridge, find which buckets it's connected to
        let bridge_connections: Vec<(usize, u64, Vec<u64>)> = bridge_indices
            .par_iter()
            .map(|&item_idx| {
                let item_row = &data[item_idx * d..(item_idx + 1) * d];
                let primary_bucket = self.bucket_ids[item_idx];

                // Compute similarity to all bucket centroids
                let mut connected: Vec<u64> = Vec::new();
                for (bucket_idx, &bucket_id) in unique_buckets.iter().enumerate() {
                    if bucket_id == primary_bucket {
                        continue; // Skip own bucket
                    }
                    let mut sim = 0.0f32;
                    for j in 0..d {
                        sim += item_row[j] * centroids[[bucket_idx, j]];
                    }
                    if sim >= connection_threshold {
                        connected.push(bucket_id);
                    }
                }

                (item_idx, primary_bucket, connected)
            })
            .collect();

        // Build adjacency matrix
        let mut adjacency = vec![0u32; n_buckets * n_buckets];
        for (_, primary_bucket, connected) in &bridge_connections {
            let primary_idx = bucket_to_idx[primary_bucket];
            for &other_bucket in connected {
                let other_idx = bucket_to_idx[&other_bucket];
                // Symmetric: bridge connects both ways
                adjacency[primary_idx * n_buckets + other_idx] += 1;
                adjacency[other_idx * n_buckets + primary_idx] += 1;
            }
        }

        BridgeAnalysis {
            bridge_indices,
            bridge_connections,
            bucket_ids: unique_buckets,
            adjacency_matrix: adjacency,
            n_buckets,
            bridge_threshold,
        }
    }
}

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

    fn generate_test_embeddings(n: usize, dim: usize, seed: u64) -> Vec<Vec<f32>> {
        let mut rng = ChaCha8Rng::seed_from_u64(seed);
        (0..n)
            .map(|_| {
                let mut emb: Vec<f32> = (0..dim).map(|_| rng.gen_range(-1.0..1.0)).collect();
                let norm: f32 = emb.iter().map(|x| x * x).sum::<f32>().sqrt();
                for x in &mut emb {
                    *x /= norm;
                }
                emb
            })
            .collect()
    }

    #[test]
    fn test_classifier_creation() {
        let classifier = DensityClassifier::with_defaults(384);
        assert!(!classifier.is_fitted());
    }

    #[test]
    fn test_fit_small_corpus() {
        let embeddings = generate_test_embeddings(100, 128, 42);
        let mut classifier = DensityClassifier::with_defaults(128);
        classifier.fit(&embeddings);

        assert!(classifier.is_fitted());
        let report = classifier.report();
        assert_eq!(report.corpus_size, 100);
        assert!(report.mean_centroid_similarity > 0.0);
        assert!(report.mean_isolation_score > 0.0);
    }

    #[test]
    fn test_raw_scores() {
        let embeddings = generate_test_embeddings(100, 128, 42);
        let mut classifier = DensityClassifier::with_defaults(128);
        classifier.fit(&embeddings);

        assert_eq!(classifier.get_bucket_sizes().len(), 100);
        assert_eq!(classifier.get_centroid_similarities().len(), 100);
        assert_eq!(classifier.get_isolation_scores().len(), 100);
        assert_eq!(classifier.get_stability_scores().len(), 100);

        // All centroid similarities should be between -1 and 1 (cosine similarity range)
        for &sim in classifier.get_centroid_similarities() {
            assert!(sim >= -1.0 && sim <= 1.0 + 1e-6, "sim={} out of range", sim);
        }

        // Isolation scores should be positive (top_k - median)
        for &iso in classifier.get_isolation_scores() {
            assert!(iso >= 0.0);
        }

        // Stability scores should be between 0 and 1
        for &stab in classifier.get_stability_scores() {
            assert!(stab >= 0.0 && stab <= 1.0, "stability={} out of range", stab);
        }
    }

    #[test]
    fn test_multi_resolution_analysis() {
        let embeddings = generate_test_embeddings(200, 128, 42);
        let mut classifier = DensityClassifier::new(128, 10, 31, 1, false);
        classifier.fit(&embeddings);

        let analysis = classifier.multi_resolution_analysis(5);
        assert_eq!(analysis.n_items, 200);
        assert_eq!(analysis.num_bits, 10);
        assert_eq!(analysis.recovery_depth.len(), 200);
        assert_eq!(analysis.recovery_ratio.len(), 200);
        assert_eq!(analysis.buckets_per_depth.len(), 10);
        assert_eq!(analysis.mean_size_per_depth.len(), 10);
        assert_eq!(analysis.dense_threshold, 5);

        // Recovery depth should be 0..=num_bits+1
        for &depth in &analysis.recovery_depth {
            assert!(depth <= 11, "depth={} should be <= num_bits+1", depth);
        }

        // Recovery ratio should be non-negative for recovered items
        for (i, &ratio) in analysis.recovery_ratio.iter().enumerate() {
            if analysis.recovery_depth[i] <= 10 {
                assert!(ratio >= 0.0, "ratio={} should be >= 0 for recovered item", ratio);
            }
        }

        // Coarser depths should have fewer buckets
        for k in 1..analysis.buckets_per_depth.len() {
            assert!(
                analysis.buckets_per_depth[k] >= analysis.buckets_per_depth[k - 1],
                "buckets should increase with depth: {} at k={} vs {} at k={}",
                analysis.buckets_per_depth[k], k + 1,
                analysis.buckets_per_depth[k - 1], k,
            );
        }
    }

    #[test]
    fn test_stability_seeds() {
        let embeddings = generate_test_embeddings(50, 64, 42);

        // With 1 seed, all stability scores should be 1.0
        let mut clf1 = DensityClassifier::new(64, 10, 31, 1, false);
        clf1.fit(&embeddings);
        for &stab in clf1.get_stability_scores() {
            assert_eq!(stab, 1.0);
        }

        // With 5 seeds, stability scores should be in valid range [0, 1]
        let mut clf5 = DensityClassifier::new(64, 10, 31, 5, false);
        clf5.fit(&embeddings);
        for &stab in clf5.get_stability_scores() {
            assert!(stab >= 0.0 && stab <= 1.0, "stability={} out of range", stab);
        }

        // Report should include mean stability
        let report = clf5.report();
        assert!(report.mean_stability_score >= 0.0 && report.mean_stability_score <= 1.0);
    }

    #[test]
    fn test_bridge_persistence_analysis() {
        let embeddings = generate_test_embeddings(200, 128, 42);
        let mut classifier = DensityClassifier::new(128, 10, 31, 1, false);
        classifier.fit(&embeddings);

        let flat: Vec<f32> = embeddings.into_iter().flatten().collect();
        let bp = classifier.bridge_persistence_analysis(&flat, 200, 128, 0.8);

        assert_eq!(bp.n_items, 200);
        assert_eq!(bp.num_bits, 10);
        assert_eq!(bp.bridge_persistence.len(), 200);
        assert_eq!(bp.max_bridge_depth.len(), 200);
        assert_eq!(bp.min_bridge_depth.len(), 200);
        assert_eq!(bp.bridge_ratio.len(), 200);
        assert_eq!(bp.bridges_per_depth.len(), 10);
        assert_eq!(bp.pairs_per_depth.len(), 10);
        assert_eq!(bp.buckets_per_depth.len(), 10);
        assert_eq!(bp.relative_threshold, 0.8);

        // Persistence count should be <= num_bits for each item
        for &p in &bp.bridge_persistence {
            assert!(p <= 10, "persistence={} should be <= num_bits=10", p);
        }

        // max_bridge_depth should be >= min_bridge_depth for items that are connectors
        for i in 0..200 {
            if bp.bridge_persistence[i] > 0 {
                assert!(bp.max_bridge_depth[i] > 0);
                assert!(bp.min_bridge_depth[i] > 0);
                assert!(bp.max_bridge_depth[i] >= bp.min_bridge_depth[i]);
                assert!(bp.max_bridge_depth[i] <= 10);
                // bridge_ratio should be >= relative_threshold for connectors
                assert!(bp.bridge_ratio[i] >= 0.8,
                    "bridge_ratio={} should be >= threshold for connector", bp.bridge_ratio[i]);
            } else {
                assert_eq!(bp.max_bridge_depth[i], 0);
                assert_eq!(bp.min_bridge_depth[i], 0);
                assert_eq!(bp.bridge_ratio[i], 0.0);
            }
        }

        // Coarser depths (higher K) should have more unique buckets
        for k in 1..bp.buckets_per_depth.len() {
            assert!(
                bp.buckets_per_depth[k] >= bp.buckets_per_depth[k - 1],
                "buckets should increase with depth: {} at k={} vs {} at k={}",
                bp.buckets_per_depth[k], k + 1,
                bp.buckets_per_depth[k - 1], k,
            );
        }

        // persistent_pairs should have valid persistence counts
        for &(a, b, persistence, min_d, max_d) in &bp.persistent_pairs {
            assert!(a <= b, "pairs should be normalized: a={} <= b={}", a, b);
            assert!(persistence > 0 && persistence <= 10);
            assert!(min_d >= 1 && min_d <= 10);
            assert!(max_d >= min_d && max_d <= 10);
        }
    }
}
