minni 0.1.1

use super::dense::DenseSearchResult;
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::path::Path;

const DEFAULT_CENTROIDS: usize = 64;
const DEFAULT_PROBES: usize = 4;
const KMEANS_ITERS: usize = 4;

#[derive(Debug, Clone, Serialize, Deserialize)]
struct AnnEntry {
    chunk_id: String,
    vector: Vec<f32>,
}

/// Lightweight ANN index using IVF-style centroid buckets.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AnnIndex {
    dim: usize,
    centroids: Vec<Vec<f32>>,
    buckets: Vec<Vec<AnnEntry>>,
}

impl AnnIndex {
    pub fn build(embeddings: &[(String, Vec<f32>)]) -> Result<Option<Self>> {
        if embeddings.is_empty() {
            return Ok(None);
        }

        let dim = embeddings[0].1.len();
        if dim == 0 {
            return Ok(None);
        }
        if embeddings.iter().any(|(_, v)| v.len() != dim) {
            anyhow::bail!("Embedding dimensions are inconsistent");
        }

        let n = embeddings.len();
        let centroid_count = DEFAULT_CENTROIDS
            .min((n as f64).sqrt().ceil() as usize)
            .max(1);

        let mut centroids = initial_centroids(embeddings, centroid_count);
        let mut assignments = vec![0usize; n];

        for _ in 0..KMEANS_ITERS {
            for (i, (_, vector)) in embeddings.iter().enumerate() {
                assignments[i] = nearest_centroid(vector, &centroids);
            }

            let mut sums = vec![vec![0.0f32; dim]; centroid_count];
            let mut counts = vec![0usize; centroid_count];

            for (i, (_, vector)) in embeddings.iter().enumerate() {
                let c = assignments[i];
                counts[c] += 1;
                for (j, value) in vector.iter().enumerate() {
                    sums[c][j] += value;
                }
            }

            for c in 0..centroid_count {
                if counts[c] == 0 {
                    continue;
                }
                for value in &mut sums[c] {
                    *value /= counts[c] as f32;
                }
                l2_normalize(&mut sums[c]);
                centroids[c] = sums[c].clone();
            }
        }

        let mut buckets = vec![Vec::new(); centroid_count];
        for (chunk_id, vector) in embeddings {
            let c = nearest_centroid(vector, &centroids);
            buckets[c].push(AnnEntry {
                chunk_id: chunk_id.clone(),
                vector: vector.clone(),
            });
        }

        Ok(Some(Self {
            dim,
            centroids,
            buckets,
        }))
    }

    pub fn save(&self, path: &Path) -> Result<()> {
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        let content = serde_json::to_string(self)?;
        std::fs::write(path, content)
            .with_context(|| format!("Failed to write ANN index at {:?}", path))?;
        Ok(())
    }

    pub fn load(path: &Path) -> Result<Option<Self>> {
        if !path.exists() {
            return Ok(None);
        }
        let content = std::fs::read_to_string(path)
            .with_context(|| format!("Failed to read ANN index at {:?}", path))?;
        let index = serde_json::from_str::<Self>(&content)
            .with_context(|| format!("Failed to parse ANN index at {:?}", path))?;
        Ok(Some(index))
    }

    pub fn search(&self, query_vector: &[f32], limit: usize) -> Vec<DenseSearchResult> {
        if limit == 0 || query_vector.len() != self.dim || self.centroids.is_empty() {
            return Vec::new();
        }

        let mut centroid_scores: Vec<(usize, f32)> = self
            .centroids
            .iter()
            .enumerate()
            .map(|(i, c)| (i, dot(query_vector, c)))
            .collect();
        centroid_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        let probes = DEFAULT_PROBES.min(centroid_scores.len()).max(1);
        let mut scored = Vec::new();
        for (idx, _) in centroid_scores.into_iter().take(probes) {
            for entry in &self.buckets[idx] {
                scored.push(DenseSearchResult {
                    chunk_id: entry.chunk_id.clone(),
                    score: dot(query_vector, &entry.vector),
                });
            }
        }

        if scored.is_empty() {
            return scored;
        }

        scored.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        scored.truncate(limit);
        scored
    }
}

fn initial_centroids(embeddings: &[(String, Vec<f32>)], k: usize) -> Vec<Vec<f32>> {
    if k == 1 {
        return vec![embeddings[0].1.clone()];
    }
    let last = embeddings.len() - 1;
    (0..k)
        .map(|i| {
            let idx = i * last / (k - 1);
            embeddings[idx].1.clone()
        })
        .collect()
}

fn nearest_centroid(vector: &[f32], centroids: &[Vec<f32>]) -> usize {
    let mut best_i = 0usize;
    let mut best_score = f32::NEG_INFINITY;
    for (i, centroid) in centroids.iter().enumerate() {
        let score = dot(vector, centroid);
        if score > best_score {
            best_score = score;
            best_i = i;
        }
    }
    best_i
}

fn dot(a: &[f32], b: &[f32]) -> f32 {
    a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
}

fn l2_normalize(v: &mut [f32]) {
    let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt().max(1e-12);
    for x in v.iter_mut() {
        *x /= norm;
    }
}