//! Pooling strategies for aggregating token embeddings.
//!
//! These utilities are available for custom embedder implementations.

#![allow(dead_code)]

use serde::{Deserialize, Serialize};

/// Strategy for pooling token embeddings into a single vector.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum PoolingStrategy {
    /// Use the [CLS] token embedding.
    Cls,
    /// Mean of all token embeddings (weighted by attention mask).
    Mean,
    /// Max pooling across all tokens.
    Max,
    /// Use the last token embedding (required for Jina Code 1.5B).
    /// This is the correct pooling for decoder-based models like Qwen2.5-Coder.
    LastToken,
}

impl Default for PoolingStrategy {
    fn default() -> Self {
        Self::Mean
    }
}

impl PoolingStrategy {
    /// Get a description of the pooling strategy.
    pub fn description(&self) -> &'static str {
        match self {
            Self::Cls => "Uses the [CLS] token embedding (first token)",
            Self::Mean => "Averages all token embeddings, weighted by attention mask",
            Self::Max => "Takes the maximum value across all tokens for each dimension",
            Self::LastToken => "Uses the last valid token embedding (for decoder models like Jina Code 1.5B)",
        }
    }

    /// Check if this strategy is recommended for sentence similarity.
    pub fn recommended_for_similarity(&self) -> bool {
        matches!(self, Self::Mean | Self::LastToken)
    }
}

/// Pool vectors using CPU operations.
pub fn pool_vectors(embeddings: &[Vec<f32>], strategy: PoolingStrategy) -> Vec<f32> {
    if embeddings.is_empty() {
        return Vec::new();
    }

    let dim = embeddings[0].len();

    match strategy {
        PoolingStrategy::Cls => {
            // First vector
            embeddings[0].clone()
        }
        PoolingStrategy::Mean => {
            // Average of all vectors
            let mut result = vec![0.0f32; dim];
            for emb in embeddings {
                for (i, &v) in emb.iter().enumerate() {
                    result[i] += v;
                }
            }
            let n = embeddings.len() as f32;
            for v in &mut result {
                *v /= n;
            }
            result
        }
        PoolingStrategy::Max => {
            // Maximum across all vectors
            let mut result = vec![f32::NEG_INFINITY; dim];
            for emb in embeddings {
                for (i, &v) in emb.iter().enumerate() {
                    if v > result[i] {
                        result[i] = v;
                    }
                }
            }
            result
        }
        PoolingStrategy::LastToken => {
            // Last token embedding (for decoder-based models)
            embeddings.last().cloned().unwrap_or_default()
        }
    }
}

/// Normalize a vector to unit length.
pub fn normalize_vector(v: &mut [f32]) {
    let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > 0.0 {
        for x in v.iter_mut() {
            *x /= norm;
        }
    }
}

/// Compute cosine similarity between two vectors.
pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    debug_assert_eq!(a.len(), b.len());

    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();

    if norm_a > 0.0 && norm_b > 0.0 {
        dot / (norm_a * norm_b)
    } else {
        0.0
    }
}

/// Compute dot product between two vectors.
/// Assumes vectors are normalized.
pub fn dot_product(a: &[f32], b: &[f32]) -> f32 {
    debug_assert_eq!(a.len(), b.len());
    a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
}

/// Compute euclidean distance between two vectors.
pub fn euclidean_distance(a: &[f32], b: &[f32]) -> f32 {
    debug_assert_eq!(a.len(), b.len());
    a.iter()
        .zip(b.iter())
        .map(|(x, y)| (x - y).powi(2))
        .sum::<f32>()
        .sqrt()
}

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

    #[test]
    fn test_mean_pooling() {
        let embeddings = vec![
            vec![1.0, 2.0, 3.0],
            vec![4.0, 5.0, 6.0],
            vec![7.0, 8.0, 9.0],
        ];

        let result = pool_vectors(&embeddings, PoolingStrategy::Mean);
        assert_eq!(result, vec![4.0, 5.0, 6.0]);
    }

    #[test]
    fn test_max_pooling() {
        let embeddings = vec![
            vec![1.0, 5.0, 3.0],
            vec![4.0, 2.0, 6.0],
            vec![7.0, 8.0, 1.0],
        ];

        let result = pool_vectors(&embeddings, PoolingStrategy::Max);
        assert_eq!(result, vec![7.0, 8.0, 6.0]);
    }

    #[test]
    fn test_cosine_similarity() {
        let a = vec![1.0, 0.0, 0.0];
        let b = vec![1.0, 0.0, 0.0];
        assert!((cosine_similarity(&a, &b) - 1.0).abs() < 0.001);

        let c = vec![0.0, 1.0, 0.0];
        assert!(cosine_similarity(&a, &c).abs() < 0.001);
    }

    #[test]
    fn test_normalize() {
        let mut v = vec![3.0, 4.0];
        normalize_vector(&mut v);

        let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!((norm - 1.0).abs() < 0.001);
    }
}
