use std::io::{BufReader, BufWriter, Read, Write};
use crate::{
    Config, IndexType, Metric, MonaVecError, SearchResult,
    error::Result,
    io::format::Header,
    quantize::{QuantPipeline, StdParams},
    search::{VectorStore, top_k, Allowlist},
};
#[cfg(feature = "ivf")]
use crate::search::IvfFlat;
#[cfg(feature = "hnsw")]
use crate::search::Hnsw;

/// MonaVec vector search index.
///
/// Supports multiple index backends selected at construction time via [`Config`]:
/// - `BruteForce` (default): exact scan, always ready after `add()`.
/// - `IvfFlat` (feature `ivf`): approximate, requires `build()` after `add()`.
/// - `Hnsw` (feature `hnsw`): approximate, built incrementally during `add()`.
pub struct MonaVec {
    config:   Config,
    pipeline: QuantPipeline,
    /// Flat store — permanent for BruteForce, staging area for IvfFlat pre-build.
    store:    VectorStore,
    /// IvfFlat index, populated after `build()`. None for BruteForce/HNSW or pre-build.
    #[cfg(feature = "ivf")]
    ivf:      Option<IvfFlat>,
    /// RHDH-transformed vectors buffered for k-means during IvfFlat build.
    /// Cleared after `build()` to release memory.
    #[cfg(feature = "ivf")]
    pre_build_transformed: Vec<Vec<f32>>,
    /// HNSW graph — built incrementally during `add()`.
    #[cfg(feature = "hnsw")]
    hnsw: Option<Hnsw>,
}

impl MonaVec {
    pub fn new(config: Config) -> Result<Self> {
        if config.bit_width != 2 && config.bit_width != 4 {
            return Err(MonaVecError::InvalidBitWidth(config.bit_width));
        }
        #[allow(unreachable_patterns)]
        match config.index_type {
            IndexType::BruteForce => {}
            #[cfg(feature = "ivf")]
            IndexType::IvfFlat => {}
            #[cfg(feature = "hnsw")]
            IndexType::Hnsw => {}
            _ => return Err(MonaVecError::IndexTypeNotImplemented(config.index_type.name())),
        }
        let pipeline     = QuantPipeline::new(&config)?;
        let packed_bytes = pipeline.packed_bytes();
        #[cfg(feature = "hnsw")]
        let hnsw = if config.index_type == IndexType::Hnsw {
            let ef_s = if config.hnsw_ef_search == 0 { config.hnsw_m * 2 } else { config.hnsw_ef_search };
            Some(Hnsw::new(config.hnsw_m, config.hnsw_ef_construction, ef_s, packed_bytes))
        } else {
            None
        };

        Ok(Self {
            config,
            pipeline,
            store: VectorStore::new(packed_bytes),
            #[cfg(feature = "ivf")]
            ivf:                    None,
            #[cfg(feature = "ivf")]
            pre_build_transformed:  Vec::new(),
            #[cfg(feature = "hnsw")]
            hnsw,
        })
    }

    /// Build the search index after all vectors have been added.
    ///
    /// - `BruteForce`: no-op (always ready after `add()`).
    /// - `IvfFlat`: runs k-means to partition vectors into `nlist` clusters.
    ///   Must be called before `search()`. Safe to call multiple times (rebuilds).
    /// - `Hnsw`: no-op (graph is built incrementally during `add()`).
    pub fn build(&mut self) -> Result<()> {
        #[allow(unreachable_patterns)]
        match self.config.index_type {
            IndexType::BruteForce => Ok(()),
            #[cfg(feature = "ivf")]
            IndexType::IvfFlat => self.build_ivf(),
            #[cfg(feature = "hnsw")]
            IndexType::Hnsw => {
                if let Some(ref mut hnsw) = self.hnsw {
                    hnsw.finalize_build();
                }
                Ok(())
            }
            _ => Err(MonaVecError::IndexTypeNotImplemented(self.config.index_type.name())),
        }
    }

    #[cfg(feature = "ivf")]
    fn build_ivf(&mut self) -> Result<()> {
        let n            = self.store.len();
        let packed_bytes = self.pipeline.packed_bytes();
        let nlist        = self.config.ivf_nlist;
        let nprobe       = self.config.ivf_nprobe;
        let transformed_dim = self.pipeline.rhdh.padded_dim;

        let ids: Vec<u64> = self.store.ids.clone();
        let packed: Vec<Vec<u8>> = (0..n)
            .map(|i| self.store.get_packed(i).to_vec())
            .collect();

        let ivf = IvfFlat::build(
            nlist, nprobe, packed_bytes, transformed_dim,
            &ids, &packed, &self.store.norms,
            &self.pre_build_transformed,
            self.config.seed,
            self.config.metric,
        )?;

        self.ivf = Some(ivf);
        // Release staging buffers — vectors now live inside ivf.lists.
        self.store = VectorStore::new(packed_bytes);
        self.pre_build_transformed.clear();
        self.pre_build_transformed.shrink_to_fit();
        Ok(())
    }

    /// Compute per-dimension standardization parameters from a representative sample.
    ///
    /// Must be called **before** `add()`. No-op for non-L2 metrics.
    /// After fitting, `encode()` and `transform_query()` subtract the per-dim mean and
    /// divide by per-dim std, making the RHDH input distribution ≈ N(0,1) so the
    /// Lloyd-Max quantization tables are valid for raw-magnitude vectors (pixels, SIFT, etc.).
    ///
    /// The parameters are persisted in `.mvec` v6+ files and restored on `load()`.
    pub fn fit(&mut self, vectors: &[Vec<f32>]) -> Result<()> {
        self.pipeline.fit(vectors, self.len())
    }

    /// Number of indexed vectors.
    pub fn len(&self) -> usize {
        #[cfg(feature = "ivf")]
        if let Some(ref ivf) = self.ivf {
            return ivf.total_vectors();
        }
        #[cfg(feature = "hnsw")]
        if let Some(ref hnsw) = self.hnsw {
            return hnsw.len();
        }
        self.store.len()
    }
    pub fn is_empty(&self)   -> bool      { self.len() == 0 }
    pub fn dim(&self)        -> usize     { self.config.dim }
    pub fn bit_width(&self)  -> u8        { self.config.bit_width }
    pub fn index_type(&self) -> IndexType { self.config.index_type }
    /// Dims at 4-bit in padded space (0 = pure 2-bit, padded_dim = pure 4-bit).
    pub fn n4_dims(&self)    -> usize     { self.config.n4_dims }
    pub fn is_mixed(&self)   -> bool      { self.config.is_mixed() }
    pub fn config(&self)     -> &Config   { &self.config }

    /// RHDH output dimension (`dim.next_power_of_two()`).
    pub fn padded_dim(&self) -> usize { self.pipeline.rhdh.padded_dim }

    /// Bytes per packed vector (depends on bit_width and n4_dims).
    pub fn packed_bytes_per_vec(&self) -> usize { self.pipeline.packed_bytes() }

    /// Bytes in the 4-bit block per vector (0 for pure 2-bit).
    pub fn pb4_per_vec(&self) -> usize { self.pipeline.pb4 }

    /// Metric as a static string slice ("cosine" / "dot" / "l2").
    pub fn metric_str(&self) -> &'static str {
        match self.config.metric {
            Metric::Cosine => "cosine",
            Metric::Dot    => "dot",
            Metric::L2     => "l2",
        }
    }

    /// Apply the full query pipeline (std → RHDH → scale) and return the
    /// transformed query as `Vec<f32>` of length `padded_dim`.
    pub fn transform_query_slice(&self, query: &[f32]) -> Result<Vec<f32>> {
        if query.len() != self.config.dim {
            return Err(MonaVecError::DimMismatch { expected: self.config.dim, got: query.len() });
        }
        Ok(self.pipeline.transform_query(query))
    }

    /// Raw vector data for GPU / external scoring (BruteForce only).
    ///
    /// Returns `(packed_data, ids, norms)` where:
    /// - `packed_data` is a flat `u8` slice of length `n × packed_bytes_per_vec`
    /// - `ids` is a `u64` slice of length `n`
    /// - `norms` is a `f32` slice of length `n` (quant norms for score adjustment)
    ///
    /// Returns `None` for IvfFlat / HNSW index types.
    pub fn raw_vectors(&self) -> Option<(&[u8], &[u64], &[f32])> {
        if self.config.index_type != IndexType::BruteForce {
            return None;
        }
        Some((&self.store.data, &self.store.ids, &self.store.norms))
    }

    /// Update HNSW ef_search at runtime without rebuilding the index.
    /// No-op for non-HNSW index types.
    #[cfg(feature = "hnsw")]
    pub fn set_ef_search(&mut self, ef: usize) {
        if let Some(ref mut hnsw) = self.hnsw {
            hnsw.ef_search = ef;
        }
    }

    /// Update IvfFlat nprobe at runtime without rebuilding the index.
    /// No-op for non-IvfFlat index types.
    #[cfg(feature = "ivf")]
    pub fn set_nprobe(&mut self, nprobe: usize) {
        if let Some(ref mut ivf) = self.ivf {
            ivf.nprobe = nprobe;
        }
    }

    /// Add vectors to the index.
    /// `ids` and `vectors` must have the same length.
    pub fn add(&mut self, ids: &[u64], vectors: &[Vec<f32>]) -> Result<()> {
        if ids.len() != vectors.len() {
            return Err(MonaVecError::CountMismatch {
                ids:     ids.len(),
                vectors: vectors.len(),
            });
        }
        for (id, v) in ids.iter().zip(vectors.iter()) {
            if v.len() != self.config.dim {
                return Err(MonaVecError::DimMismatch {
                    expected: self.config.dim,
                    got:      v.len(),
                });
            }
            let packed = self.pipeline.encode(v);
            let norm   = self.pipeline.compute_norm(&packed);

            match self.config.index_type {
                #[cfg(feature = "hnsw")]
                IndexType::Hnsw => {
                    if let Some(ref mut hnsw) = self.hnsw {
                        let tq       = self.pipeline.transform_query(v);
                        let pipeline = &self.pipeline;
                        let metric   = self.config.metric;
                        hnsw.insert(*id, packed, norm, &tq, metric, &|q, p| pipeline.score(q, p));
                    }
                }
                #[cfg(feature = "ivf")]
                IndexType::IvfFlat => {
                    self.pre_build_transformed.push(self.pipeline.transform_query(v));
                    self.store.add(*id, &packed, norm);
                }
                _ => {
                    self.store.add(*id, &packed, norm);
                }
            }
        }
        Ok(())
    }

    /// Add vectors from a flat, contiguous f32 buffer.
    ///
    /// Identical semantics to [`add`](Self::add), but reads each vector as a
    /// `dim`-length slice of `flat` instead of a `Vec<f32>`. This lets callers
    /// (e.g. the Python binding) ingest a contiguous NumPy matrix without
    /// allocating an intermediate `Vec<Vec<f32>>`. `flat.len()` must equal
    /// `ids.len() * dim`. Encoding, ordering, and determinism are unchanged.
    pub fn add_flat(&mut self, ids: &[u64], flat: &[f32]) -> Result<()> {
        let dim = self.config.dim;
        if dim == 0 || flat.len() != ids.len() * dim {
            return Err(MonaVecError::CountMismatch {
                ids:     ids.len(),
                vectors: if dim > 0 { flat.len() / dim } else { 0 },
            });
        }
        for (id, v) in ids.iter().zip(flat.chunks_exact(dim)) {
            let packed = self.pipeline.encode(v);
            let norm   = self.pipeline.compute_norm(&packed);

            match self.config.index_type {
                #[cfg(feature = "hnsw")]
                IndexType::Hnsw => {
                    if let Some(ref mut hnsw) = self.hnsw {
                        let tq       = self.pipeline.transform_query(v);
                        let pipeline = &self.pipeline;
                        let metric   = self.config.metric;
                        hnsw.insert(*id, packed, norm, &tq, metric, &|q, p| pipeline.score(q, p));
                    }
                }
                #[cfg(feature = "ivf")]
                IndexType::IvfFlat => {
                    self.pre_build_transformed.push(self.pipeline.transform_query(v));
                    self.store.add(*id, &packed, norm);
                }
                _ => {
                    self.store.add(*id, &packed, norm);
                }
            }
        }
        Ok(())
    }

    /// Search for the K nearest neighbors of `query`.
    ///
    /// `allowlist`: if provided, only these IDs are candidates (pre-filtered).
    /// Note: for `IvfFlat`, allowlist is applied as a post-filter on the probed cells.
    pub fn search(
        &self,
        query:     &[f32],
        k:         usize,
        allowlist: Option<&[u64]>,
    ) -> Result<Vec<SearchResult>> {
        if query.len() != self.config.dim {
            return Err(MonaVecError::DimMismatch {
                expected: self.config.dim,
                got:      query.len(),
            });
        }
        if self.is_empty() {
            return Err(MonaVecError::EmptyIndex);
        }

        let tq = self.pipeline.transform_query(query);

        #[cfg(feature = "ivf")]
        if self.config.index_type == IndexType::IvfFlat {
            return self.search_ivf(&tq, k, allowlist);
        }

        #[cfg(feature = "hnsw")]
        if self.config.index_type == IndexType::Hnsw {
            return self.search_hnsw(&tq, k, allowlist);
        }

        let al      = allowlist.map(Allowlist::new);
        let results = top_k(
            &self.store,
            k,
            al.as_ref(),
            self.config.metric,
            |packed| self.pipeline.score(&tq, packed),
        );
        Ok(results)
    }

    #[cfg(feature = "ivf")]
    fn search_ivf(
        &self,
        tq:        &[f32],
        k:         usize,
        allowlist: Option<&[u64]>,
    ) -> Result<Vec<SearchResult>> {
        let ivf = self.ivf.as_ref()
            .ok_or(MonaVecError::BuildRequired)?;

        let mut results = ivf.search(tq, k, |packed| self.pipeline.score(tq, packed));

        // Post-filter for allowlist (IvfFlat pre-filter would require per-list scan).
        if let Some(al) = allowlist {
            results.retain(|r| al.contains(&r.id));
        }
        Ok(results)
    }

    #[cfg(feature = "hnsw")]
    fn search_hnsw(
        &self,
        tq:        &[f32],
        k:         usize,
        allowlist: Option<&[u64]>,
    ) -> Result<Vec<SearchResult>> {
        let hnsw = self.hnsw.as_ref()
            .ok_or(MonaVecError::IndexTypeNotImplemented("Hnsw"))?;

        let pipeline = &self.pipeline;
        let mut results = hnsw.search(tq, k, self.config.metric, &|q, p| pipeline.score(q, p));

        if let Some(al) = allowlist {
            results.retain(|r| al.contains(&r.id));
        }
        Ok(results)
    }

    /// Persist the index to a `.mvec` binary file.
    /// IvfFlat: `build()` must be called first. Saves centroid + list data after norms.
    pub fn save(&self, path: &str) -> Result<()> {
        #[cfg(feature = "ivf")]
        if self.config.index_type == IndexType::IvfFlat {
            return self.save_ivf(path);
        }
        #[cfg(feature = "hnsw")]
        if self.config.index_type == IndexType::Hnsw {
            return self.save_hnsw(path);
        }

        let file   = std::fs::File::create(path)?;
        let mut w  = BufWriter::new(file);
        let mut header = Header::from_config(&self.config, self.store.len());
        if self.pipeline.std_params.is_some() { header.has_std = true; }
        header.write(&mut w)?;
        write_std_params_if_set(&mut w, &self.pipeline.std_params)?;
        w.write_all(&self.store.data)?;
        for &id in &self.store.ids {
            w.write_all(&id.to_le_bytes())?;
        }
        for &n in &self.store.norms {
            w.write_all(&n.to_le_bytes())?;
        }
        Ok(())
    }

    #[cfg(feature = "hnsw")]
    fn save_hnsw(&self, path: &str) -> Result<()> {
        let hnsw   = self.hnsw.as_ref().ok_or(MonaVecError::IndexTypeNotImplemented("Hnsw"))?;
        let file   = std::fs::File::create(path)?;
        let mut w  = BufWriter::new(file);
        let mut header = Header::from_config(&self.config, hnsw.len());
        if self.pipeline.std_params.is_some() { header.has_std = true; }
        header.write(&mut w)?;
        write_std_params_if_set(&mut w, &self.pipeline.std_params)?;
        hnsw.write_data(&mut w).map_err(MonaVecError::Io)?;
        Ok(())
    }

    #[cfg(feature = "ivf")]
    fn save_ivf(&self, path: &str) -> Result<()> {
        let ivf = self.ivf.as_ref().ok_or(MonaVecError::BuildRequired)?;
        let file   = std::fs::File::create(path)?;
        let mut w  = BufWriter::new(file);
        let mut header = Header::from_config(&self.config, ivf.total_vectors());
        if self.pipeline.std_params.is_some() { header.has_std = true; }
        header.write(&mut w)?;
        write_std_params_if_set(&mut w, &self.pipeline.std_params)?;
        ivf.write_data(&mut w)?;
        Ok(())
    }

    /// Load an index from a `.mvec` file (read into memory).
    /// Supports v1–v6. For IvfFlat, reads centroid+list data directly — no rebuild.
    pub fn load(path: &str) -> Result<Self> {
        let file   = std::fs::File::open(path)?;
        let mut r  = BufReader::new(file);
        let header = Header::read(&mut r)?;
        let config = header.to_config()?;
        let pb     = {
            let tmp = QuantPipeline::new(&config)?;
            tmp.packed_bytes()
        };

        // Read std params if present (v6+), then build the pipeline.
        let std_params = if header.has_std {
            Some(read_std_params(&mut r, config.dim)?)
        } else {
            None
        };
        let mut pipeline = QuantPipeline::new(&config)?;
        if let Some(p) = std_params { pipeline.set_std_params(p); }

        #[cfg(feature = "ivf")]
        if config.index_type == IndexType::IvfFlat {
            let ivf = IvfFlat::read_data(&mut r, pb, config.ivf_nprobe, config.metric)
                .map_err(MonaVecError::Io)?;
            return Ok(Self {
                config,
                pipeline,
                store: VectorStore::new(pb),
                ivf:                   Some(ivf),
                pre_build_transformed: Vec::new(),
                #[cfg(feature = "hnsw")]
                hnsw: None,
            });
        }

        #[cfg(feature = "hnsw")]
        if config.index_type == IndexType::Hnsw {
            // v4 HNSW stored n_nb as u8 (max 255); v5 changed to u16. Loading v4 HNSW
            // would silently read corrupt neighbor counts — reject it explicitly.
            if header.version == 4 {
                return Err(MonaVecError::InvalidFormat(
                    "HNSW index version 4 is incompatible with v5 format (neighbor count type changed u8→u16). Rebuild the index with the current version.".into()
                ));
            }
            let ef_s = if config.hnsw_ef_search == 0 { config.hnsw_m * 2 } else { config.hnsw_ef_search };
            let hnsw = Hnsw::read_data(&mut r, pb, config.hnsw_m, config.hnsw_ef_construction, ef_s)
                .map_err(MonaVecError::Io)?;
            return Ok(Self {
                config,
                pipeline,
                store: VectorStore::new(pb),
                #[cfg(feature = "ivf")]
                ivf:                   None,
                #[cfg(feature = "ivf")]
                pre_build_transformed: Vec::new(),
                hnsw: Some(hnsw),
            });
        }

        let count = header.count as usize;
        let mut data  = vec![0u8; count * pb];
        r.read_exact(&mut data)?;

        let mut ids    = Vec::with_capacity(count);
        let mut id_buf = [0u8; 8];
        for _ in 0..count {
            r.read_exact(&mut id_buf)?;
            ids.push(u64::from_le_bytes(id_buf));
        }

        // v2+: read per-vector norms; v1: norms stay empty (uncorrected scoring).
        let mut norms = Vec::with_capacity(count);
        if header.version >= 2 {
            let mut norm_buf = [0u8; 4];
            for _ in 0..count {
                r.read_exact(&mut norm_buf)?;
                norms.push(f32::from_le_bytes(norm_buf));
            }
        }

        Ok(Self {
            config,
            pipeline,
            store: VectorStore { data, ids, norms, packed_bytes: pb },
            #[cfg(feature = "ivf")]
            ivf:                    None,
            #[cfg(feature = "ivf")]
            pre_build_transformed:  Vec::new(),
            #[cfg(feature = "hnsw")]
            hnsw: None,
        })
    }

    /// Load an index using memory-mapped IO. Same result as `load` but the
    /// OS maps the file and pages data in on demand — better for large indexes.
    /// IvfFlat and HNSW fall back to buffered `load` (their data isn't a flat contiguous slice).
    pub fn load_mmap(path: &str) -> Result<Self> {
        // IvfFlat / HNSW are not flat contiguous slices, so they fall back to
        // buffered load(). Only peek at the header when those features are built;
        // for a plain BruteForce build the peek is unnecessary.
        #[cfg(any(feature = "ivf", feature = "hnsw"))]
        {
            let file = std::fs::File::open(path)?;
            let mut r = BufReader::new(file);
            let header = Header::read(&mut r)?;
            let config = header.to_config()?;
            #[cfg(feature = "ivf")]
            if config.index_type == IndexType::IvfFlat {
                return Self::load(path);
            }
            #[cfg(feature = "hnsw")]
            if config.index_type == IndexType::Hnsw {
                return Self::load(path);
            }
        }

        use crate::io::mmap::MmapView;
        use std::io::Cursor;

        let view   = MmapView::open(path)?;
        let bytes  = &view.map[..];
        let mut r  = Cursor::new(bytes);
        let header = Header::read(&mut r)?;
        let config = header.to_config()?;
        let pb     = {
            let tmp = QuantPipeline::new(&config)?;
            tmp.packed_bytes()
        };
        let count = header.count as usize;

        // Read std params from bytes (if v6 + has_std).
        // header.byte_size() already includes the std params bytes, so data_start is correct.
        let std_params: Option<StdParams> = if header.has_std {
            use crate::io::format::BYTE_SIZE_V4;
            let fixed = BYTE_SIZE_V4; // std params start right after the 56B fixed header
            let dim   = config.dim;
            let mean: Vec<f32>    = bytes[fixed..fixed + dim * 4]
                .chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect();
            let inv_std: Vec<f32> = bytes[fixed + dim * 4..fixed + dim * 8]
                .chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect();
            Some(StdParams { mean, inv_std })
        } else {
            None
        };
        let mut pipeline = QuantPipeline::new(&config)?;
        if let Some(p) = std_params { pipeline.set_std_params(p); }

        let data_start = header.byte_size();
        let id_start   = data_start + count * pb;
        let norm_start = id_start + count * 8;

        let data: Vec<u8> = bytes[data_start..id_start].to_vec();
        let mut ids = Vec::with_capacity(count);
        for chunk in bytes[id_start..norm_start].chunks_exact(8) {
            ids.push(u64::from_le_bytes(chunk.try_into().unwrap()));
        }

        // v2+: read norms; v1: empty (uncorrected scoring).
        let mut norms = Vec::with_capacity(count);
        if header.version >= 2 && bytes.len() >= norm_start + count * 4 {
            for chunk in bytes[norm_start..norm_start + count * 4].chunks_exact(4) {
                norms.push(f32::from_le_bytes(chunk.try_into().unwrap()));
            }
        }

        Ok(Self {
            config,
            pipeline,
            store: VectorStore { data, ids, norms, packed_bytes: pb },
            #[cfg(feature = "ivf")]
            ivf:                    None,
            #[cfg(feature = "ivf")]
            pre_build_transformed:  Vec::new(),
            #[cfg(feature = "hnsw")]
            hnsw: None,
        })
    }
}

fn write_std_params_if_set<W: Write>(w: &mut W, params: &Option<StdParams>) -> Result<()> {
    if let Some(p) = params {
        for &m in &p.mean    { w.write_all(&m.to_le_bytes())?; }
        for &s in &p.inv_std { w.write_all(&s.to_le_bytes())?; }
    }
    Ok(())
}

fn read_std_params<R: Read>(r: &mut R, dim: usize) -> Result<StdParams> {
    let mut f32_buf = [0u8; 4];
    let mut mean    = Vec::with_capacity(dim);
    let mut inv_std = Vec::with_capacity(dim);
    for _ in 0..dim {
        r.read_exact(&mut f32_buf)?;
        mean.push(f32::from_le_bytes(f32_buf));
    }
    for _ in 0..dim {
        r.read_exact(&mut f32_buf)?;
        inv_std.push(f32::from_le_bytes(f32_buf));
    }
    Ok(StdParams { mean, inv_std })
}

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

    fn rand_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 add_and_search_basic() {
        let config = Config::new(64, Metric::Cosine, 4);
        let mut index = MonaVec::new(config).unwrap();

        let ids: Vec<u64> = (0..10).collect();
        let vecs: Vec<Vec<f32>> = (0..10).map(|i| rand_vec(64, i)).collect();

        index.add(&ids, &vecs).unwrap();
        assert_eq!(index.len(), 10);

        let results = index.search(&vecs[0], 3, None).unwrap();
        assert_eq!(results.len(), 3);
        assert_eq!(results[0].id, 0);
    }

    #[test]
    fn search_with_allowlist() {
        let config = Config::new(64, Metric::Cosine, 4);
        let mut index = MonaVec::new(config).unwrap();

        let ids: Vec<u64> = (0..20).collect();
        let vecs: Vec<Vec<f32>> = (0..20).map(|i| rand_vec(64, i)).collect();
        index.add(&ids, &vecs).unwrap();

        let allowlist = vec![1u64, 5, 10, 15];
        let results = index.search(&vecs[0], 10, Some(&allowlist)).unwrap();

        for r in &results {
            assert!(allowlist.contains(&r.id), "unexpected id: {}", r.id);
        }
        assert!(results.len() <= 4);
    }

    #[test]
    fn dim_mismatch_error() {
        let config = Config::new(64, Metric::Cosine, 4);
        let mut index = MonaVec::new(config).unwrap();
        let v = rand_vec(32, 0);
        let result = index.add(&[0], &[v]);
        assert!(matches!(result, Err(MonaVecError::DimMismatch { .. })));
    }

    #[test]
    fn empty_index_search_error() {
        let config = Config::new(64, Metric::Cosine, 4);
        let index  = MonaVec::new(config).unwrap();
        let q      = rand_vec(64, 0);
        let result = index.search(&q, 5, None);
        assert!(matches!(result, Err(MonaVecError::EmptyIndex)));
    }

    #[test]
    fn save_load_roundtrip() {
        let config = Config::new(64, Metric::Cosine, 4);
        let mut index = MonaVec::new(config).unwrap();
        let ids  = vec![10u64, 20, 30];
        let vecs = vec![rand_vec(64, 1), rand_vec(64, 2), rand_vec(64, 3)];
        index.add(&ids, &vecs).unwrap();

        let query   = rand_vec(64, 99);
        let before  = index.search(&query, 3, None).unwrap();

        let path = "/tmp/monavec_test_roundtrip.mvec";
        index.save(path).unwrap();

        let loaded = MonaVec::load(path).unwrap();
        assert_eq!(loaded.len(), 3);
        let after = loaded.search(&query, 3, None).unwrap();

        assert_eq!(before.len(), after.len());
        for (b, a) in before.iter().zip(after.iter()) {
            assert_eq!(b.id, a.id, "id mismatch after load");
            assert!((b.score - a.score).abs() < 1e-5, "score mismatch: {b:?} vs {a:?}");
        }
        std::fs::remove_file(path).ok();
    }

    #[test]
    fn save_load_mmap_roundtrip() {
        let config = Config::new(64, Metric::Cosine, 4);
        let mut index = MonaVec::new(config).unwrap();
        let ids  = vec![7u64, 8, 9];
        let vecs = vec![rand_vec(64, 7), rand_vec(64, 8), rand_vec(64, 9)];
        index.add(&ids, &vecs).unwrap();

        let path  = "/tmp/monavec_test_mmap.mvec";
        index.save(path).unwrap();
        let loaded = MonaVec::load_mmap(path).unwrap();
        assert_eq!(loaded.len(), 3);
        std::fs::remove_file(path).ok();
    }

    #[test]
    fn two_bit_encoding() {
        let config = Config::new(64, Metric::Cosine, 2);
        let mut index = MonaVec::new(config).unwrap();
        let ids  = vec![0u64, 1, 2];
        let vecs = vec![rand_vec(64, 0), rand_vec(64, 1), rand_vec(64, 2)];
        index.add(&ids, &vecs).unwrap();
        let results = index.search(&vecs[0], 1, None).unwrap();
        assert_eq!(results[0].id, 0);
    }

    #[test]
    fn mixed_encoding_search_and_roundtrip() {
        let config = Config::new(64, Metric::Cosine, 4).with_mixed_ratio(0.5);
        assert!(config.is_mixed());

        let mut index = MonaVec::new(config).unwrap();
        let ids  = vec![0u64, 1, 2, 3, 4];
        let vecs: Vec<Vec<f32>> = (0..5).map(|i| rand_vec(64, i)).collect();
        index.add(&ids, &vecs).unwrap();

        let r = index.search(&vecs[0], 1, None).unwrap();
        assert_eq!(r[0].id, 0, "top-1 should be self");

        let path = "/tmp/monavec_test_mixed.mvec";
        index.save(path).unwrap();
        let loaded = MonaVec::load(path).unwrap();
        assert_eq!(loaded.len(), 5);
        assert!(loaded.is_mixed());
        assert_eq!(loaded.n4_dims(), index.n4_dims());

        let r2 = loaded.search(&vecs[0], 1, None).unwrap();
        assert_eq!(r2[0].id, r[0].id);
        assert!((r2[0].score - r[0].score).abs() < 1e-5, "score mismatch after load");
        std::fs::remove_file(path).ok();
    }

    // --- L2 fit() tests ---

    fn raw_l2_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);
        // Non-unit, non-zero-mean — simulates pixel/SIFT data
        (0..dim).map(|i| {
            let base = (i % 8) as f32 * 30.0; // per-dim mean variation
            base + (rng.next_u32() as f32 / u32::MAX as f32) * 50.0
        }).collect()
    }

    #[test]
    fn fit_l2_self_score_highest() {
        let dim = 64usize;
        let config = Config::new(dim, Metric::L2, 4);
        let mut idx = MonaVec::new(config).unwrap();

        let sample: Vec<Vec<f32>> = (0..100).map(|s| raw_l2_vec(dim, s)).collect();
        idx.fit(&sample).unwrap();

        let ids:  Vec<u64>      = (0..20).collect();
        let vecs: Vec<Vec<f32>> = (0..20).map(|i| raw_l2_vec(dim, i + 200)).collect();
        idx.add(&ids, &vecs).unwrap();

        for (i, v) in vecs.iter().enumerate() {
            let results = idx.search(v, 1, None).unwrap();
            assert_eq!(results[0].id, i as u64, "top-1 should be self for L2 + fit");
        }
    }

    #[test]
    fn fit_noop_for_cosine() {
        let config = Config::new(32, Metric::Cosine, 4);
        let mut idx = MonaVec::new(config).unwrap();
        let sample: Vec<Vec<f32>> = (0..10).map(|s| raw_l2_vec(32, s)).collect();
        // Should not error and should not change behaviour
        idx.fit(&sample).unwrap();
        assert!(idx.pipeline.std_params.is_none());
    }

    #[test]
    fn fit_after_add_is_error() {
        let config = Config::new(32, Metric::L2, 4);
        let mut idx = MonaVec::new(config).unwrap();
        let v = raw_l2_vec(32, 0);
        idx.add(&[0], &[v.clone()]).unwrap();
        let sample = vec![v];
        let result = idx.fit(&sample);
        assert!(matches!(result, Err(MonaVecError::FitAfterAdd)));
    }

    #[test]
    fn fit_l2_save_load_roundtrip() {
        let dim = 64usize;
        let config = Config::new(dim, Metric::L2, 4);
        let mut idx = MonaVec::new(config).unwrap();

        let sample: Vec<Vec<f32>> = (0..50).map(|s| raw_l2_vec(dim, s)).collect();
        idx.fit(&sample).unwrap();

        let ids:  Vec<u64>      = (0..10).collect();
        let vecs: Vec<Vec<f32>> = (0..10).map(|i| raw_l2_vec(dim, i + 300)).collect();
        idx.add(&ids, &vecs).unwrap();

        let query  = raw_l2_vec(dim, 999);
        let before = idx.search(&query, 5, None).unwrap();

        let path = "/tmp/monavec_test_l2_fit.mvec";
        idx.save(path).unwrap();

        let loaded = MonaVec::load(path).unwrap();
        assert!(loaded.pipeline.std_params.is_some(), "std_params should survive save/load");
        let after = loaded.search(&query, 5, None).unwrap();

        assert_eq!(before.len(), after.len());
        for (b, a) in before.iter().zip(after.iter()) {
            assert_eq!(b.id, a.id, "id mismatch after L2 fit save/load");
            assert!((b.score - a.score).abs() < 1e-5, "score mismatch: {:?} vs {:?}", b, a);
        }
        std::fs::remove_file(path).ok();
    }

    #[test]
    fn fit_l2_mmap_roundtrip() {
        let dim = 64usize;
        let config = Config::new(dim, Metric::L2, 4);
        let mut idx = MonaVec::new(config).unwrap();

        let sample: Vec<Vec<f32>> = (0..50).map(|s| raw_l2_vec(dim, s)).collect();
        idx.fit(&sample).unwrap();

        let ids:  Vec<u64>      = (0..10).collect();
        let vecs: Vec<Vec<f32>> = (0..10).map(|i| raw_l2_vec(dim, i + 400)).collect();
        idx.add(&ids, &vecs).unwrap();

        let path = "/tmp/monavec_test_l2_mmap.mvec";
        idx.save(path).unwrap();

        let loaded = MonaVec::load_mmap(path).unwrap();
        assert!(loaded.pipeline.std_params.is_some(), "std_params should survive mmap load");

        let query  = raw_l2_vec(dim, 1001);
        let r1 = idx.search(&query, 5, None).unwrap();
        let r2 = loaded.search(&query, 5, None).unwrap();
        assert_eq!(r1.len(), r2.len());
        for (a, b) in r1.iter().zip(r2.iter()) {
            assert_eq!(a.id, b.id, "mmap: id mismatch after L2 fit load");
        }
        std::fs::remove_file(path).ok();
    }

    #[cfg(feature = "ivf")]
    #[test]
    fn ivf_build_required_before_search() {
        let config = Config::new(64, Metric::Cosine, 4).with_ivf(8, 2);
        let mut idx = MonaVec::new(config).unwrap();
        let ids:  Vec<u64>      = (0..20).collect();
        let vecs: Vec<Vec<f32>> = (0..20).map(|i| rand_vec(64, i)).collect();
        idx.add(&ids, &vecs).unwrap();
        let r = idx.search(&vecs[0], 3, None);
        assert!(matches!(r, Err(crate::MonaVecError::BuildRequired)));
    }

    #[cfg(feature = "ivf")]
    #[test]
    fn ivf_recall_at_10() {
        let dim = 64usize;
        let n   = 500usize;
        let k   = 10usize;

        let config = Config::new(dim, Metric::Cosine, 4).with_ivf(8, 4); // 50% cells probed
        let mut idx = MonaVec::new(config).unwrap();
        let ids:  Vec<u64>      = (0..n as u64).collect();
        let vecs: Vec<Vec<f32>> = (0..n).map(|i| rand_vec(dim, i as u64)).collect();
        idx.add(&ids, &vecs).unwrap();
        idx.build().unwrap();

        let bf_config = Config::new(dim, Metric::Cosine, 4);
        let mut bf = MonaVec::new(bf_config).unwrap();
        bf.add(&ids, &vecs).unwrap();

        let mut hits = 0usize;
        let queries = 20usize;
        for qi in 0..queries {
            let q = rand_vec(dim, (n + qi) as u64);
            let ivf_res: Vec<u64> = idx.search(&q, k, None).unwrap().iter().map(|r| r.id).collect();
            let bf_res:  Vec<u64> = bf.search(&q, k, None).unwrap().iter().map(|r| r.id).collect();
            hits += ivf_res.iter().filter(|id| bf_res.contains(id)).count();
        }
        let recall = hits as f32 / (queries * k) as f32;
        assert!(recall >= 0.60, "IvfFlat recall@10 = {:.2}, expected >= 0.60", recall);
    }

    #[cfg(feature = "ivf")]
    #[test]
    fn ivf_save_load_roundtrip() {
        let dim = 32usize;
        let n   = 100usize;

        let config = Config::new(dim, Metric::Cosine, 4).with_ivf(8, 3);
        let mut idx = MonaVec::new(config).unwrap();
        let ids:  Vec<u64>      = (0..n as u64).collect();
        let vecs: Vec<Vec<f32>> = (0..n).map(|i| rand_vec(dim, i as u64)).collect();
        idx.add(&ids, &vecs).unwrap();
        idx.build().unwrap();

        let path = "/tmp/monavec_test_ivf.mvec";
        idx.save(path).unwrap();
        let loaded = MonaVec::load(path).unwrap();
        assert_eq!(loaded.len(), n);

        let q = rand_vec(dim, 9999);
        let r1 = idx.search(&q, 5, None).unwrap();
        let r2 = loaded.search(&q, 5, None).unwrap();
        assert_eq!(r1.len(), r2.len());
        for (a, b) in r1.iter().zip(r2.iter()) {
            assert_eq!(a.id, b.id, "id mismatch after IvfFlat load");
        }
        std::fs::remove_file(path).ok();
    }

    #[cfg(feature = "hnsw")]
    #[test]
    fn hnsw_self_search_top1() {
        let dim = 64usize;
        let n   = 100usize;

        let config = Config::new(dim, Metric::Cosine, 4).with_hnsw(16, 100);
        let mut idx = MonaVec::new(config).unwrap();
        let ids:  Vec<u64>      = (0..n as u64).collect();
        let vecs: Vec<Vec<f32>> = (0..n).map(|i| rand_vec(dim, i as u64)).collect();
        idx.add(&ids, &vecs).unwrap();

        let mut correct = 0;
        for (i, v) in vecs.iter().enumerate() {
            let r = idx.search(v, 1, None).unwrap();
            if !r.is_empty() && r[0].id == i as u64 { correct += 1; }
        }
        // HNSW with quantization: threshold at 60% (quantization + greedy search introduce noise)
        assert!(correct >= n * 6 / 10, "HNSW self-search: {}/{} correct", correct, n);
    }

    #[cfg(feature = "hnsw")]
    #[test]
    fn hnsw_save_load_roundtrip() {
        let dim = 32usize;
        let n   = 50usize;

        let config = Config::new(dim, Metric::Cosine, 4).with_hnsw(4, 20);
        let mut idx = MonaVec::new(config).unwrap();
        let ids:  Vec<u64>      = (0..n as u64).collect();
        let vecs: Vec<Vec<f32>> = (0..n).map(|i| rand_vec(dim, i as u64)).collect();
        idx.add(&ids, &vecs).unwrap();
        assert_eq!(idx.len(), n);

        let path = "/tmp/monavec_test_hnsw.mvec";
        idx.save(path).unwrap();
        let loaded = MonaVec::load(path).unwrap();
        assert_eq!(loaded.len(), n);

        let q  = rand_vec(dim, 9999);
        let r1 = idx.search(&q, 5, None).unwrap();
        let r2 = loaded.search(&q, 5, None).unwrap();
        assert_eq!(r1.len(), r2.len());
        for (a, b) in r1.iter().zip(r2.iter()) {
            assert_eq!(a.id, b.id, "id mismatch after HNSW load");
        }
        std::fs::remove_file(path).ok();
    }
}
