feox-ann 0.1.0

use std::collections::HashMap;

use crate::graph::{Node, ScoredNode};
use crate::math::{dot, stable_hash};
use crate::{AnnConfig, AnnError, Result};

pub(crate) const MAX_ENTRY_POINTS: usize = 256;

#[derive(Clone)]
pub struct AnnIndex {
    pub(crate) config: AnnConfig,
    pub(crate) nodes: Vec<Node>,
    pub(crate) ids: HashMap<String, usize>,
    pub(crate) entry: Option<usize>,
    pub(crate) entry_points: Vec<usize>,
    pub(crate) max_level: usize,
    pub(crate) active: usize,
}

impl AnnIndex {
    pub fn new(config: AnnConfig) -> Result<Self> {
        if config.dimensions == 0 {
            return Err(AnnError::InvalidConfig(
                "dimensions must be greater than zero".to_string(),
            ));
        }
        if config.max_neighbors == 0 {
            return Err(AnnError::InvalidConfig(
                "max_neighbors must be greater than zero".to_string(),
            ));
        }
        if config.max_base_neighbors < config.max_neighbors {
            return Err(AnnError::InvalidConfig(
                "max_base_neighbors must be at least max_neighbors".to_string(),
            ));
        }
        if config.ef_construction < config.max_neighbors {
            return Err(AnnError::InvalidConfig(
                "ef_construction must be at least max_neighbors".to_string(),
            ));
        }
        if config.ef_search == 0 {
            return Err(AnnError::InvalidConfig(
                "ef_search must be greater than zero".to_string(),
            ));
        }

        Ok(Self {
            config,
            nodes: Vec::new(),
            ids: HashMap::new(),
            entry: None,
            entry_points: Vec::new(),
            max_level: 0,
            active: 0,
        })
    }

    pub fn reserve(&mut self, additional: usize) {
        self.nodes.reserve(additional);
        self.ids.reserve(additional);
        self.entry_points.reserve(additional.min(MAX_ENTRY_POINTS));
    }

    pub fn dimensions(&self) -> usize {
        self.config.dimensions
    }

    pub fn len(&self) -> usize {
        self.active
    }

    pub fn is_empty(&self) -> bool {
        self.active == 0
    }

    pub fn delete(&mut self, id: &str) -> bool {
        let Some(index) = self.ids.remove(id) else {
            return false;
        };
        if self.nodes[index].deleted {
            return false;
        }
        self.nodes[index].deleted = true;
        self.active = self.active.saturating_sub(1);
        true
    }

    pub(crate) fn validate_vector(&self, vector: &[f32]) -> Result<()> {
        if vector.len() != self.config.dimensions {
            return Err(AnnError::InvalidVector(format!(
                "expected {} dimensions, got {}",
                self.config.dimensions,
                vector.len()
            )));
        }
        if vector.iter().any(|value| !value.is_finite()) {
            return Err(AnnError::InvalidVector(
                "vector values must be finite".to_string(),
            ));
        }
        Ok(())
    }

    pub(crate) fn best_candidate(&self, query_index: usize, layer: usize) -> Option<usize> {
        self.nodes[query_index]
            .neighbors
            .get(layer)
            .and_then(|neighbors| neighbors.first().copied())
    }

    pub(crate) fn layer_neighbors(&self, index: usize, layer: usize) -> &[usize] {
        self.nodes[index]
            .neighbors
            .get(layer)
            .map(Vec::as_slice)
            .unwrap_or(&[])
    }

    pub(crate) fn add_entry_point(&mut self, index: usize) {
        if self.entry_points.len() < MAX_ENTRY_POINTS {
            self.entry_points.push(index);
        }
    }

    pub(crate) fn best_entry_point(&self, query: &[f32]) -> Option<usize> {
        let mut best = None;
        for &index in &self.entry_points {
            if self.nodes[index].deleted {
                continue;
            }
            let score = dot(query, &self.nodes[index].vector);
            if best
                .map(|candidate: ScoredNode| score > candidate.score)
                .unwrap_or(true)
            {
                best = Some(ScoredNode { index, score });
            }
        }
        best.map(|candidate| candidate.index).or(self.entry)
    }

    pub(crate) fn neighbor_limit(&self, layer: usize) -> usize {
        if layer == 0 {
            self.config.max_base_neighbors
        } else {
            self.config.max_neighbors
        }
    }

    pub(crate) fn level_for_id(&self, id: &str) -> usize {
        let mut hash = stable_hash(id.as_bytes());
        let mut level = 0;
        while level < self.config.max_level && (hash & 0b11) == 0 {
            level += 1;
            hash >>= 2;
        }
        level
    }
}