feox-ann 0.1.0

use std::thread;

use crate::graph::{Node, ScoredNode, SearchScratch};
use crate::math::{dot, normalize_in_place};
use crate::{AnnIndex, Result};

const BULK_BATCH: usize = 256;
const MAX_BULK_THREADS: usize = 64;

struct PreparedInsert {
    id: String,
    vector: Vec<f32>,
    level: usize,
    layer_candidates: Vec<Vec<ScoredNode>>,
}

impl AnnIndex {
    pub fn bulk_load<I>(&mut self, records: I, threads: usize) -> Result<()>
    where
        I: IntoIterator<Item = (String, Vec<f32>)>,
    {
        let records: Vec<(String, Vec<f32>)> = records.into_iter().collect();
        for (_, vector) in &records {
            self.validate_vector(vector)?;
        }
        let threads = threads.clamp(1, MAX_BULK_THREADS);
        self.reserve(records.len());

        let mut items = records.into_iter();
        loop {
            let batch: Vec<(String, Vec<f32>)> = items.by_ref().take(BULK_BATCH).collect();
            if batch.is_empty() {
                return Ok(());
            }
            let prepared = self.prepare_batch(batch, threads);
            self.apply_batch(prepared);
        }
    }

    fn prepare_batch(&self, batch: Vec<(String, Vec<f32>)>, threads: usize) -> Vec<PreparedInsert> {
        if threads <= 1 || batch.len() <= 1 {
            let mut scratch = SearchScratch::default();
            return batch
                .into_iter()
                .map(|(id, vector)| self.prepare_insert(id, vector, &mut scratch))
                .collect();
        }

        let chunk_size = batch.len().div_ceil(threads);
        let mut chunks: Vec<Vec<(String, Vec<f32>)>> = Vec::new();
        let mut items = batch.into_iter();
        loop {
            let chunk: Vec<(String, Vec<f32>)> = items.by_ref().take(chunk_size).collect();
            if chunk.is_empty() {
                break;
            }
            chunks.push(chunk);
        }

        let mut prepared: Vec<Vec<PreparedInsert>> = Vec::with_capacity(chunks.len());
        thread::scope(|scope| {
            let handles: Vec<_> = chunks
                .into_iter()
                .map(|chunk| {
                    scope.spawn(move || {
                        let mut scratch = SearchScratch::default();
                        chunk
                            .into_iter()
                            .map(|(id, vector)| self.prepare_insert(id, vector, &mut scratch))
                            .collect::<Vec<_>>()
                    })
                })
                .collect();
            for handle in handles {
                prepared.push(handle.join().expect("bulk load worker panicked"));
            }
        });
        prepared.into_iter().flatten().collect()
    }

    fn prepare_insert(
        &self,
        id: String,
        mut vector: Vec<f32>,
        scratch: &mut SearchScratch,
    ) -> PreparedInsert {
        normalize_in_place(&mut vector);
        let level = self.level_for_id(&id);
        let mut layer_candidates = vec![Vec::new(); level + 1];

        if let Some(mut entry) = self.entry {
            for layer in ((level + 1)..=self.max_level).rev() {
                entry = self.greedy_search_vector(&vector, entry, layer);
            }
            let top_layer = level.min(self.max_level);
            for layer in (0..=top_layer).rev() {
                self.search_layer_vector(
                    &vector,
                    entry,
                    layer,
                    self.config.ef_construction,
                    scratch,
                );
                layer_candidates[layer] = scratch.found.clone();
                if let Some(best) = scratch.found.iter().copied().max() {
                    entry = best.index;
                }
            }
        }

        PreparedInsert {
            id,
            vector,
            level,
            layer_candidates,
        }
    }

    fn apply_batch(&mut self, prepared: Vec<PreparedInsert>) {
        let mut applied: Vec<usize> = Vec::with_capacity(prepared.len());
        for item in prepared {
            if let Some(existing) = self.ids.remove(&item.id) {
                if !self.nodes[existing].deleted {
                    self.nodes[existing].deleted = true;
                    self.active = self.active.saturating_sub(1);
                }
            }

            let index = self.nodes.len();
            self.nodes.push(Node {
                id: item.id.clone(),
                vector: item.vector,
                deleted: false,
                neighbors: vec![Vec::new(); item.level + 1],
            });
            self.ids.insert(item.id, index);
            self.add_entry_point(index);
            self.active += 1;

            let Some(_) = self.entry else {
                self.entry = Some(index);
                self.max_level = item.level;
                applied.push(index);
                continue;
            };

            let top_layer = item.level.min(self.max_level);
            for layer in (0..=top_layer).rev() {
                let mut candidates = item
                    .layer_candidates
                    .get(layer)
                    .cloned()
                    .unwrap_or_default();
                for &peer in &applied {
                    if peer != index
                        && !self.nodes[peer].deleted
                        && self.nodes[peer].neighbors.len() > layer
                    {
                        candidates.push(ScoredNode {
                            index: peer,
                            score: dot(&self.nodes[index].vector, &self.nodes[peer].vector),
                        });
                    }
                }
                let selected = self.select_neighbors(index, layer, &mut candidates);
                for &neighbor in &selected {
                    self.nodes[neighbor].neighbors[layer].push(index);
                    self.prune_neighbors(neighbor, layer);
                }
                self.nodes[index].neighbors[layer] = selected;
            }

            if item.level > self.max_level {
                self.entry = Some(index);
                self.max_level = item.level;
            }
            applied.push(index);
        }
    }
}