next_plaid/

update.rs

//! Index update functionality for adding new documents.
//!
//! This module provides functions to incrementally update an existing PLAID index
//! with new documents, matching fast-plaid's behavior:
//! - Buffer mechanism for small updates
//! - Centroid expansion for outliers
//! - Cluster threshold updates

use std::cell::RefCell;
use std::collections::HashMap;
use std::fs;
use std::fs::File;
use std::io::{BufReader, BufWriter, Write};
use std::path::Path;

use serde::{Deserialize, Serialize};

use ndarray::{s, Array1, Array2, Axis};
use rayon::prelude::*;

use crate::codec::ResidualCodec;
use crate::error::Error;
use crate::error::Result;
use crate::index::Metadata;
use crate::kmeans::compute_kmeans;
use crate::kmeans::ComputeKmeansConfig;
use crate::utils::atomic_write_file;
use crate::utils::quantile;

/// Default batch size for processing documents (matches fast-plaid).
const DEFAULT_BATCH_SIZE: usize = 50_000;

/// Thread-local progress callback invoked with `(stage, message)` during an update.
type ProgressCallback = Box<dyn Fn(&str, &str)>;

thread_local! {
    static UPDATE_PROGRESS: RefCell<Option<ProgressCallback>> = const { RefCell::new(None) };
}

struct ProgressGuard;

impl Drop for ProgressGuard {
    fn drop(&mut self) {
        UPDATE_PROGRESS.with(|slot| {
            *slot.borrow_mut() = None;
        });
    }
}

/// Run an update operation with a thread-local progress callback.
///
/// The callback is intentionally thread-local because update work runs inside one blocking
/// worker thread; this avoids plumbing API-specific state through the public update config.
pub fn with_update_progress<F, R>(callback: F, operation: impl FnOnce() -> R) -> R
where
    F: Fn(&str, &str) + 'static,
{
    UPDATE_PROGRESS.with(|slot| {
        *slot.borrow_mut() = Some(Box::new(callback));
    });
    let _guard = ProgressGuard;
    operation()
}

fn emit_update_progress(stage: &str, message: &str) {
    UPDATE_PROGRESS.with(|slot| {
        if let Some(callback) = slot.borrow().as_ref() {
            callback(stage, message);
        }
    });
}

/// Configuration for index updates.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UpdateConfig {
    /// Batch size for processing documents (default: 50,000)
    pub batch_size: usize,
    /// Number of K-means iterations for centroid expansion (default: 4)
    pub kmeans_niters: usize,
    /// Max points per centroid for K-means (default: 256)
    pub max_points_per_centroid: usize,
    /// Number of samples for K-means (default: auto-calculated)
    pub n_samples_kmeans: Option<usize>,
    /// Random seed (default: 42)
    pub seed: u64,
    /// If index has fewer docs than this, rebuild from scratch (default: 999)
    pub start_from_scratch: usize,
    /// Buffer size before triggering centroid expansion (default: 100)
    pub buffer_size: usize,
    /// Force CPU execution for K-means even when CUDA feature is enabled.
    #[serde(default)]
    pub force_cpu: bool,
}

impl Default for UpdateConfig {
    fn default() -> Self {
        Self {
            batch_size: DEFAULT_BATCH_SIZE,
            kmeans_niters: 4,
            max_points_per_centroid: 256,
            n_samples_kmeans: None,
            seed: 42,
            start_from_scratch: crate::default_start_from_scratch(),
            buffer_size: 100,
            force_cpu: false,
        }
    }
}

impl UpdateConfig {
    /// Convert to ComputeKmeansConfig for centroid expansion.
    pub fn to_kmeans_config(&self) -> ComputeKmeansConfig {
        ComputeKmeansConfig {
            kmeans_niters: self.kmeans_niters,
            max_points_per_centroid: self.max_points_per_centroid,
            seed: self.seed,
            n_samples_kmeans: self.n_samples_kmeans,
            num_partitions: None,
            force_cpu: self.force_cpu,
        }
    }
}

// ============================================================================
// Buffer Management
// ============================================================================

/// Load buffered embeddings from buffer.npy.
///
/// Returns an empty vector if buffer.npy doesn't exist.
/// Uses buffer_lengths.json to split the flattened array back into per-document arrays.
pub fn load_buffer(index_path: &Path) -> Result<Vec<Array2<f32>>> {
    use ndarray_npy::ReadNpyExt;

    let buffer_path = index_path.join("buffer.npy");
    let lengths_path = index_path.join("buffer_lengths.json");

    if !buffer_path.exists() {
        return Ok(Vec::new());
    }

    // Load the flattened embeddings array
    let flat: Array2<f32> = match Array2::read_npy(File::open(&buffer_path)?) {
        Ok(arr) => arr,
        Err(_) => return Ok(Vec::new()),
    };

    // Load lengths to split back into per-document arrays
    if lengths_path.exists() {
        let lengths: Vec<i64> =
            serde_json::from_reader(BufReader::new(File::open(&lengths_path)?))?;

        let mut result = Vec::with_capacity(lengths.len());
        let mut offset = 0;

        for &len in &lengths {
            let len_usize = len as usize;
            if offset + len_usize > flat.nrows() {
                break;
            }
            let doc_emb = flat.slice(s![offset..offset + len_usize, ..]).to_owned();
            result.push(doc_emb);
            offset += len_usize;
        }

        return Ok(result);
    }

    // Fallback: if no lengths file, return as single document (legacy behavior)
    Ok(vec![flat])
}

/// Save embeddings to buffer.npy.
///
/// Also saves buffer_info.json with the number of documents for deletion tracking.
pub fn save_buffer(index_path: &Path, embeddings: &[Array2<f32>]) -> Result<()> {
    use ndarray_npy::WriteNpyExt;

    let buffer_path = index_path.join("buffer.npy");

    // For simplicity, concatenate all embeddings into one array
    // and store the lengths separately
    if embeddings.is_empty() {
        return Ok(());
    }

    let dim = embeddings[0].ncols();
    let total_rows: usize = embeddings.iter().map(|e| e.nrows()).sum();

    let mut flat = Array2::<f32>::zeros((total_rows, dim));
    let mut offset = 0;
    let mut lengths = Vec::new();

    for emb in embeddings {
        let n = emb.nrows();
        flat.slice_mut(s![offset..offset + n, ..]).assign(emb);
        lengths.push(n as i64);
        offset += n;
    }

    atomic_write_file(&buffer_path, |file| {
        flat.write_npy(file)?;
        Ok(())
    })?;

    // Save lengths
    let lengths_path = index_path.join("buffer_lengths.json");
    atomic_write_file(&lengths_path, |file| {
        let mut writer = BufWriter::new(file);
        serde_json::to_writer(&mut writer, &lengths)?;
        writer.flush()?;
        Ok(())
    })?;

    // Save buffer info for deletion tracking (number of documents)
    let info_path = index_path.join("buffer_info.json");
    let buffer_info = serde_json::json!({ "num_docs": embeddings.len() });
    atomic_write_file(&info_path, |file| {
        let mut writer = BufWriter::new(file);
        serde_json::to_writer(&mut writer, &buffer_info)?;
        writer.flush()?;
        Ok(())
    })?;

    Ok(())
}

/// Load buffer info (number of buffered documents).
///
/// Returns 0 if buffer_info.json doesn't exist.
pub fn load_buffer_info(index_path: &Path) -> Result<usize> {
    let info_path = index_path.join("buffer_info.json");
    if !info_path.exists() {
        return Ok(0);
    }

    let info: serde_json::Value = serde_json::from_reader(BufReader::new(File::open(&info_path)?))?;

    Ok(info.get("num_docs").and_then(|v| v.as_u64()).unwrap_or(0) as usize)
}

/// Clear buffer files.
pub fn clear_buffer(index_path: &Path) -> Result<()> {
    let buffer_path = index_path.join("buffer.npy");
    let lengths_path = index_path.join("buffer_lengths.json");
    let info_path = index_path.join("buffer_info.json");

    if buffer_path.exists() {
        fs::remove_file(&buffer_path)?;
    }
    if lengths_path.exists() {
        fs::remove_file(&lengths_path)?;
    }
    if info_path.exists() {
        fs::remove_file(&info_path)?;
    }

    Ok(())
}

/// Load embeddings stored for rebuild (embeddings.npy + embeddings_lengths.json).
///
/// This function loads raw embeddings that were saved for start-from-scratch rebuilds.
/// The embeddings are stored in a flat 2D array with a separate lengths file.
pub fn load_embeddings_npy(index_path: &Path) -> Result<Vec<Array2<f32>>> {
    use ndarray_npy::ReadNpyExt;

    let emb_path = index_path.join("embeddings.npy");
    let lengths_path = index_path.join("embeddings_lengths.json");

    if !emb_path.exists() {
        return Ok(Vec::new());
    }

    // Load flat embeddings array
    let flat: Array2<f32> = Array2::read_npy(File::open(&emb_path)?)?;

    // Load lengths to split back into per-document arrays
    if lengths_path.exists() {
        let lengths: Vec<i64> =
            serde_json::from_reader(BufReader::new(File::open(&lengths_path)?))?;

        let mut result = Vec::with_capacity(lengths.len());
        let mut offset = 0;

        for &len in &lengths {
            let len_usize = len as usize;
            if offset + len_usize > flat.nrows() {
                break;
            }
            let doc_emb = flat.slice(s![offset..offset + len_usize, ..]).to_owned();
            result.push(doc_emb);
            offset += len_usize;
        }

        return Ok(result);
    }

    // Fallback: if no lengths file, return as single document
    Ok(vec![flat])
}

/// Save embeddings for potential rebuild (start-from-scratch mode).
///
/// Stores embeddings in embeddings.npy (flat array) + embeddings_lengths.json.
/// This matches fast-plaid's behavior of storing raw embeddings when the index
/// is below the start_from_scratch threshold.
pub fn save_embeddings_npy(index_path: &Path, embeddings: &[Array2<f32>]) -> Result<()> {
    use ndarray_npy::WriteNpyExt;

    if embeddings.is_empty() {
        return Ok(());
    }

    let dim = embeddings[0].ncols();
    let total_rows: usize = embeddings.iter().map(|e| e.nrows()).sum();

    let mut flat = Array2::<f32>::zeros((total_rows, dim));
    let mut offset = 0;
    let mut lengths = Vec::with_capacity(embeddings.len());

    for emb in embeddings {
        let n = emb.nrows();
        flat.slice_mut(s![offset..offset + n, ..]).assign(emb);
        lengths.push(n as i64);
        offset += n;
    }

    // Save flat embeddings
    let emb_path = index_path.join("embeddings.npy");
    atomic_write_file(&emb_path, |file| {
        flat.write_npy(file)?;
        Ok(())
    })?;

    // Save lengths for reconstruction
    let lengths_path = index_path.join("embeddings_lengths.json");
    atomic_write_file(&lengths_path, |file| {
        let mut writer = BufWriter::new(file);
        serde_json::to_writer(&mut writer, &lengths)?;
        writer.flush()?;
        Ok(())
    })?;

    Ok(())
}

/// Clear embeddings.npy and embeddings_lengths.json.
pub fn clear_embeddings_npy(index_path: &Path) -> Result<()> {
    let emb_path = index_path.join("embeddings.npy");
    let lengths_path = index_path.join("embeddings_lengths.json");

    if emb_path.exists() {
        fs::remove_file(&emb_path)?;
    }
    if lengths_path.exists() {
        fs::remove_file(&lengths_path)?;
    }
    Ok(())
}

/// Check if embeddings.npy exists for start-from-scratch mode.
pub fn embeddings_npy_exists(index_path: &Path) -> bool {
    index_path.join("embeddings.npy").exists()
}

// ============================================================================
// Cluster Threshold Management
// ============================================================================

/// Load cluster threshold from cluster_threshold.npy.
pub fn load_cluster_threshold(index_path: &Path) -> Result<f32> {
    use ndarray_npy::ReadNpyExt;

    let thresh_path = index_path.join("cluster_threshold.npy");
    if !thresh_path.exists() {
        return Err(Error::Update("cluster_threshold.npy not found".into()));
    }

    let arr: Array1<f32> = Array1::read_npy(File::open(&thresh_path)?)?;
    Ok(arr[0])
}

/// Update cluster_threshold.npy with weighted average.
pub fn update_cluster_threshold(
    index_path: &Path,
    new_residual_norms: &Array1<f32>,
    old_total_embeddings: usize,
) -> Result<()> {
    use ndarray_npy::{ReadNpyExt, WriteNpyExt};

    let new_count = new_residual_norms.len();
    if new_count == 0 {
        return Ok(());
    }

    let new_threshold = quantile(new_residual_norms, 0.75);

    let thresh_path = index_path.join("cluster_threshold.npy");
    let final_threshold = if thresh_path.exists() {
        let old_arr: Array1<f32> = Array1::read_npy(File::open(&thresh_path)?)?;
        let old_threshold = old_arr[0];
        let total = old_total_embeddings + new_count;
        (old_threshold * old_total_embeddings as f32 + new_threshold * new_count as f32)
            / total as f32
    } else {
        new_threshold
    };

    atomic_write_file(&thresh_path, |file| {
        Array1::from_vec(vec![final_threshold]).write_npy(file)?;
        Ok(())
    })?;

    Ok(())
}

// ============================================================================
// Centroid Expansion
// ============================================================================

const OUTLIER_CENTROID_TILE: usize = 8;
const OUTLIER_EMBEDDING_TILE: usize = 64;
const OUTLIER_BLOCKS_MIN_LEN: usize = 4;
const OUTLIER_THRESHOLD_RECHECK_REL_EPS: f32 = 1e-5;

#[inline]
fn squared_norm(row: &[f32]) -> f32 {
    let mut sum0 = 0.0f32;
    let mut sum1 = 0.0f32;
    let mut sum2 = 0.0f32;
    let mut sum3 = 0.0f32;

    let mut i = 0;
    while i + 4 <= row.len() {
        sum0 += row[i] * row[i];
        sum1 += row[i + 1] * row[i + 1];
        sum2 += row[i + 2] * row[i + 2];
        sum3 += row[i + 3] * row[i + 3];
        i += 4;
    }

    let mut total = sum0 + sum1 + sum2 + sum3;
    while i < row.len() {
        total += row[i] * row[i];
        i += 1;
    }

    total
}

/// Recompute min L2² distance in f64 for borderline candidates.
/// The f32 tiled loop can produce rounding errors near the threshold;
/// promoting to f64 here eliminates false positives without slowing
/// down the fast path (only called for ~1% of embeddings).
#[inline]
fn min_distance_sq_precise(row: &[f32], centroids_flat: &[f32], dim: usize) -> f32 {
    let mut min_dist_sq = f32::INFINITY;

    for centroid in centroids_flat.chunks_exact(dim) {
        let mut dist_sq = 0.0f64;
        let mut d = 0;
        while d < dim {
            let diff = row[d] as f64 - centroid[d] as f64;
            dist_sq += diff * diff;
            d += 1;
        }

        min_dist_sq = min_dist_sq.min(dist_sq as f32);
    }

    min_dist_sq
}

/// Find outliers with a blocked raw-slice kernel over centroid tiles.
///
/// This avoids materializing the intermediate [batch, num_centroids] scores matrix,
/// removes iterator overhead from the hot loop, and parallelizes across embedding blocks.
/// Find embeddings whose minimum L2² distance to any centroid exceeds `threshold_sq`.
///
/// Uses a tiled blocked loop instead of a full [n, k] similarity matrix:
///   - Outer: blocks of OUTLIER_EMBEDDING_TILE embeddings (parallelized via rayon)
///   - Inner: tiles of OUTLIER_CENTROID_TILE centroids
///
/// This avoids allocating O(n*k) memory and keeps the working set in L1/L2 cache.
/// The dot-product loop (`dim_idx`) is the innermost to maximize sequential memory
/// access across the embedding dimension. Borderline results (within 1% of the
/// threshold) are rechecked in f64 to avoid false positives from f32 rounding.
#[allow(clippy::needless_range_loop)]
fn find_outliers(
    flat_embeddings: &Array2<f32>,
    centroids: &Array2<f32>,
    threshold_sq: f32,
) -> Vec<usize> {
    let n = flat_embeddings.nrows();
    let k = centroids.nrows();
    let dim = flat_embeddings.ncols();

    if n == 0 || k == 0 {
        return Vec::new();
    }

    let embeddings_owned;
    let embeddings_flat = if let Some(slice) = flat_embeddings.as_slice_memory_order() {
        slice
    } else {
        embeddings_owned = flat_embeddings.as_standard_layout().to_owned();
        embeddings_owned
            .as_slice_memory_order()
            .expect("standard-layout embeddings should be contiguous")
    };
    let centroids_owned;
    let centroids_flat = if let Some(slice) = centroids.as_slice_memory_order() {
        slice
    } else {
        centroids_owned = centroids.as_standard_layout().to_owned();
        centroids_owned
            .as_slice_memory_order()
            .expect("standard-layout centroids should be contiguous")
    };

    let centroid_norms_sq: Vec<f32> = centroids_flat
        .par_chunks_exact(dim)
        .map(squared_norm)
        .collect();

    let row_stride = dim * OUTLIER_EMBEDDING_TILE;
    embeddings_flat
        .par_chunks(row_stride)
        .with_min_len(OUTLIER_BLOCKS_MIN_LEN)
        .enumerate()
        .flat_map_iter(|(block_idx, rows_block)| {
            let row_count = rows_block.len() / dim;
            let row_offset = block_idx * OUTLIER_EMBEDDING_TILE;

            let mut min_dists = vec![f32::INFINITY; row_count];
            let emb_norms: Vec<f32> = rows_block.chunks_exact(dim).map(squared_norm).collect();

            let mut centroid_idx = 0;
            while centroid_idx + OUTLIER_CENTROID_TILE <= k {
                let centroid_bases: [usize; OUTLIER_CENTROID_TILE] =
                    std::array::from_fn(|j| (centroid_idx + j) * dim);

                let mut dots = [[0.0f32; OUTLIER_CENTROID_TILE]; OUTLIER_EMBEDDING_TILE];

                let mut dim_idx = 0;
                while dim_idx < dim {
                    let centroid_vals: [f32; OUTLIER_CENTROID_TILE] =
                        std::array::from_fn(|j| centroids_flat[centroid_bases[j] + dim_idx]);

                    for row_idx in 0..row_count {
                        let x = rows_block[row_idx * dim + dim_idx];
                        for j in 0..OUTLIER_CENTROID_TILE {
                            dots[row_idx][j] += x * centroid_vals[j];
                        }
                    }

                    dim_idx += 1;
                }

                for row_idx in 0..row_count {
                    let emb_norm_sq = emb_norms[row_idx];
                    for j in 0..OUTLIER_CENTROID_TILE {
                        let dist_sq = emb_norm_sq + centroid_norms_sq[centroid_idx + j]
                            - 2.0 * dots[row_idx][j];
                        min_dists[row_idx] = min_dists[row_idx].min(dist_sq);
                    }
                }

                centroid_idx += OUTLIER_CENTROID_TILE;
            }

            while centroid_idx < k {
                let centroid = &centroids_flat[centroid_idx * dim..(centroid_idx + 1) * dim];
                for row_idx in 0..row_count {
                    let row = &rows_block[row_idx * dim..(row_idx + 1) * dim];
                    let mut dot = 0.0f32;
                    let mut dim_idx = 0;
                    while dim_idx < dim {
                        dot += row[dim_idx] * centroid[dim_idx];
                        dim_idx += 1;
                    }

                    let dist_sq = emb_norms[row_idx] + centroid_norms_sq[centroid_idx] - 2.0 * dot;
                    min_dists[row_idx] = min_dists[row_idx].min(dist_sq);
                }

                centroid_idx += 1;
            }

            min_dists
                .into_iter()
                .enumerate()
                .filter_map(move |(row_idx, min_dist_sq)| {
                    let final_min_dist_sq = if (min_dist_sq - threshold_sq).abs()
                        <= threshold_sq.abs().max(1.0) * OUTLIER_THRESHOLD_RECHECK_REL_EPS
                    {
                        let row = &rows_block[row_idx * dim..(row_idx + 1) * dim];
                        min_distance_sq_precise(row, centroids_flat, dim)
                    } else {
                        min_dist_sq
                    };

                    (final_min_dist_sq > threshold_sq).then_some(row_offset + row_idx)
                })
        })
        .collect()
}

/// Expand centroids by clustering embeddings far from existing centroids.
///
/// This implements fast-plaid's update_centroids() function:
/// 1. Flatten all new embeddings
/// 2. Find outliers (distance > cluster_threshold²)
/// 3. Cluster outliers to get new centroids
/// 4. Append new centroids to centroids.npy
/// 5. Extend ivf_lengths.npy with zeros
/// 6. Update metadata.json num_partitions
///
/// Returns the number of new centroids added.
pub fn update_centroids(
    index_path: &Path,
    new_embeddings: &[Array2<f32>],
    cluster_threshold: f32,
    config: &UpdateConfig,
) -> Result<usize> {
    use ndarray_npy::{ReadNpyExt, WriteNpyExt};

    let centroids_path = index_path.join("centroids.npy");
    if !centroids_path.exists() {
        return Ok(0);
    }

    // Load existing centroids
    let existing_centroids: Array2<f32> = Array2::read_npy(File::open(&centroids_path)?)?;

    // Flatten all new embeddings
    let dim = existing_centroids.ncols();
    let total_tokens: usize = new_embeddings.iter().map(|e| e.nrows()).sum();

    if total_tokens == 0 {
        return Ok(0);
    }

    let mut flat_embeddings = Array2::<f32>::zeros((total_tokens, dim));
    let mut offset = 0;

    for emb in new_embeddings {
        let n = emb.nrows();
        flat_embeddings
            .slice_mut(s![offset..offset + n, ..])
            .assign(emb);
        offset += n;
    }

    // Find outliers
    emit_update_progress(
        "centroid_expansion",
        "finding embeddings outside existing centroids",
    );
    let threshold_sq = cluster_threshold * cluster_threshold;
    let outlier_indices = find_outliers(&flat_embeddings, &existing_centroids, threshold_sq);

    let num_outliers = outlier_indices.len();
    if num_outliers == 0 {
        return Ok(0);
    }

    // Extract outlier embeddings
    let mut outliers = Array2::<f32>::zeros((num_outliers, dim));
    for (i, &idx) in outlier_indices.iter().enumerate() {
        outliers.row_mut(i).assign(&flat_embeddings.row(idx));
    }

    // Compute number of new centroids
    // k_update = max(1, ceil(num_outliers / max_points_per_centroid) * 4)
    let target_k =
        ((num_outliers as f64 / config.max_points_per_centroid as f64).ceil() as usize).max(1) * 4;
    let k_update = target_k.min(num_outliers); // Can't have more centroids than points

    // Cluster outliers to get new centroids
    let kmeans_config = ComputeKmeansConfig {
        kmeans_niters: config.kmeans_niters,
        max_points_per_centroid: config.max_points_per_centroid,
        seed: config.seed,
        n_samples_kmeans: config.n_samples_kmeans,
        num_partitions: Some(k_update),
        force_cpu: config.force_cpu,
    };

    // Convert outliers to vector of single-token "documents" for compute_kmeans
    let outlier_docs: Vec<Array2<f32>> = outlier_indices
        .iter()
        .map(|&idx| flat_embeddings.slice(s![idx..idx + 1, ..]).to_owned())
        .collect();

    emit_update_progress("kmeans", "clustering outlier embeddings");
    let new_centroids = compute_kmeans(&outlier_docs, &kmeans_config)?;
    let k_new = new_centroids.nrows();

    // Concatenate centroids
    let new_num_centroids = existing_centroids.nrows() + k_new;
    let mut final_centroids = Array2::<f32>::zeros((new_num_centroids, dim));
    final_centroids
        .slice_mut(s![..existing_centroids.nrows(), ..])
        .assign(&existing_centroids);
    final_centroids
        .slice_mut(s![existing_centroids.nrows().., ..])
        .assign(&new_centroids);

    // Save updated centroids
    emit_update_progress("index_write", "writing updated centroids");
    atomic_write_file(&centroids_path, |file| {
        final_centroids.write_npy(file)?;
        Ok(())
    })?;

    // Extend ivf_lengths.npy with zeros for new centroids
    let ivf_lengths_path = index_path.join("ivf_lengths.npy");
    if ivf_lengths_path.exists() {
        let old_lengths: Array1<i32> = Array1::read_npy(File::open(&ivf_lengths_path)?)?;
        let mut new_lengths = Array1::<i32>::zeros(new_num_centroids);
        new_lengths
            .slice_mut(s![..old_lengths.len()])
            .assign(&old_lengths);
        atomic_write_file(&ivf_lengths_path, |file| {
            new_lengths.write_npy(file)?;
            Ok(())
        })?;
    }

    // Update metadata.json num_partitions
    let meta_path = index_path.join("metadata.json");
    if meta_path.exists() {
        let mut meta: serde_json::Value =
            serde_json::from_reader(BufReader::new(File::open(&meta_path)?))?;

        if let Some(obj) = meta.as_object_mut() {
            obj.insert("num_partitions".to_string(), new_num_centroids.into());
        }

        atomic_write_file(&meta_path, |file| {
            let mut writer = BufWriter::new(file);
            serde_json::to_writer_pretty(&mut writer, &meta)?;
            writer.flush()?;
            Ok(())
        })?;
    }

    Ok(k_new)
}

// ============================================================================
// Low-Level Index Update
// ============================================================================

/// Update an existing index with new documents.
///
/// # Arguments
///
/// * `embeddings` - List of new document embeddings, each of shape `[num_tokens, dim]`
/// * `index_path` - Path to the existing index directory
/// * `codec` - The loaded ResidualCodec for compression
/// * `batch_size` - Optional batch size for processing (default: 50,000)
/// * `update_threshold` - Whether to update the cluster threshold
/// * `force_cpu` - Force CPU execution even when CUDA is available
///
/// # Returns
///
/// The number of new documents added
pub fn update_index(
    embeddings: &[Array2<f32>],
    index_path: &str,
    codec: &ResidualCodec,
    batch_size: Option<usize>,
    update_threshold: bool,
    force_cpu: bool,
) -> Result<usize> {
    emit_update_progress("index_write", "writing index chunks");
    let batch_size = batch_size.unwrap_or(DEFAULT_BATCH_SIZE);
    let index_dir = Path::new(index_path);

    // Load existing metadata (infers num_documents from doclens if not present)
    let metadata_path = index_dir.join("metadata.json");
    let metadata = Metadata::load_from_path(index_dir)?;

    let num_existing_chunks = metadata.num_chunks;
    let old_num_documents = metadata.num_documents;
    let old_total_embeddings = metadata.num_embeddings;
    let num_centroids = codec.num_centroids();
    let embedding_dim = codec.embedding_dim();
    let nbits = metadata.nbits;

    // Determine starting chunk index
    let mut start_chunk_idx = num_existing_chunks;
    let mut append_to_last = false;
    let mut current_emb_offset = old_total_embeddings;

    // Check if we should append to the last chunk (if it has < 2000 documents)
    if start_chunk_idx > 0 {
        let last_idx = start_chunk_idx - 1;
        let last_meta_path = index_dir.join(format!("{}.metadata.json", last_idx));

        if last_meta_path.exists() {
            let last_meta: serde_json::Value =
                serde_json::from_reader(BufReader::new(File::open(&last_meta_path).map_err(
                    |e| Error::IndexLoad(format!("Failed to open chunk metadata: {}", e)),
                )?))?;

            if let Some(nd) = last_meta.get("num_documents").and_then(|x| x.as_u64()) {
                if nd < 2000 {
                    start_chunk_idx = last_idx;
                    append_to_last = true;

                    if let Some(off) = last_meta.get("embedding_offset").and_then(|x| x.as_u64()) {
                        current_emb_offset = off as usize;
                    } else {
                        let embs_in_last = last_meta
                            .get("num_embeddings")
                            .and_then(|x| x.as_u64())
                            .unwrap_or(0) as usize;
                        current_emb_offset = old_total_embeddings - embs_in_last;
                    }
                }
            }
        }
    }

    // Process new documents
    let num_new_documents = embeddings.len();
    let n_new_chunks = (num_new_documents as f64 / batch_size as f64).ceil() as usize;

    let mut new_codes_accumulated: Vec<Vec<usize>> = Vec::new();
    let mut new_doclens_accumulated: Vec<i64> = Vec::new();
    let mut all_residual_norms: Vec<f32> = Vec::new();

    let packed_dim = embedding_dim * nbits / 8;

    for i in 0..n_new_chunks {
        let global_chunk_idx = start_chunk_idx + i;
        let chk_offset = i * batch_size;
        let chk_end = (chk_offset + batch_size).min(num_new_documents);
        let chunk_docs = &embeddings[chk_offset..chk_end];

        // Collect document lengths
        let mut chk_doclens: Vec<i64> = chunk_docs.iter().map(|d| d.nrows() as i64).collect();
        let total_tokens: usize = chk_doclens.iter().sum::<i64>() as usize;

        // Concatenate all embeddings in the chunk for batch processing
        let mut batch_embeddings = ndarray::Array2::<f32>::zeros((total_tokens, embedding_dim));
        let mut offset = 0;
        for doc in chunk_docs {
            let n = doc.nrows();
            batch_embeddings
                .slice_mut(s![offset..offset + n, ..])
                .assign(doc);
            offset += n;
        }

        // BATCH: Compress all embeddings at once
        // Use CPU-only version when force_cpu is set to avoid CUDA initialization overhead
        let batch_codes = if force_cpu {
            codec.compress_into_codes_cpu(&batch_embeddings)
        } else {
            codec.compress_into_codes(&batch_embeddings)
        };

        // BATCH: Compute residuals using parallel subtraction
        let mut batch_residuals = batch_embeddings;
        {
            let centroids = &codec.centroids;
            batch_residuals
                .axis_iter_mut(Axis(0))
                .into_par_iter()
                .zip(batch_codes.as_slice().unwrap().par_iter())
                .for_each(|(mut row, &code)| {
                    let centroid = centroids.row(code);
                    row.iter_mut()
                        .zip(centroid.iter())
                        .for_each(|(r, c)| *r -= c);
                });
        }

        // Collect residual norms if updating threshold
        if update_threshold {
            for row in batch_residuals.axis_iter(Axis(0)) {
                let norm = row.dot(&row).sqrt();
                all_residual_norms.push(norm);
            }
        }

        // BATCH: Quantize all residuals at once
        let batch_packed = codec.quantize_residuals(&batch_residuals)?;

        // Convert to lists for chunk saving
        let mut chk_codes_list: Vec<usize> = batch_codes.iter().copied().collect();
        let mut chk_residuals_list: Vec<u8> = batch_packed.iter().copied().collect();

        // Split codes back into per-document arrays for IVF building
        let mut code_offset = 0;
        for &len in &chk_doclens {
            let len_usize = len as usize;
            let codes: Vec<usize> = batch_codes
                .slice(s![code_offset..code_offset + len_usize])
                .iter()
                .copied()
                .collect();
            new_codes_accumulated.push(codes);
            new_doclens_accumulated.push(len);
            code_offset += len_usize;
        }

        // Handle appending to last chunk
        if i == 0 && append_to_last {
            use ndarray_npy::ReadNpyExt;

            let old_doclens_path = index_dir.join(format!("doclens.{}.json", global_chunk_idx));

            if old_doclens_path.exists() {
                let old_doclens: Vec<i64> =
                    serde_json::from_reader(BufReader::new(File::open(&old_doclens_path)?))?;

                let old_codes_path = index_dir.join(format!("{}.codes.npy", global_chunk_idx));
                let old_residuals_path =
                    index_dir.join(format!("{}.residuals.npy", global_chunk_idx));

                let old_codes: Array1<i64> = Array1::read_npy(File::open(&old_codes_path)?)?;
                let old_residuals: Array2<u8> = Array2::read_npy(File::open(&old_residuals_path)?)?;

                // Prepend old data
                let mut combined_codes: Vec<usize> =
                    old_codes.iter().map(|&x| x as usize).collect();
                combined_codes.extend(chk_codes_list);
                chk_codes_list = combined_codes;

                let mut combined_residuals: Vec<u8> = old_residuals.iter().copied().collect();
                combined_residuals.extend(chk_residuals_list);
                chk_residuals_list = combined_residuals;

                let mut combined_doclens = old_doclens;
                combined_doclens.extend(chk_doclens);
                chk_doclens = combined_doclens;
            }
        }

        // Save chunk data
        {
            use ndarray_npy::WriteNpyExt;

            let codes_arr: Array1<i64> = chk_codes_list.iter().map(|&x| x as i64).collect();
            let codes_path = index_dir.join(format!("{}.codes.npy", global_chunk_idx));
            atomic_write_file(&codes_path, |file| {
                codes_arr.write_npy(file)?;
                Ok(())
            })?;

            let num_tokens = chk_codes_list.len();
            let residuals_arr =
                Array2::from_shape_vec((num_tokens, packed_dim), chk_residuals_list)
                    .map_err(|e| Error::Shape(format!("Failed to reshape residuals: {}", e)))?;
            let residuals_path = index_dir.join(format!("{}.residuals.npy", global_chunk_idx));
            atomic_write_file(&residuals_path, |file| {
                residuals_arr.write_npy(file)?;
                Ok(())
            })?;
        }

        // Save doclens
        let doclens_path = index_dir.join(format!("doclens.{}.json", global_chunk_idx));
        atomic_write_file(&doclens_path, |file| {
            let mut writer = BufWriter::new(file);
            serde_json::to_writer(&mut writer, &chk_doclens)?;
            writer.flush()?;
            Ok(())
        })?;

        // Save chunk metadata
        let chk_meta = serde_json::json!({
            "num_documents": chk_doclens.len(),
            "num_embeddings": chk_codes_list.len(),
            "embedding_offset": current_emb_offset,
        });
        current_emb_offset += chk_codes_list.len();

        let meta_path = index_dir.join(format!("{}.metadata.json", global_chunk_idx));
        atomic_write_file(&meta_path, |file| {
            let mut writer = BufWriter::new(file);
            serde_json::to_writer_pretty(&mut writer, &chk_meta)?;
            writer.flush()?;
            Ok(())
        })?;
    }

    // Update cluster threshold if requested
    if update_threshold && !all_residual_norms.is_empty() {
        let norms = Array1::from_vec(all_residual_norms);
        update_cluster_threshold(index_dir, &norms, old_total_embeddings)?;
    }

    // Build new partial IVF
    let mut partition_pids_map: HashMap<usize, Vec<i64>> = HashMap::new();

    for (pid_counter, doc_codes) in (old_num_documents as i64..).zip(new_codes_accumulated.iter()) {
        for &code in doc_codes {
            partition_pids_map
                .entry(code)
                .or_default()
                .push(pid_counter);
        }
    }

    // Load old IVF and merge
    {
        use ndarray_npy::{ReadNpyExt, WriteNpyExt};

        let ivf_path = index_dir.join("ivf.npy");
        let ivf_lengths_path = index_dir.join("ivf_lengths.npy");

        let old_ivf: Array1<i64> = if ivf_path.exists() {
            Array1::read_npy(File::open(&ivf_path)?)?
        } else {
            Array1::zeros(0)
        };

        let old_ivf_lengths: Array1<i32> = if ivf_lengths_path.exists() {
            Array1::read_npy(File::open(&ivf_lengths_path)?)?
        } else {
            Array1::zeros(num_centroids)
        };

        // Compute old offsets
        let mut old_offsets = vec![0i64];
        for &len in old_ivf_lengths.iter() {
            old_offsets.push(old_offsets.last().unwrap() + len as i64);
        }

        // Merge IVF
        let mut new_ivf_data: Vec<i64> = Vec::new();
        let mut new_ivf_lengths: Vec<i32> = Vec::with_capacity(num_centroids);

        for centroid_id in 0..num_centroids {
            // Get old PIDs for this centroid
            let old_start = old_offsets[centroid_id] as usize;
            let old_len = if centroid_id < old_ivf_lengths.len() {
                old_ivf_lengths[centroid_id] as usize
            } else {
                0
            };

            let mut pids: Vec<i64> = if old_len > 0 && old_start + old_len <= old_ivf.len() {
                old_ivf.slice(s![old_start..old_start + old_len]).to_vec()
            } else {
                Vec::new()
            };

            // Add new PIDs
            if let Some(new_pids) = partition_pids_map.get(&centroid_id) {
                pids.extend(new_pids);
            }

            // Deduplicate and sort
            pids.sort_unstable();
            pids.dedup();

            new_ivf_lengths.push(pids.len() as i32);
            new_ivf_data.extend(pids);
        }

        // Save updated IVF
        let new_ivf = Array1::from_vec(new_ivf_data);
        atomic_write_file(&ivf_path, |file| {
            new_ivf.write_npy(file)?;
            Ok(())
        })?;

        let new_lengths = Array1::from_vec(new_ivf_lengths);
        atomic_write_file(&ivf_lengths_path, |file| {
            new_lengths.write_npy(file)?;
            Ok(())
        })?;
    }

    // Update global metadata
    let new_total_chunks = start_chunk_idx + n_new_chunks;
    let new_tokens_count: i64 = new_doclens_accumulated.iter().sum();
    let num_embeddings = old_total_embeddings + new_tokens_count as usize;
    let total_num_documents = old_num_documents + num_new_documents;

    let new_avg_doclen = if total_num_documents > 0 {
        let old_sum = metadata.avg_doclen * old_num_documents as f64;
        (old_sum + new_tokens_count as f64) / total_num_documents as f64
    } else {
        0.0
    };

    let new_metadata = Metadata {
        num_chunks: new_total_chunks,
        nbits,
        num_partitions: num_centroids,
        num_embeddings,
        avg_doclen: new_avg_doclen,
        num_documents: total_num_documents,
        embedding_dim,
        next_plaid_compatible: true,
    };

    emit_update_progress("metadata_write", "writing index metadata");
    atomic_write_file(&metadata_path, |file| {
        let mut writer = BufWriter::new(file);
        serde_json::to_writer_pretty(&mut writer, &new_metadata)?;
        writer.flush()?;
        Ok(())
    })?;

    // Clear merged files to force regeneration on next load.
    // This ensures the merged files are consistent with the new chunk data.
    crate::mmap::clear_merged_files(index_dir)?;

    Ok(num_new_documents)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    };

    #[test]
    fn test_update_config_default() {
        let config = UpdateConfig::default();
        assert_eq!(config.batch_size, 50_000);
        assert_eq!(config.buffer_size, 100);
        assert_eq!(
            config.start_from_scratch,
            crate::default_start_from_scratch()
        );
    }

    #[test]
    fn test_update_progress_clears_after_panic() {
        let stale_count = Arc::new(AtomicUsize::new(0));
        let stale_count_for_callback = Arc::clone(&stale_count);
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            with_update_progress(
                move |_, _| {
                    stale_count_for_callback.fetch_add(1, Ordering::SeqCst);
                },
                || panic!("forced progress panic"),
            );
        }));

        assert!(result.is_err());

        emit_update_progress("after_panic", "must not hit stale callback");
        assert_eq!(stale_count.load(Ordering::SeqCst), 0);

        let fresh_count = Arc::new(AtomicUsize::new(0));
        let fresh_count_for_callback = Arc::clone(&fresh_count);
        with_update_progress(
            move |_, _| {
                fresh_count_for_callback.fetch_add(1, Ordering::SeqCst);
            },
            || emit_update_progress("fresh", "fresh callback"),
        );
        assert_eq!(fresh_count.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_find_outliers() {
        // Create centroids at (0,0), (1,1)
        let centroids = Array2::from_shape_vec((2, 2), vec![0.0, 0.0, 1.0, 1.0]).unwrap();

        // Create embeddings: one close to (0,0), one close to (1,1), one far away at (5,5)
        let embeddings =
            Array2::from_shape_vec((3, 2), vec![0.1, 0.1, 0.9, 0.9, 5.0, 5.0]).unwrap();

        // Threshold of 1.0 squared = 1.0
        let outliers = find_outliers(&embeddings, &centroids, 1.0);

        // Only the point at (5,5) should be an outlier
        assert_eq!(outliers.len(), 1);
        assert_eq!(outliers[0], 2);
    }

    #[test]
    fn test_buffer_roundtrip() {
        use tempfile::TempDir;

        let dir = TempDir::new().unwrap();

        // Create 3 documents with different numbers of embeddings
        let embeddings = vec![
            Array2::from_shape_vec((3, 4), (0..12).map(|x| x as f32).collect()).unwrap(),
            Array2::from_shape_vec((2, 4), (12..20).map(|x| x as f32).collect()).unwrap(),
            Array2::from_shape_vec((5, 4), (20..40).map(|x| x as f32).collect()).unwrap(),
        ];

        // Save buffer
        save_buffer(dir.path(), &embeddings).unwrap();

        // Load buffer and verify we get 3 documents, not 1
        let loaded = load_buffer(dir.path()).unwrap();

        assert_eq!(loaded.len(), 3, "Should have 3 documents, not 1");
        assert_eq!(loaded[0].nrows(), 3, "First doc should have 3 rows");
        assert_eq!(loaded[1].nrows(), 2, "Second doc should have 2 rows");
        assert_eq!(loaded[2].nrows(), 5, "Third doc should have 5 rows");

        // Verify content matches
        assert_eq!(loaded[0], embeddings[0]);
        assert_eq!(loaded[1], embeddings[1]);
        assert_eq!(loaded[2], embeddings[2]);
    }

    #[test]
    fn test_buffer_info_matches_buffer_len() {
        use tempfile::TempDir;

        let dir = TempDir::new().unwrap();

        // Create 5 documents
        let embeddings: Vec<Array2<f32>> = (0..5)
            .map(|i| {
                let rows = i + 2; // 2, 3, 4, 5, 6 rows
                Array2::from_shape_fn((rows, 4), |(r, c)| (r * 4 + c) as f32)
            })
            .collect();

        save_buffer(dir.path(), &embeddings).unwrap();

        // Verify buffer_info.json matches actual document count
        let info_count = load_buffer_info(dir.path()).unwrap();
        let loaded = load_buffer(dir.path()).unwrap();

        assert_eq!(info_count, 5, "buffer_info should report 5 docs");
        assert_eq!(
            loaded.len(),
            5,
            "load_buffer should return 5 docs to match buffer_info"
        );
    }
}