/// IvfFlat (Inverted File Index) — clustering-based approximate nearest neighbour search.
///
/// Build: k-means partitions all vectors into `nlist` Voronoi cells.
/// Search: query is compared to all centroids; only the `nprobe` nearest cells are scanned.
///
/// Vectors stay quantized (packed bytes) in the inverted lists — same format as BruteForce.
/// Centroids are stored as f32 in the *transformed* (RHDH-rotated) space so that distance
/// to a centroid can be computed with a plain dot product against the transformed query.

use crate::{Metric, SearchResult, error::{MonaVecError, Result}};
use std::collections::BinaryHeap;
use std::cmp::Ordering;

// ---------------------------------------------------------------------------
// Centroid list
// ---------------------------------------------------------------------------

/// One Voronoi cell: a centroid (f32 in transformed space) + packed-byte vectors.
pub struct IvfList {
    pub centroid: Vec<f32>,
    /// Packed quantized vectors concatenated, each `packed_bytes` long.
    pub data:     Vec<u8>,
    pub ids:      Vec<u64>,
    pub norms:    Vec<f32>,
}

impl IvfList {
    fn new(centroid: Vec<f32>) -> Self {
        Self { centroid, data: Vec::new(), ids: Vec::new(), norms: Vec::new() }
    }

    fn add(&mut self, id: u64, packed: &[u8], norm: f32) {
        self.ids.push(id);
        self.data.extend_from_slice(packed);
        self.norms.push(norm);
    }

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

// ---------------------------------------------------------------------------
// IvfFlat index
// ---------------------------------------------------------------------------

pub struct IvfFlat {
    pub lists:       Vec<IvfList>,
    pub packed_bytes: usize,
    pub nprobe:      usize,
    /// Transformed dimension (padded_dim = dim.next_power_of_two()).
    pub transformed_dim: usize,
    pub metric:      Metric,
}

impl IvfFlat {
    /// Build from a flat store of (id, packed, norm, transformed_vec) tuples.
    ///
    /// `transformed_vecs`: each entry is the full RHDH-transformed f32 vector for one item.
    /// These are used only for k-means centroid computation; they are not stored after build.
    pub fn build(
        nlist:            usize,
        nprobe:           usize,
        packed_bytes:     usize,
        transformed_dim:  usize,
        ids:              &[u64],
        packed:           &[Vec<u8>],
        norms:            &[f32],
        transformed_vecs: &[Vec<f32>],
        seed:             u64,
        metric:           Metric,
    ) -> Result<Self> {
        let n = ids.len();
        if n == 0 {
            return Err(MonaVecError::EmptyIndex);
        }
        let nlist = nlist.min(n); // can't have more clusters than vectors

        let centroids = kmeans(transformed_vecs, nlist, 25, seed, metric);

        // Assign each vector to its nearest centroid.
        let assignments: Vec<usize> = transformed_vecs.iter()
            .map(|v| nearest_centroid(&centroids, v, metric))
            .collect();

        let mut lists: Vec<IvfList> = centroids.into_iter().map(IvfList::new).collect();
        for i in 0..n {
            lists[assignments[i]].add(ids[i], &packed[i], norms[i]);
        }

        Ok(Self { lists, packed_bytes, nprobe, transformed_dim, metric })
    }

    /// Search for top-K nearest neighbours using the pre-transformed query.
    ///
    /// `transformed_query`: RHDH-transformed f32 query (padded_dim length).
    /// `score_fn(packed) -> f32`: scores a packed vector against the query.
    pub fn search<F>(
        &self,
        transformed_query: &[f32],
        k:                 usize,
        score_fn:          F,
    ) -> Vec<SearchResult>
    where
        F: Fn(&[u8]) -> f32,
    {
        if k == 0 { return Vec::new(); }

        // Rank all centroids by proximity to the query using the correct metric.
        // L2: score = 2*dot(q,c) - ||c||² (≡ -||q-c||² up to constant ||q||²).
        // Cosine/Dot: score = dot(q, c).
        let mut centroid_scores: Vec<(usize, f32)> = self.lists.iter().enumerate()
            .map(|(i, list)| {
                let s = match self.metric {
                    Metric::L2 => 2.0 * dot(transformed_query, &list.centroid)
                                  - dot(&list.centroid, &list.centroid),
                    _          => dot(transformed_query, &list.centroid),
                };
                (i, s)
            })
            .collect();
        centroid_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal));

        let nprobe = self.nprobe.min(self.lists.len());
        let mut heap: BinaryHeap<Candidate> = BinaryHeap::with_capacity(k + 1);

        for (list_idx, _) in centroid_scores.iter().take(nprobe) {
            let list = &self.lists[*list_idx];
            for (pos, id) in list.ids.iter().enumerate() {
                let start  = pos * self.packed_bytes;
                let packed = &list.data[start..start + self.packed_bytes];
                let raw    = score_fn(packed);
                let score  = if list.norms.is_empty() { raw } else { self.metric.adjust_score(raw, list.norms[pos]) };

                if heap.len() < k {
                    heap.push(Candidate { score, id: *id });
                } else if let Some(min) = heap.peek() {
                    if score > min.score {
                        heap.pop();
                        heap.push(Candidate { score, id: *id });
                    }
                }
            }
        }

        let mut results: Vec<SearchResult> = heap.into_iter()
            .map(|c| SearchResult { id: c.id, score: c.score })
            .collect();
        results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(Ordering::Equal));
        results
    }

    pub fn total_vectors(&self) -> usize {
        self.lists.iter().map(|l| l.len()).sum()
    }

    /// Serialize IvfFlat index data (centroids + lists) after the main header+vectors section.
    ///
    /// Layout:
    ///   [NLIST:    u32             ] number of lists
    ///   [TDIM:     u32             ] transformed dimension (centroid length)
    /// For each list:
    ///   [CENTROID: f32 × tdim      ]
    ///   [COUNT:    u32             ] vectors in this list
    ///   for each vector in list:
    ///     [ID:     u64             ]
    ///     [PACKED: u8 × pb        ]
    ///     [NORM:   f32             ]
    pub fn write_data<W: std::io::Write>(&self, w: &mut W) -> std::io::Result<()> {
        let nlist = self.lists.len() as u32;
        let tdim  = if self.lists.is_empty() { 0u32 } else { self.lists[0].centroid.len() as u32 };
        w.write_all(&nlist.to_le_bytes())?;
        w.write_all(&tdim.to_le_bytes())?;
        for list in &self.lists {
            for &x in &list.centroid {
                w.write_all(&x.to_le_bytes())?;
            }
            let count = list.ids.len() as u32;
            w.write_all(&count.to_le_bytes())?;
            for i in 0..list.ids.len() {
                w.write_all(&list.ids[i].to_le_bytes())?;
                let start = i * self.packed_bytes;
                w.write_all(&list.data[start..start + self.packed_bytes])?;
                w.write_all(&list.norms[i].to_le_bytes())?;
            }
        }
        Ok(())
    }

    /// Deserialize IvfFlat index data written by `write_data`.
    pub fn read_data<R: std::io::Read>(
        r:           &mut R,
        packed_bytes: usize,
        nprobe:       usize,
        metric:       Metric,
    ) -> std::io::Result<Self> {
        let mut b4 = [0u8; 4];
        r.read_exact(&mut b4)?;
        let nlist = u32::from_le_bytes(b4) as usize;
        r.read_exact(&mut b4)?;
        let tdim  = u32::from_le_bytes(b4) as usize;

        let mut lists = Vec::with_capacity(nlist);
        for _ in 0..nlist {
            let mut centroid = vec![0f32; tdim];
            for x in centroid.iter_mut() {
                r.read_exact(&mut b4)?;
                *x = f32::from_le_bytes(b4);
            }
            r.read_exact(&mut b4)?;
            let count = u32::from_le_bytes(b4) as usize;
            let mut list = IvfList::new(centroid);
            let mut b8 = [0u8; 8];
            for _ in 0..count {
                r.read_exact(&mut b8)?;
                let id = u64::from_le_bytes(b8);
                let mut packed = vec![0u8; packed_bytes];
                r.read_exact(&mut packed)?;
                r.read_exact(&mut b4)?;
                let norm = f32::from_le_bytes(b4);
                list.add(id, &packed, norm);
            }
            lists.push(list);
        }
        let transformed_dim = if lists.is_empty() { 0 } else { lists[0].centroid.len() };
        Ok(Self { lists, packed_bytes, nprobe, transformed_dim, metric })
    }
}

// ---------------------------------------------------------------------------
// K-means (Lloyd's algorithm, cosine-friendly: normalize centroids each iter)
// ---------------------------------------------------------------------------

fn kmeans(vecs: &[Vec<f32>], k: usize, max_iter: usize, seed: u64, metric: Metric) -> Vec<Vec<f32>> {
    let dim = vecs[0].len();
    let n   = vecs.len();

    // Seeded centroid init: pick k vectors deterministically by stride.
    let stride = (n / k).max(1);
    let mut centroids: Vec<Vec<f32>> = (0..k)
        .map(|i| vecs[(i * stride + (seed as usize % stride.max(1))) % n].clone())
        .collect();

    let mut assignments = vec![0usize; n];

    for _ in 0..max_iter {
        // Assign each vector to its nearest centroid.
        let mut changed = false;
        for (i, v) in vecs.iter().enumerate() {
            let nearest = nearest_centroid(&centroids, v, metric);
            if nearest != assignments[i] {
                assignments[i] = nearest;
                changed = true;
            }
        }
        if !changed { break; }

        // Update: mean of assigned vectors.
        let mut sums   = vec![vec![0f32; dim]; k];
        let mut counts = vec![0usize; k];
        for (i, v) in vecs.iter().enumerate() {
            let c = assignments[i];
            for (s, x) in sums[c].iter_mut().zip(v.iter()) {
                *s += x;
            }
            counts[c] += 1;
        }

        for c in 0..k {
            if counts[c] == 0 {
                // Empty cluster: reinitialise from a seeded vector.
                centroids[c] = vecs[(c * 7919 + seed as usize) % n].clone();
                // Cosine: centroid must be normalized so dot product = cosine similarity.
                if metric == Metric::Cosine {
                    l2_normalize(&mut centroids[c]);
                }
            } else {
                let cnt = counts[c] as f32;
                for (centroid_d, sum) in centroids[c].iter_mut().zip(sums[c].iter()) {
                    *centroid_d = sum / cnt;
                }
                // Cosine: normalize so centroid direction is the representative direction.
                // Dot/L2: keep raw mean so magnitude is preserved.
                if metric == Metric::Cosine {
                    l2_normalize(&mut centroids[c]);
                }
            }
        }
    }

    centroids
}

fn nearest_centroid(centroids: &[Vec<f32>], v: &[f32], metric: Metric) -> usize {
    centroids.iter()
        .enumerate()
        .map(|(i, c)| {
            let score = match metric {
                // L2: argmin ||v-c||² = argmax (2*dot(v,c) - ||c||²)
                Metric::L2 => 2.0 * dot(v, c) - dot(c, c),
                // Cosine/Dot: argmax dot(v, c)
                _          => dot(v, c),
            };
            (i, score)
        })
        .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal))
        .map(|(i, _)| i)
        .unwrap_or(0)
}

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

fn l2_normalize(v: &mut Vec<f32>) {
    let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > 1e-9 {
        v.iter_mut().for_each(|x| *x /= norm);
    }
}

// ---------------------------------------------------------------------------
// Min-heap candidate (same pattern as brute.rs)
// ---------------------------------------------------------------------------

#[derive(PartialEq)]
struct Candidate { score: f32, id: u64 }
impl Eq for Candidate {}
impl PartialOrd for Candidate {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> { Some(self.cmp(other)) }
}
impl Ord for Candidate {
    fn cmp(&self, other: &Self) -> Ordering {
        other.score.partial_cmp(&self.score)
            .unwrap_or(Ordering::Equal)
            .then(self.id.cmp(&other.id))
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    fn unit_vec(dim: usize, seed: u64) -> Vec<f32> {
        use rand::SeedableRng;
        use rand::RngCore;
        use rand_chacha::ChaCha20Rng;
        let mut rng = ChaCha20Rng::seed_from_u64(seed);
        let mut v: Vec<f32> = (0..dim)
            .map(|_| (rng.next_u32() as f32 / u32::MAX as f32) * 2.0 - 1.0)
            .collect();
        let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        v.iter_mut().for_each(|x| *x /= norm);
        v
    }

    #[test]
    fn build_and_search_basic() {
        let dim = 32usize;
        let n   = 200usize;
        let ids: Vec<u64>      = (0..n as u64).collect();
        let vecs: Vec<Vec<f32>> = (0..n).map(|i| unit_vec(dim, i as u64)).collect();
        let norms               = vec![1.0f32; n];
        // Use identity packed (1 byte per vec, just the index) — score_fn ignores it.
        let packed: Vec<Vec<u8>> = (0..n).map(|_| vec![0u8; 4]).collect();

        let index = IvfFlat::build(
            16, 4, 4, dim,
            &ids, &packed, &norms, &vecs, 42, Metric::Cosine,
        ).unwrap();

        assert_eq!(index.total_vectors(), n);
        assert_eq!(index.lists.len(), 16);

        // Search: score_fn returns dot(query, vec) via transformed vecs — use a closure
        // that always returns 1.0 (just tests top-K machinery).
        let query  = unit_vec(dim, 0);
        let results = index.search(&query, 5, |_| 1.0f32);
        assert_eq!(results.len(), 5);
    }

    #[test]
    fn empty_index_errors() {
        let result = IvfFlat::build(
            4, 1, 4, 8,
            &[], &[], &[], &[], 0, Metric::Cosine,
        );
        assert!(matches!(result, Err(MonaVecError::EmptyIndex)));
    }

    #[test]
    fn nlist_capped_at_n() {
        let dim = 8usize;
        let ids: Vec<u64>       = vec![0, 1, 2];
        let vecs: Vec<Vec<f32>> = (0..3).map(|i| unit_vec(dim, i as u64)).collect();
        let packed: Vec<Vec<u8>> = vec![vec![0u8; 4]; 3];
        let norms = vec![1.0f32; 3];

        // nlist=100 but only 3 vectors → capped to 3
        let index = IvfFlat::build(100, 2, 4, dim, &ids, &packed, &norms, &vecs, 0, Metric::Cosine).unwrap();
        assert!(index.lists.len() <= 3);
        assert_eq!(index.total_vectors(), 3);
    }
}
