aprender-core 0.29.2

//! `HuggingFace` weight comparison module (GH-121).
//!
// Allow format_push_string: cleaner code for string building without I/O concerns
#![allow(clippy::format_push_string, clippy::uninlined_format_args)]
//!
//! Compares model weights between APR and SafeTensors/HuggingFace sources:
//! - Tensor value comparison (L2 norm, max diff)
//! - Tensor shape comparison
//! - Tensor name mapping
//! - Statistical summary
//!
//! # Example
//!
//! ```rust
//! use aprender::format::compare::{TensorComparison, WeightComparer, CompareConfig};
//!
//! let comparison = TensorComparison {
//!     name: "encoder.weight".to_string(),
//!     shape_match: true,
//!     source_shape: vec![512, 768],
//!     target_shape: vec![512, 768],
//!     l2_diff: Some(1.23e-7),
//!     max_diff: Some(4.56e-8),
//!     mean_diff: Some(1.0e-8),
//! };
//! assert!(comparison.is_match());
//! ```
//!
//! # PMAT Compliance
//!
//! - Zero `unwrap()` calls
//! - All public APIs return `Result<T, E>`

use std::collections::{HashMap, HashSet};

// ============================================================================
// Configuration
// ============================================================================

/// Configuration for weight comparison
#[derive(Debug, Clone)]
pub struct CompareConfig {
    /// Tolerance for L2 norm difference (default: 1e-5)
    pub l2_tolerance: f64,
    /// Tolerance for max element-wise difference (default: 1e-5)
    pub max_tolerance: f64,
    /// Whether to allow shape broadcasting
    pub allow_broadcast: bool,
    /// Whether to normalize tensors before comparison
    pub normalize_first: bool,
    /// Prefix to strip from source tensor names
    pub source_prefix: Option<String>,
    /// Prefix to strip from target tensor names
    pub target_prefix: Option<String>,
}

impl Default for CompareConfig {
    fn default() -> Self {
        Self {
            l2_tolerance: 1e-5,
            max_tolerance: 1e-5,
            allow_broadcast: false,
            normalize_first: false,
            source_prefix: None,
            target_prefix: None,
        }
    }
}

impl CompareConfig {
    /// Strict comparison (no tolerance for differences)
    #[must_use]
    pub fn strict() -> Self {
        Self {
            l2_tolerance: 0.0,
            max_tolerance: 0.0,
            ..Self::default()
        }
    }

    /// Relaxed comparison for quantized models
    #[must_use]
    pub fn quantized() -> Self {
        Self {
            l2_tolerance: 1e-2,
            max_tolerance: 1e-2,
            ..Self::default()
        }
    }
}

// ============================================================================
// Comparison Results
// ============================================================================

/// Result of comparing a single tensor
#[derive(Debug, Clone)]
pub struct TensorComparison {
    /// Tensor name
    pub name: String,
    /// Whether shapes match
    pub shape_match: bool,
    /// Source tensor shape
    pub source_shape: Vec<usize>,
    /// Target tensor shape
    pub target_shape: Vec<usize>,
    /// L2 norm of difference (if shapes match)
    pub l2_diff: Option<f64>,
    /// Maximum element-wise difference (if shapes match)
    pub max_diff: Option<f64>,
    /// Mean element-wise difference (if shapes match)
    pub mean_diff: Option<f64>,
}

impl TensorComparison {
    /// Check if tensors match within tolerance
    #[must_use]
    pub fn is_match_with_tolerance(&self, l2_tol: f64, max_tol: f64) -> bool {
        self.shape_match
            && self.l2_diff.is_some_and(|d| d <= l2_tol)
            && self.max_diff.is_some_and(|d| d <= max_tol)
    }

    /// Check if tensors match with default tolerance
    #[must_use]
    pub fn is_match(&self) -> bool {
        self.is_match_with_tolerance(1e-5, 1e-5)
    }

    /// Get element count from shape
    #[must_use]
    pub fn element_count(&self) -> usize {
        self.source_shape.iter().product()
    }
}

/// Summary of full model comparison
#[derive(Debug, Clone)]
pub struct CompareReport {
    /// Individual tensor comparisons
    pub tensors: Vec<TensorComparison>,
    /// Tensors only in source
    pub source_only: Vec<String>,
    /// Tensors only in target
    pub target_only: Vec<String>,
    /// Total L2 norm across all tensors
    pub total_l2_diff: f64,
    /// Maximum difference across all tensors
    pub global_max_diff: f64,
    /// Configuration used
    pub config: CompareConfig,
}

impl CompareReport {
    /// Check if all tensors match within tolerance
    #[must_use]
    pub fn all_match(&self) -> bool {
        self.source_only.is_empty()
            && self.target_only.is_empty()
            && self.tensors.iter().all(|t| {
                t.is_match_with_tolerance(self.config.l2_tolerance, self.config.max_tolerance)
            })
    }

    /// Count matching tensors
    #[must_use]
    pub fn match_count(&self) -> usize {
        self.tensors
            .iter()
            .filter(|t| {
                t.is_match_with_tolerance(self.config.l2_tolerance, self.config.max_tolerance)
            })
            .count()
    }

    /// Count mismatched tensors
    #[must_use]
    pub fn mismatch_count(&self) -> usize {
        self.tensors.len() - self.match_count()
    }

    /// Get summary string
    #[must_use]
    pub fn summary(&self) -> String {
        let mut result = String::new();
        result.push_str("Weight Comparison Report\n");
        result.push_str("========================\n\n");

        result.push_str(&format!("Tensors compared: {}\n", self.tensors.len()));
        result.push_str(&format!("  Matching: {}\n", self.match_count()));
        result.push_str(&format!("  Mismatched: {}\n", self.mismatch_count()));
        result.push_str(&format!("  Source only: {}\n", self.source_only.len()));
        result.push_str(&format!("  Target only: {}\n", self.target_only.len()));
        result.push('\n');
        result.push_str(&format!("Total L2 diff: {:.6e}\n", self.total_l2_diff));
        result.push_str(&format!("Global max diff: {:.6e}\n", self.global_max_diff));
        result.push_str(&format!("L2 tolerance: {:.6e}\n", self.config.l2_tolerance));
        result.push_str(&format!(
            "Max tolerance: {:.6e}\n",
            self.config.max_tolerance
        ));

        if !self.source_only.is_empty() {
            result.push_str("\nTensors only in source:\n");
            for name in &self.source_only {
                result.push_str(&format!("  - {name}\n"));
            }
        }

        if !self.target_only.is_empty() {
            result.push_str("\nTensors only in target:\n");
            for name in &self.target_only {
                result.push_str(&format!("  - {name}\n"));
            }
        }

        let mismatched: Vec<_> = self
            .tensors
            .iter()
            .filter(|t| {
                !t.is_match_with_tolerance(self.config.l2_tolerance, self.config.max_tolerance)
            })
            .collect();

        if !mismatched.is_empty() {
            result.push_str("\nMismatched tensors:\n");
            for t in mismatched {
                let shape_info = if t.shape_match {
                    format!("{:?}", t.source_shape)
                } else {
                    format!("{:?} vs {:?}", t.source_shape, t.target_shape)
                };
                let diff_info = t
                    .l2_diff
                    .map_or_else(|| "shape mismatch".to_string(), |d| format!("L2={:.6e}", d));
                result.push_str(&format!("  - {} [{}]: {}\n", t.name, shape_info, diff_info));
            }
        }

        result
    }
}

// ============================================================================
// Weight Comparer
// ============================================================================

/// Compares weights between two tensor collections.
///
/// # Example
///
/// ```rust
/// use aprender::format::compare::{WeightComparer, CompareConfig};
///
/// let comparer = WeightComparer::new(CompareConfig::default());
///
/// // Compare two f32 tensors
/// let source = vec![1.0_f32, 2.0, 3.0];
/// let target = vec![1.0_f32, 2.0, 3.0];
/// let comparison = comparer.compare_tensors(
///     "test",
///     &source,
///     &[3],
///     &target,
///     &[3],
/// );
/// assert!(comparison.is_match());
/// ```
#[derive(Debug, Clone)]
pub struct WeightComparer {
    config: CompareConfig,
}

impl WeightComparer {
    /// Create new weight comparer
    #[must_use]
    pub fn new(config: CompareConfig) -> Self {
        Self { config }
    }

    /// Compare two tensors
    #[must_use]
    pub fn compare_tensors(
        &self,
        name: &str,
        source_data: &[f32],
        source_shape: &[usize],
        target_data: &[f32],
        target_shape: &[usize],
    ) -> TensorComparison {
        let shape_match = source_shape == target_shape;

        let (l2_diff, max_diff, mean_diff) =
            if shape_match && source_data.len() == target_data.len() {
                let diff_stats = self.compute_diff_stats(source_data, target_data);
                (Some(diff_stats.0), Some(diff_stats.1), Some(diff_stats.2))
            } else {
                (None, None, None)
            };

        TensorComparison {
            name: name.to_string(),
            shape_match,
            source_shape: source_shape.to_vec(),
            target_shape: target_shape.to_vec(),
            l2_diff,
            max_diff,
            mean_diff,
        }
    }

    /// Compute L2 norm, max diff, and mean diff
    #[allow(clippy::unused_self)]
    fn compute_diff_stats(&self, source: &[f32], target: &[f32]) -> (f64, f64, f64) {
        if source.is_empty() {
            return (0.0, 0.0, 0.0);
        }

        let mut sum_sq = 0.0_f64;
        let mut max_diff = 0.0_f64;
        let mut sum_diff = 0.0_f64;
        let mut count = 0_usize;

        for (s, t) in source.iter().zip(target.iter()) {
            if s.is_finite() && t.is_finite() {
                let diff = (f64::from(*s) - f64::from(*t)).abs();
                sum_sq += diff * diff;
                if diff > max_diff {
                    max_diff = diff;
                }
                sum_diff += diff;
                count += 1;
            }
        }

        let l2 = sum_sq.sqrt();
        let mean = if count > 0 {
            sum_diff / count as f64
        } else {
            0.0
        };

        (l2, max_diff, mean)
    }

    /// Compare two models represented as tensor maps
    #[must_use]
    pub fn compare_models(
        &self,
        source: &HashMap<String, (Vec<f32>, Vec<usize>)>,
        target: &HashMap<String, (Vec<f32>, Vec<usize>)>,
    ) -> CompareReport {
        let source_names: HashSet<&str> = source.keys().map(String::as_str).collect();
        let target_names: HashSet<&str> = target.keys().map(String::as_str).collect();

        // Normalize names if prefixes configured
        let normalize_name = |name: &str, prefix: &Option<String>| -> String {
            prefix
                .as_ref()
                .and_then(|p| name.strip_prefix(p.as_str()))
                .unwrap_or(name)
                .to_string()
        };

        let source_only: Vec<String> = source_names
            .difference(&target_names)
            .map(|s| (*s).to_string())
            .collect();

        let target_only: Vec<String> = target_names
            .difference(&source_names)
            .map(|s| (*s).to_string())
            .collect();

        let common: Vec<&str> = source_names.intersection(&target_names).copied().collect();

        let mut tensors = Vec::new();
        let mut total_l2_sq = 0.0_f64;
        let mut global_max = 0.0_f64;

        for name in common {
            let (source_data, source_shape) = source.get(name).expect("checked intersection");
            let (target_data, target_shape) = target.get(name).expect("checked intersection");

            let normalized_name = normalize_name(name, &self.config.source_prefix);

            let comparison = self.compare_tensors(
                &normalized_name,
                source_data,
                source_shape,
                target_data,
                target_shape,
            );

            if let Some(l2) = comparison.l2_diff {
                total_l2_sq += l2 * l2;
            }
            if let Some(max) = comparison.max_diff {
                if max > global_max {
                    global_max = max;
                }
            }

            tensors.push(comparison);
        }

        CompareReport {
            tensors,
            source_only,
            target_only,
            total_l2_diff: total_l2_sq.sqrt(),
            global_max_diff: global_max,
            config: self.config.clone(),
        }
    }
}

// ============================================================================
// Utility Functions
// ============================================================================

/// Compare two f32 slices and return L2 norm of difference.
///
/// Returns None if lengths differ or data contains NaN/Inf.
#[must_use]
pub fn l2_diff(a: &[f32], b: &[f32]) -> Option<f64> {
    if a.len() != b.len() {
        return None;
    }

    let mut sum_sq = 0.0_f64;
    for (x, y) in a.iter().zip(b.iter()) {
        if !x.is_finite() || !y.is_finite() {
            return None;
        }
        let diff = f64::from(*x) - f64::from(*y);
        sum_sq += diff * diff;
    }

    Some(sum_sq.sqrt())
}

/// Compare two f32 slices and return max absolute difference.
#[must_use]
pub fn max_diff(a: &[f32], b: &[f32]) -> Option<f64> {
    if a.len() != b.len() {
        return None;
    }

    let mut max = 0.0_f64;
    for (x, y) in a.iter().zip(b.iter()) {
        if !x.is_finite() || !y.is_finite() {
            return None;
        }
        let diff = (f64::from(*x) - f64::from(*y)).abs();
        if diff > max {
            max = diff;
        }
    }

    Some(max)
}

/// Compute relative L2 error: ||a - b|| / ||a||
#[must_use]
pub fn relative_l2_error(a: &[f32], b: &[f32]) -> Option<f64> {
    let diff_norm = l2_diff(a, b)?;
    let a_norm: f64 = a.iter().map(|x| f64::from(*x).powi(2)).sum::<f64>().sqrt();

    if a_norm < f64::EPSILON {
        if diff_norm < f64::EPSILON {
            Some(0.0)
        } else {
            Some(f64::INFINITY)
        }
    } else {
        Some(diff_norm / a_norm)
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
#[path = "compare_tests.rs"]
mod tests;