pub mod rotation;
pub mod scalar;
pub mod tables;

use crate::{Config, Metric, error::Result, error::MonaVecError};
use rotation::Rhdh;
use scalar::{encode_4bit, encode_2bit, packed_bytes_4bit, packed_bytes_2bit};
use tables::{quant_norm_4bit, quant_norm_2bit, quant_norm_mixed};

/// Global standardization parameters for L2 metric.
///
/// Uses a single (scalar) mean and inv_std computed across ALL values in ALL
/// dimensions of the sample corpus. This is a uniform scaling:
///   `x_std = (x - mean) / std`
///
/// Crucially, applying the same scalar factor to every dimension preserves
/// relative L2 distances — nearest-neighbour ordering in the standardized
/// space is identical to the original space. Per-dimension whitening would
/// change the metric (Mahalanobis distance) and degrade L2 recall.
///
/// Stored in `.mvec` v6+ files as `mean[dim] || inv_std[dim]` where all
/// elements of each array hold the same scalar value.
pub struct StdParams {
    pub mean:    Vec<f32>,
    pub inv_std: Vec<f32>,
}

impl StdParams {
    /// Compute global (scalar) mean and inverse-std over all values in `samples`.
    ///
    /// Both output vectors have length `dim` with every element equal to the
    /// same scalar, so the existing `apply()` / save / load paths are unchanged.
    pub fn compute(samples: &[impl AsRef<[f32]>], dim: usize) -> Self {
        let mut sum  = 0f64;
        let mut sum2 = 0f64;
        let mut count = 0usize;
        for s in samples {
            for &x in s.as_ref().iter().take(dim) {
                let xf = x as f64;
                sum  += xf;
                sum2 += xf * xf;
                count += 1;
            }
        }
        let n        = count as f64;
        let mean_val = sum / n;
        let var      = (sum2 / n - mean_val * mean_val).max(0.0);
        let inv_std_val = if var > 1e-10 { (1.0 / var.sqrt()) as f32 } else { 1.0 };
        let mean_f = mean_val as f32;
        Self {
            mean:    vec![mean_f;       dim],
            inv_std: vec![inv_std_val;  dim],
        }
    }

    /// Apply `(x - mean) / std` element-wise.
    #[inline]
    pub fn apply(&self, v: &[f32]) -> Vec<f32> {
        v.iter()
            .zip(self.mean.iter())
            .zip(self.inv_std.iter())
            .map(|((&x, &m), &s)| (x - m) * s)
            .collect()
    }
}

/// Full quantization pipeline: normalize → RHDH rotate → scale → pack.
///
/// Supports three modes:
/// - Pure 4-bit:  `n4 == padded_dim` — all dims at 4-bit (default)
/// - Pure 2-bit:  `n4 == 0`          — all dims at 2-bit
/// - Mixed:       `0 < n4 < padded_dim` — first n4 dims at 4-bit, rest at 2-bit
///
/// Mixed mode packed layout per vector: [4-bit block (pb4 bytes) | 2-bit block (pb2 bytes)]
pub struct QuantPipeline {
    pub dim:       usize,
    pub bit_width: u8,
    pub metric:    Metric,
    pub rhdh:      Rhdh,
    /// Multiplier to normalize rotated coords from N(0, 1/padded_dim) to N(0, 1).
    scale:         f32,
    /// Dims at 4-bit in padded space. 0 = pure 2-bit, padded_dim = pure 4-bit.
    pub n4:        usize,
    /// Dims at 2-bit in padded space.
    pub n2:        usize,
    /// Bytes for the 4-bit block per vector.
    pub pb4:       usize,
    /// Per-dimension standardization for L2 metric. None = raw pass-through.
    pub std_params: Option<StdParams>,
}

impl QuantPipeline {
    pub fn new(config: &Config) -> Result<Self> {
        let rhdh       = Rhdh::new(config.dim, config.seed);
        let padded     = rhdh.padded_dim;
        let n4_dims    = config.n4_dims;

        // n4_dims is the single source of truth for quantization mode.
        // bit_width is kept for file format and legacy API compatibility only.
        if n4_dims > padded {
            return Err(MonaVecError::InvalidBitWidth(config.bit_width));
        }
        // Mixed mode requires n4_dims to be a multiple of 4 for byte alignment.
        if n4_dims > 0 && n4_dims < padded && n4_dims % 4 != 0 {
            return Err(MonaVecError::InvalidBitWidth(config.bit_width));
        }

        let scale = (padded as f32).sqrt();
        let n4    = n4_dims;
        let n2    = padded - n4;
        let pb4   = packed_bytes_4bit(n4);

        Ok(Self { dim: config.dim, bit_width: config.bit_width, metric: config.metric, rhdh, scale, n4, n2, pb4, std_params: None })
    }

    /// True when some dims are at 4-bit and others at 2-bit.
    pub fn is_mixed(&self) -> bool {
        self.n4 > 0 && self.n2 > 0
    }

    /// Returns the number of bytes one packed vector occupies.
    pub fn packed_bytes(&self) -> usize {
        self.pb4 + packed_bytes_2bit(self.n2)
    }

    /// Encode a raw f32 vector through the full pipeline.
    /// Cosine: normalizes to unit length before rotation (magnitude is irrelevant).
    /// Dot / L2: uses raw vector (magnitude carries signal).
    pub fn encode(&self, v: &[f32]) -> Vec<u8> {
        debug_assert_eq!(v.len(), self.dim);
        let input   = self.prepare_input(v);
        let rotated = self.rhdh.apply(&input);
        let scaled: Vec<f32> = rotated.iter().map(|&x| x * self.scale).collect();

        if !self.is_mixed() {
            return if self.n4 > 0 {
                encode_4bit(&scaled)
            } else {
                encode_2bit(&scaled)
            };
        }

        // Mixed: [4-bit block | 2-bit block]
        let mut out = encode_4bit(&scaled[..self.n4]);
        out.extend_from_slice(&encode_2bit(&scaled[self.n4..]));
        out
    }

    /// Transform a query vector for asymmetric scoring.
    /// Returns padded_dim f32 values (matches packed vector dimensionality).
    pub fn transform_query(&self, query: &[f32]) -> Vec<f32> {
        debug_assert_eq!(query.len(), self.dim);
        let input   = self.prepare_input(query);
        let rotated = self.rhdh.apply(&input);
        rotated.iter().map(|&x| x * self.scale).collect()
    }

    /// Compute per-dimension standardization params from `samples` and store them.
    ///
    /// Must be called before `encode()` / `transform_query()`. No-op for non-L2 metrics.
    /// `n_indexed` must be 0 — calling after vectors have been added is an error.
    pub fn fit(&mut self, samples: &[impl AsRef<[f32]>], n_indexed: usize) -> Result<()> {
        if self.metric != Metric::L2 {
            return Ok(());
        }
        if n_indexed > 0 {
            return Err(MonaVecError::FitAfterAdd);
        }
        if samples.is_empty() {
            return Err(MonaVecError::EmptySample);
        }
        for (i, s) in samples.iter().enumerate() {
            if s.as_ref().len() != self.dim {
                return Err(MonaVecError::DimMismatch { expected: self.dim, got: s.as_ref().len() });
            }
            let _ = i;
        }
        self.std_params = Some(StdParams::compute(samples, self.dim));
        Ok(())
    }

    /// Directly inject pre-computed standardization params (used by `load()`).
    pub fn set_std_params(&mut self, params: StdParams) {
        self.std_params = Some(params);
    }

    /// Cosine: normalize to unit length. L2: apply per-dim standardization if fitted. Dot: raw.
    #[inline]
    fn prepare_input(&self, v: &[f32]) -> Vec<f32> {
        match self.metric {
            Metric::Cosine => crate::metric::cosine::normalized_copy(v),
            Metric::Dot    => v.to_vec(),
            Metric::L2     => match &self.std_params {
                Some(p) => p.apply(v),
                None    => v.to_vec(),
            },
        }
    }

    /// Score a transformed query against one packed DB vector.
    /// Dispatches to AVX-512/AVX2/NEON/scalar based on runtime CPU features.
    pub fn score(&self, transformed_query: &[f32], db_packed: &[u8]) -> f32 {
        if self.is_mixed() {
            return crate::simd::score_mixed(transformed_query, db_packed, self.n4);
        }
        if self.n4 > 0 {
            crate::simd::score_4bit(transformed_query, db_packed)
        } else {
            crate::simd::score_2bit(transformed_query, db_packed)
        }
    }

    /// Compute the quantized-vector norm for length-renormalized scoring.
    pub fn compute_norm(&self, packed: &[u8]) -> f32 {
        if self.is_mixed() {
            return quant_norm_mixed(packed, self.pb4);
        }
        if self.n4 > 0 {
            quant_norm_4bit(packed)
        } else {
            quant_norm_2bit(packed)
        }
    }

    /// Profile per-dimension variance of RHDH-rotated embeddings.
    ///
    /// Rotates each sample vector through normalize → RHDH → scale and computes
    /// empirical variance per output dimension. With ideal RHDH all values will be
    /// close to 1.0; higher values identify dimensions that benefit most from 4-bit.
    ///
    /// `sample` must be non-empty; all vectors must have length `config.dim`.
    pub fn profile_variance(config: &Config, sample: &[impl AsRef<[f32]>]) -> Result<Vec<f32>> {
        if sample.is_empty() {
            return Err(MonaVecError::EmptySample);
        }
        let rhdh  = Rhdh::new(config.dim, config.seed);
        let scale = (rhdh.padded_dim as f32).sqrt();
        let pad   = rhdh.padded_dim;
        let n     = sample.len() as f64;

        let mut sum1 = vec![0f64; pad];
        let mut sum2 = vec![0f64; pad];

        for v in sample {
            let v = v.as_ref();
            if v.len() != config.dim {
                return Err(MonaVecError::DimMismatch { expected: config.dim, got: v.len() });
            }
            let input   = match config.metric {
                Metric::Cosine => crate::metric::cosine::normalized_copy(v),
                Metric::Dot | Metric::L2 => v.to_vec(),
            };
            let rotated = rhdh.apply(&input);
            for (i, &x) in rotated.iter().enumerate() {
                let xs = (x * scale) as f64;
                sum1[i] += xs;
                sum2[i] += xs * xs;
            }
        }

        Ok((0..pad)
            .map(|i| {
                let mean = sum1[i] / n;
                ((sum2[i] / n) - mean * mean).max(0.0) as f32
            })
            .collect())
    }

    /// Compute the optimal `n4_dims` for a given average bit budget via water-filling.
    ///
    /// Water-filling on the {2, 4}-bit discrete alphabet assigns 4-bit capacity to
    /// the highest-variance dimensions first. The formula for the total count is:
    ///   `n4 = padded_dim × (avg_bits − 2) / 2`
    ///
    /// `avg_bits` is clamped to [2.0, 4.0]. Result is rounded to a multiple of 4
    /// (required for byte alignment).
    ///
    /// Pass the output of `profile_variance` as `variances`.
    pub fn optimal_n4_dims(variances: &[f32], avg_bits: f32) -> usize {
        let padded   = variances.len();
        let avg_bits = avg_bits.clamp(2.0, 4.0);
        let n4_raw   = (padded as f32 * (avg_bits - 2.0) / 2.0).round() as usize;
        (n4_raw / 4) * 4
    }

    /// Returns dimension indices sorted by variance descending.
    ///
    /// The first `n4_dims` entries should receive 4-bit encoding for optimal allocation.
    /// Currently informational; a future format version will store this permutation so
    /// high-variance dims are physically placed in the 4-bit block.
    pub fn variance_perm(variances: &[f32]) -> Vec<u16> {
        let mut idx: Vec<u16> = (0..variances.len() as u16).collect();
        idx.sort_unstable_by(|&a, &b| {
            variances[b as usize]
                .partial_cmp(&variances[a as usize])
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        idx
    }
}  // end impl QuantPipeline

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{Config, Metric};
    fn random_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(|_| {
            let u = rng.next_u32() as f32 / u32::MAX as f32;
            u * 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 encode_decode_same_length() {
        let config = Config::new(64, Metric::Cosine, 4);
        let pipe   = QuantPipeline::new(&config).unwrap();
        let v      = random_unit_vec(64, 0);
        let packed = pipe.encode(&v);
        assert_eq!(packed.len(), pipe.packed_bytes());
    }

    #[test]
    fn self_score_highest() {
        let dim    = 128;
        let config = Config::new(dim, Metric::Cosine, 4);
        let pipe   = QuantPipeline::new(&config).unwrap();

        let query   = random_unit_vec(dim, 1);
        let other1  = random_unit_vec(dim, 2);
        let other2  = random_unit_vec(dim, 3);

        let tq      = pipe.transform_query(&query);
        let packed_self   = pipe.encode(&query);
        let packed_other1 = pipe.encode(&other1);
        let packed_other2 = pipe.encode(&other2);

        let s_self   = pipe.score(&tq, &packed_self);
        let s_other1 = pipe.score(&tq, &packed_other1);
        let s_other2 = pipe.score(&tq, &packed_other2);

        assert!(s_self >= s_other1, "self={s_self}, other1={s_other1}");
        assert!(s_self >= s_other2, "self={s_self}, other2={s_other2}");
    }

    #[test]
    fn mixed_encode_correct_length() {
        let dim    = 128;
        let config = Config::new(dim, Metric::Cosine, 4).with_mixed_ratio(0.5);
        assert!(config.is_mixed(), "should be mixed");
        let pipe   = QuantPipeline::new(&config).unwrap();
        let v      = random_unit_vec(dim, 0);
        let packed = pipe.encode(&v);
        assert_eq!(packed.len(), pipe.packed_bytes());
        let pipe4  = QuantPipeline::new(&Config::new(dim, Metric::Cosine, 4)).unwrap();
        let pipe2  = QuantPipeline::new(&Config::new(dim, Metric::Cosine, 2)).unwrap();
        assert!(packed.len() < pipe4.packed_bytes());
        assert!(packed.len() > pipe2.packed_bytes());
    }

    #[test]
    fn mixed_self_score_highest() {
        let dim    = 128;
        let config = Config::new(dim, Metric::Cosine, 4).with_mixed_ratio(0.5);
        let pipe   = QuantPipeline::new(&config).unwrap();

        let query  = random_unit_vec(dim, 10);
        let other1 = random_unit_vec(dim, 11);
        let other2 = random_unit_vec(dim, 12);

        let tq = pipe.transform_query(&query);
        let s_self  = pipe.score(&tq, &pipe.encode(&query));
        let s_oth1  = pipe.score(&tq, &pipe.encode(&other1));
        let s_oth2  = pipe.score(&tq, &pipe.encode(&other2));

        assert!(s_self >= s_oth1, "self={s_self}, other1={s_oth1}");
        assert!(s_self >= s_oth2, "self={s_self}, other2={s_oth2}");
    }

    #[test]
    fn mixed_ratios_produce_correct_n4() {
        let c = Config::new(128, Metric::Cosine, 4).with_mixed_ratio(0.75);
        assert_eq!(c.n4_dims, 96);

        let c = Config::new(128, Metric::Cosine, 4).with_mixed_ratio(0.0);
        assert_eq!(c.n4_dims, 0);

        let c = Config::new(128, Metric::Cosine, 4).with_mixed_ratio(1.0);
        assert_eq!(c.n4_dims, 128);
    }

    #[test]
    fn profile_variance_approx_uniform() {
        let dim    = 128;
        let config = Config::new(dim, Metric::Cosine, 4);
        let sample: Vec<Vec<f32>> = (0..200).map(|s| random_unit_vec(dim, s)).collect();

        let variances = QuantPipeline::profile_variance(&config, &sample).unwrap();
        assert_eq!(variances.len(), dim); // padded_dim == dim when dim is power-of-2

        for (i, &v) in variances.iter().enumerate() {
            assert!(
                (v - 1.0).abs() < 0.3,
                "dim {i}: variance {v:.3} deviates too far from 1.0"
            );
        }
    }

    #[test]
    fn optimal_n4_dims_budget() {
        let padded = 128usize;
        let fake_variances = vec![1.0f32; padded];

        // avg 3.0 bits → n4 = 128 * 0.5 = 64
        assert_eq!(QuantPipeline::optimal_n4_dims(&fake_variances, 3.0), 64);
        // avg 3.5 bits → n4 = 128 * 0.75 = 96
        assert_eq!(QuantPipeline::optimal_n4_dims(&fake_variances, 3.5), 96);
        // avg 2.0 bits → pure 2-bit → 0
        assert_eq!(QuantPipeline::optimal_n4_dims(&fake_variances, 2.0), 0);
        // avg 4.0 bits → pure 4-bit → 128
        assert_eq!(QuantPipeline::optimal_n4_dims(&fake_variances, 4.0), 128);
    }

    #[test]
    fn variance_perm_sorted_desc() {
        let variances = vec![0.5f32, 2.0, 1.0, 3.0];
        let perm = QuantPipeline::variance_perm(&variances);
        // Expected order: dim 3 (3.0), dim 1 (2.0), dim 2 (1.0), dim 0 (0.5)
        assert_eq!(perm, vec![3u16, 1, 2, 0]);
    }

    #[test]
    fn with_n4_dims_builder() {
        let c = Config::new(128, Metric::Cosine, 4).with_n4_dims(96);
        assert_eq!(c.n4_dims, 96);
        // rounds down to mult-of-4
        let c2 = Config::new(128, Metric::Cosine, 4).with_n4_dims(97);
        assert_eq!(c2.n4_dims, 96);
        // clamps to padded_dim
        let c3 = Config::new(128, Metric::Cosine, 4).with_n4_dims(9999);
        assert_eq!(c3.n4_dims, 128);
    }
}
