feox-ann 0.1.0

use std::cell::RefCell;
use std::cmp::Reverse;

use crate::graph::{compare_scored_desc, ScoredNode, SearchScratch};
use crate::math::{dot, normalize};
use crate::{AnnFilter, AnnIndex, AnnQuery, Result};

thread_local! {
    static QUERY_SCRATCH: RefCell<SearchScratch> = RefCell::new(SearchScratch::default());
}

impl AnnIndex {
    pub fn query(&self, input: AnnQuery<'_>) -> Result<Vec<crate::AnnCandidate>> {
        self.validate_vector(input.vector)?;
        if input.top_k == 0 || self.active == 0 {
            return Ok(Vec::new());
        }

        let query = normalize(input.vector);
        let Some(mut entry) = self.best_entry_point(&query) else {
            return Ok(Vec::new());
        };

        for layer in (1..=self.max_level).rev() {
            entry = self.greedy_search_vector(&query, entry, layer);
        }

        let ef_search = input
            .ef_search
            .unwrap_or(self.config.ef_search)
            .max(input.top_k);
        QUERY_SCRATCH.with(|cell| match cell.try_borrow_mut() {
            Ok(mut scratch) => self.finish_query(&query, entry, ef_search, input, &mut scratch),
            Err(_) => {
                let mut scratch = SearchScratch::default();
                self.finish_query(&query, entry, ef_search, input, &mut scratch)
            }
        })
    }

    fn finish_query(
        &self,
        query: &[f32],
        entry: usize,
        ef_search: usize,
        input: AnnQuery<'_>,
        scratch: &mut SearchScratch,
    ) -> Result<Vec<crate::AnnCandidate>> {
        self.search_layer_vector(query, entry, 0, ef_search, scratch);
        scratch
            .found
            .retain(|candidate| self.result_allowed(candidate.index, input.filter));
        scratch.found.sort_by(compare_scored_desc);
        scratch.found.truncate(input.top_k);

        Ok(scratch
            .found
            .iter()
            .map(|candidate| crate::AnnCandidate {
                id: self.nodes[candidate.index].id.clone(),
                score: candidate.score,
            })
            .collect())
    }

    pub(crate) fn greedy_search(&self, query_index: usize, entry: usize, layer: usize) -> usize {
        let query = &self.nodes[query_index].vector;
        self.greedy_search_vector(query, entry, layer)
    }

    pub(crate) fn greedy_search_vector(&self, query: &[f32], entry: usize, layer: usize) -> usize {
        let mut current = entry;
        let mut current_score = dot(query, &self.nodes[current].vector);

        loop {
            let mut changed = false;
            for &neighbor in self.layer_neighbors(current, layer) {
                let score = dot(query, &self.nodes[neighbor].vector);
                if score > current_score {
                    current = neighbor;
                    current_score = score;
                    changed = true;
                }
            }
            if !changed {
                return current;
            }
        }
    }

    pub(crate) fn search_layer(
        &self,
        query_index: usize,
        entry: usize,
        layer: usize,
        ef: usize,
        scratch: &mut SearchScratch,
    ) {
        let query = &self.nodes[query_index].vector;
        self.search_layer_vector(query, entry, layer, ef, scratch);
    }

    pub(crate) fn search_layer_vector(
        &self,
        query: &[f32],
        entry: usize,
        layer: usize,
        ef: usize,
        scratch: &mut SearchScratch,
    ) {
        scratch.reset(self.nodes.len());
        let entry_score = dot(query, &self.nodes[entry].vector);
        let entry_node = ScoredNode {
            index: entry,
            score: entry_score,
        };
        scratch.candidates.push(entry_node);
        scratch.results.push(Reverse(entry_node));
        scratch.visit(entry);

        while let Some(candidate) = scratch.candidates.pop() {
            if let Some(Reverse(worst)) = scratch.results.peek() {
                if scratch.results.len() >= ef && candidate.score < worst.score {
                    break;
                }
            }

            for &neighbor in self.layer_neighbors(candidate.index, layer) {
                if !scratch.visit(neighbor) {
                    continue;
                }
                let score = dot(query, &self.nodes[neighbor].vector);
                let scored = ScoredNode {
                    index: neighbor,
                    score,
                };
                let should_keep = scratch
                    .results
                    .peek()
                    .map(|Reverse(worst)| scratch.results.len() < ef || score > worst.score)
                    .unwrap_or(true);
                if should_keep {
                    scratch.candidates.push(scored);
                    scratch.results.push(Reverse(scored));
                    if scratch.results.len() > ef {
                        scratch.results.pop();
                    }
                }
            }
        }

        scratch
            .found
            .extend(scratch.results.iter().map(|Reverse(item)| *item));
    }

    fn result_allowed(&self, index: usize, filter: Option<&dyn AnnFilter>) -> bool {
        if self.nodes[index].deleted {
            return false;
        }
        filter
            .map(|filter| filter.accept(&self.nodes[index].id))
            .unwrap_or(true)
    }
}