trustformers-core 0.2.1

//! Cross-framework numerical validation for TrustformeRS.
//!
//! Only frameworks that are actually reachable from this crate can be
//! validated against:
//!
//! * **ONNX** — real: the in-tree pure-Rust ONNX CPU executor
//!   ([`crate::export::onnx_runtime`]) runs a reference `.onnx` graph and its
//!   outputs are compared element-wise against the TrustformeRS outputs the
//!   caller recorded on each test case.
//! * **PyTorch / TensorFlow / JAX** — no bindings are compiled into this crate,
//!   so these report themselves unavailable. They never report `passed`, and
//!   they never synthesise `max_diff` / `mean_diff` statistics.

use crate::{
    errors::{Result, TrustformersError},
    tensor::Tensor,
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;

/// Supported ML frameworks for cross-validation
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum Framework {
    /// PyTorch framework
    PyTorch,
    /// TensorFlow framework
    TensorFlow,
    /// JAX/Flax framework
    Jax,
    /// ONNX Runtime
    OnnxRuntime,
    /// TrustformeRS (our implementation)
    TrustformeRS,
}

impl Framework {
    /// Get the string representation of the framework
    pub fn as_str(&self) -> &'static str {
        match self {
            Framework::PyTorch => "pytorch",
            Framework::TensorFlow => "tensorflow",
            Framework::Jax => "jax",
            Framework::OnnxRuntime => "onnx",
            Framework::TrustformeRS => "trustformers",
        }
    }
}

/// Configuration for cross-framework validation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationConfig {
    /// Absolute tolerance for numerical comparison
    pub atol: f64,
    /// Relative tolerance for numerical comparison
    pub rtol: f64,
    /// Maximum number of mismatched elements to report
    pub max_errors: usize,
    /// Whether to validate gradients (if available)
    pub validate_gradients: bool,
    /// Frameworks to validate against
    pub target_frameworks: Vec<Framework>,
    /// Model architecture to test
    pub model_architecture: String,
    /// Model parameters/weights to use
    #[serde(skip)]
    pub model_params: Option<HashMap<String, Tensor>>,
    /// Path to the reference `.onnx` model executed by the ONNX validation
    /// path. Without it, ONNX validation reports that it has nothing to
    /// compare against.
    #[serde(default)]
    pub onnx_reference_model: Option<PathBuf>,
}

impl Default for ValidationConfig {
    fn default() -> Self {
        Self {
            atol: 1e-5,
            rtol: 1e-4,
            max_errors: 10,
            validate_gradients: false,
            target_frameworks: vec![Framework::PyTorch, Framework::TensorFlow],
            model_architecture: "transformer".to_string(),
            model_params: None,
            onnx_reference_model: None,
        }
    }
}

impl ValidationConfig {
    /// Point the ONNX validation path at a reference model on disk.
    pub fn with_onnx_reference_model(mut self, path: impl Into<PathBuf>) -> Self {
        self.onnx_reference_model = Some(path.into());
        self
    }
}

/// Results of cross-framework validation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationResult {
    /// Framework this result is for
    pub framework: Framework,
    /// Whether validation passed
    pub passed: bool,
    /// Maximum absolute difference found
    pub max_diff: f64,
    /// Mean absolute difference
    pub mean_diff: f64,
    /// Number of mismatched elements
    pub mismatch_count: usize,
    /// Total number of elements compared
    pub total_elements: usize,
    /// Execution time in milliseconds
    pub execution_time_ms: f64,
    /// Additional metrics
    pub metrics: HashMap<String, f64>,
    /// Error messages (if any)
    pub errors: Vec<String>,
}

impl ValidationResult {
    /// Create a new validation result
    pub fn new(framework: Framework) -> Self {
        Self {
            framework,
            passed: false,
            max_diff: 0.0,
            mean_diff: 0.0,
            mismatch_count: 0,
            total_elements: 0,
            execution_time_ms: 0.0,
            metrics: HashMap::new(),
            errors: Vec::new(),
        }
    }

    /// Calculate pass rate as percentage
    pub fn pass_rate(&self) -> f64 {
        if self.total_elements == 0 {
            0.0
        } else {
            100.0 * (self.total_elements - self.mismatch_count) as f64 / self.total_elements as f64
        }
    }

    /// Add a custom metric
    pub fn add_metric(&mut self, name: String, value: f64) {
        self.metrics.insert(name, value);
    }

    /// Add an error message
    pub fn add_error(&mut self, error: String) {
        self.errors.push(error);
    }
}

/// Test case for cross-framework validation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationTestCase {
    /// Name of the test case
    pub name: String,
    /// Input tensors for the test
    #[serde(skip)]
    pub inputs: Vec<Tensor>,
    /// Expected output shape
    pub expected_shape: Vec<usize>,
    /// Model configuration
    pub model_config: HashMap<String, serde_json::Value>,
    /// Test-specific configuration overrides
    #[serde(skip)]
    pub config_overrides: Option<ValidationConfig>,
    /// Graph input names matching `inputs`, in order.
    ///
    /// When empty, the reference graph's own input order is used.
    #[serde(default)]
    pub input_names: Vec<String>,
    /// The TrustformeRS outputs to compare the reference framework against,
    /// keyed by graph output name.
    ///
    /// A test case with no expected outputs cannot be validated; the validator
    /// records an error for it instead of reporting a pass.
    #[serde(skip)]
    pub expected_outputs: HashMap<String, Tensor>,
}

impl ValidationTestCase {
    /// Create a new validation test case
    pub fn new(name: String, inputs: Vec<Tensor>) -> Self {
        Self {
            name,
            inputs,
            expected_shape: Vec::new(),
            model_config: HashMap::new(),
            config_overrides: None,
            input_names: Vec::new(),
            expected_outputs: HashMap::new(),
        }
    }

    /// Name the graph inputs corresponding to `inputs`.
    pub fn with_input_names(mut self, names: Vec<String>) -> Self {
        self.input_names = names;
        self
    }

    /// Record the TrustformeRS output the reference framework must reproduce.
    pub fn with_expected_output(mut self, name: impl Into<String>, tensor: Tensor) -> Self {
        self.expected_outputs.insert(name.into(), tensor);
        self
    }

    /// Set expected output shape
    pub fn with_expected_shape(mut self, shape: Vec<usize>) -> Self {
        self.expected_shape = shape;
        self
    }

    /// Add model configuration
    pub fn with_model_config(mut self, key: String, value: serde_json::Value) -> Self {
        self.model_config.insert(key, value);
        self
    }

    /// Set configuration overrides
    pub fn with_config_overrides(mut self, config: ValidationConfig) -> Self {
        self.config_overrides = Some(config);
        self
    }
}

/// Cross-framework validation suite
#[derive(Debug)]
pub struct CrossFrameworkValidator {
    /// Base configuration
    config: ValidationConfig,
    /// Available frameworks
    available_frameworks: Vec<Framework>,
    /// Test cases
    test_cases: Vec<ValidationTestCase>,
}

impl CrossFrameworkValidator {
    /// Create a new cross-framework validator
    pub fn new(config: ValidationConfig) -> Self {
        Self {
            config,
            available_frameworks: vec![Framework::TrustformeRS], // Always available
            test_cases: Vec::new(),
        }
    }

    /// Create validator with default configuration
    pub fn with_defaults() -> Self {
        Self::new(ValidationConfig::default())
    }

    /// Detect available frameworks
    pub fn detect_frameworks(&mut self) -> Result<()> {
        self.available_frameworks.clear();
        self.available_frameworks.push(Framework::TrustformeRS);

        // Check for PyTorch
        if Self::check_pytorch_available() {
            self.available_frameworks.push(Framework::PyTorch);
        }

        // Check for TensorFlow
        if Self::check_tensorflow_available() {
            self.available_frameworks.push(Framework::TensorFlow);
        }

        // Check for JAX
        if Self::check_jax_available() {
            self.available_frameworks.push(Framework::Jax);
        }

        // Check for ONNX Runtime
        if Self::check_onnx_available() {
            self.available_frameworks.push(Framework::OnnxRuntime);
        }

        Ok(())
    }

    /// Check if PyTorch is available
    fn check_pytorch_available() -> bool {
        // No PyTorch backend is compiled into trustformers-core (the `tch` dependency was
        // removed); PyTorch cross-framework validation is therefore never available.
        false
    }

    /// Check if TensorFlow is available.
    ///
    /// No TensorFlow bindings are compiled into `trustformers-core`, so this is
    /// always false. It deliberately does *not* consult an environment
    /// variable: an env var says nothing about whether this process can execute
    /// a TensorFlow graph.
    fn check_tensorflow_available() -> bool {
        false
    }

    /// Check if JAX is available. See [`Self::check_tensorflow_available`].
    fn check_jax_available() -> bool {
        false
    }

    /// Check if an ONNX executor is available.
    ///
    /// Always true: `trustformers-core` ships a pure-Rust ONNX protobuf reader
    /// and CPU graph executor ([`crate::export::onnx_runtime`]).
    fn check_onnx_available() -> bool {
        true
    }

    /// Add a test case
    pub fn add_test_case(&mut self, test_case: ValidationTestCase) {
        self.test_cases.push(test_case);
    }

    /// Run validation against all available frameworks
    pub fn validate_all(&self) -> Result<HashMap<Framework, ValidationResult>> {
        let mut results = HashMap::new();

        for &framework in &self.available_frameworks {
            if framework == Framework::TrustformeRS {
                continue; // Skip self-validation
            }

            if self.config.target_frameworks.contains(&framework) {
                let result = self.validate_framework(framework)?;
                results.insert(framework, result);
            }
        }

        Ok(results)
    }

    /// Validate against a specific framework
    pub fn validate_framework(&self, framework: Framework) -> Result<ValidationResult> {
        let mut result = ValidationResult::new(framework);
        let start_time = std::time::Instant::now();

        match framework {
            Framework::PyTorch => self.validate_pytorch(&mut result)?,
            Framework::TensorFlow => self.validate_tensorflow(&mut result)?,
            Framework::Jax => self.validate_jax(&mut result)?,
            Framework::OnnxRuntime => self.validate_onnx(&mut result)?,
            Framework::TrustformeRS => {
                return Err(TrustformersError::invalid_input(
                    "Cannot validate against self".to_string(),
                ))
            },
        }

        result.execution_time_ms = start_time.elapsed().as_secs_f64() * 1000.0;
        Ok(result)
    }

    /// Validate against PyTorch
    fn validate_pytorch(&self, result: &mut ValidationResult) -> Result<()> {
        // No PyTorch backend is compiled into trustformers-core (the `tch` dependency was
        // removed), so PyTorch cross-framework validation always reports unavailable.
        result.add_error("PyTorch not available".to_string());

        Ok(())
    }

    /// Validate against TensorFlow.
    ///
    /// No TensorFlow backend is compiled into `trustformers-core`, so this
    /// reports unavailable. It never sets `passed` and never synthesises diff
    /// statistics.
    fn validate_tensorflow(&self, result: &mut ValidationResult) -> Result<()> {
        result.add_error(
            "TensorFlow validation is not available: no TensorFlow bindings are compiled into \
             trustformers-core, so no tensor can be compared"
                .to_string(),
        );
        Ok(())
    }

    /// Validate against JAX. See [`Self::validate_tensorflow`].
    fn validate_jax(&self, result: &mut ValidationResult) -> Result<()> {
        result.add_error(
            "JAX validation is not available: no JAX bindings are compiled into \
             trustformers-core, so no tensor can be compared"
                .to_string(),
        );
        Ok(())
    }

    /// Validate against a reference ONNX graph using the in-tree CPU executor.
    ///
    /// Every registered test case is run through the reference graph and its
    /// outputs are compared element-wise against the TrustformeRS outputs
    /// recorded on the case. `passed` is only ever set from a real comparison.
    fn validate_onnx(&self, result: &mut ValidationResult) -> Result<()> {
        use crate::export::onnx_runtime::ONNXRuntimeBackend;

        let Some(model_path) = self.config.onnx_reference_model.as_ref() else {
            result.add_error(
                "ONNX validation needs a reference model: set \
                 ValidationConfig::onnx_reference_model to a .onnx file"
                    .to_string(),
            );
            return Ok(());
        };

        let backend = ONNXRuntimeBackend::new();
        let session = match backend.load_model(model_path) {
            Ok(session) => session,
            Err(error) => {
                result.add_error(format!(
                    "failed to load reference ONNX model {}: {}",
                    model_path.display(),
                    error
                ));
                return Ok(());
            },
        };

        let unsupported = session.unsupported_operators();
        if !unsupported.is_empty() {
            result.add_error(format!(
                "reference ONNX model uses operators the CPU executor does not implement: {}",
                unsupported.join(", ")
            ));
            return Ok(());
        }

        if self.test_cases.is_empty() {
            result.add_error(
                "ONNX validation has no test cases to run; add one with add_test_case".to_string(),
            );
            return Ok(());
        }

        let mut max_diff: f64 = 0.0;
        let mut weighted_diff_sum: f64 = 0.0;
        let mut mismatch_count = 0usize;
        let mut total_elements = 0usize;
        let mut compared_tensors = 0usize;

        for test_case in &self.test_cases {
            if test_case.expected_outputs.is_empty() {
                result.add_error(format!(
                    "test case '{}' records no TrustformeRS outputs, so there is nothing to \
                     compare the ONNX result against",
                    test_case.name
                ));
                continue;
            }

            let input_names: Vec<String> = if test_case.input_names.is_empty() {
                session.input_names().to_vec()
            } else {
                test_case.input_names.clone()
            };

            if input_names.len() != test_case.inputs.len() {
                result.add_error(format!(
                    "test case '{}' provides {} input tensors but the graph expects {}",
                    test_case.name,
                    test_case.inputs.len(),
                    input_names.len()
                ));
                continue;
            }

            let inputs: HashMap<String, Tensor> =
                input_names.into_iter().zip(test_case.inputs.iter().cloned()).collect();

            let outputs = match session.run(inputs) {
                Ok(outputs) => outputs,
                Err(error) => {
                    result.add_error(format!(
                        "test case '{}': reference ONNX execution failed: {}",
                        test_case.name, error
                    ));
                    continue;
                },
            };

            for (output_name, expected) in &test_case.expected_outputs {
                let Some(actual) = outputs.get(output_name) else {
                    result.add_error(format!(
                        "test case '{}': the reference graph produced no output named '{}'",
                        test_case.name, output_name
                    ));
                    continue;
                };

                if actual.shape() != expected.shape() {
                    result.add_error(format!(
                        "test case '{}' output '{}': shape mismatch {:?} vs {:?}",
                        test_case.name,
                        output_name,
                        actual.shape(),
                        expected.shape()
                    ));
                    continue;
                }

                let comparison = self.compare_tensor_values(actual, expected)?;
                max_diff = max_diff.max(comparison.max_diff);
                weighted_diff_sum += comparison.mean_diff * comparison.total_elements as f64;
                mismatch_count += comparison.mismatch_count;
                total_elements += comparison.total_elements;
                compared_tensors += 1;
            }
        }

        result.max_diff = max_diff;
        result.mean_diff =
            if total_elements > 0 { weighted_diff_sum / total_elements as f64 } else { 0.0 };
        result.mismatch_count = mismatch_count;
        result.total_elements = total_elements;
        result.add_metric("compared_tensors".to_string(), compared_tensors as f64);
        result.passed = compared_tensors > 0 && mismatch_count == 0 && result.errors.is_empty();

        Ok(())
    }

    /// Compare two tensors with the given tolerances
    pub fn compare_tensors(&self, tensor1: &Tensor, tensor2: &Tensor) -> Result<ValidationResult> {
        let mut result = ValidationResult::new(Framework::TrustformeRS);

        // Check shapes match
        if tensor1.shape() != tensor2.shape() {
            result.add_error(format!(
                "Shape mismatch: {:?} vs {:?}",
                tensor1.shape(),
                tensor2.shape()
            ));
            return Ok(result);
        }

        // Compare data types
        if tensor1.dtype() != tensor2.dtype() {
            result.add_error(format!(
                "Data type mismatch: {:?} vs {:?}",
                tensor1.dtype(),
                tensor2.dtype()
            ));
            return Ok(result);
        }

        // Compare values
        let comparison = self.compare_tensor_values(tensor1, tensor2)?;
        result.max_diff = comparison.max_diff;
        result.mean_diff = comparison.mean_diff;
        result.mismatch_count = comparison.mismatch_count;
        result.total_elements = comparison.total_elements;
        result.passed = comparison.mismatch_count == 0;

        Ok(result)
    }

    /// Compare tensor values
    fn compare_tensor_values(
        &self,
        tensor1: &Tensor,
        tensor2: &Tensor,
    ) -> Result<TensorComparison> {
        match (tensor1, tensor2) {
            (Tensor::F32(a1), Tensor::F32(a2)) => {
                let s1 =
                    a1.as_slice().ok_or_else(|| anyhow::anyhow!("F32 tensor is not contiguous"))?;
                let s2 =
                    a2.as_slice().ok_or_else(|| anyhow::anyhow!("F32 tensor is not contiguous"))?;
                self.compare_f32_arrays(s1, s2)
            },
            (Tensor::F64(a1), Tensor::F64(a2)) => {
                let s1 =
                    a1.as_slice().ok_or_else(|| anyhow::anyhow!("F64 tensor is not contiguous"))?;
                let s2 =
                    a2.as_slice().ok_or_else(|| anyhow::anyhow!("F64 tensor is not contiguous"))?;
                self.compare_f64_arrays(s1, s2)
            },
            _ => {
                // For other types, convert to f32 and compare
                let data1 = tensor1.to_vec_f32()?;
                let data2 = tensor2.to_vec_f32()?;
                self.compare_f32_arrays(&data1, &data2)
            },
        }
    }

    /// Compare f32 arrays
    fn compare_f32_arrays(&self, arr1: &[f32], arr2: &[f32]) -> Result<TensorComparison> {
        let mut max_diff: f64 = 0.0;
        let mut sum_diff: f64 = 0.0;
        let mut mismatch_count = 0;
        let total_elements = arr1.len();

        for (&v1, &v2) in arr1.iter().zip(arr2.iter()) {
            let diff = (v1 - v2).abs();
            let rel_diff = if v2.abs() > 0.0 { diff / v2.abs() } else { diff };

            if diff > self.config.atol as f32 && rel_diff > self.config.rtol as f32 {
                mismatch_count += 1;
                if mismatch_count <= self.config.max_errors {
                    // Log the mismatch (in a real implementation)
                }
            }

            max_diff = max_diff.max(diff as f64);
            sum_diff += diff as f64;
        }

        Ok(TensorComparison {
            max_diff,
            mean_diff: sum_diff / total_elements as f64,
            mismatch_count,
            total_elements,
        })
    }

    /// Compare f64 arrays
    fn compare_f64_arrays(&self, arr1: &[f64], arr2: &[f64]) -> Result<TensorComparison> {
        let mut max_diff: f64 = 0.0;
        let mut sum_diff: f64 = 0.0;
        let mut mismatch_count = 0;
        let total_elements = arr1.len();

        for (&v1, &v2) in arr1.iter().zip(arr2.iter()) {
            let diff = (v1 - v2).abs();
            let rel_diff = if v2.abs() > 0.0 { diff / v2.abs() } else { diff };

            if diff > self.config.atol && rel_diff > self.config.rtol {
                mismatch_count += 1;
            }

            max_diff = max_diff.max(diff);
            sum_diff += diff;
        }

        Ok(TensorComparison {
            max_diff,
            mean_diff: sum_diff / total_elements as f64,
            mismatch_count,
            total_elements,
        })
    }

    /// Generate a comprehensive validation report
    pub fn generate_report(&self, results: &HashMap<Framework, ValidationResult>) -> String {
        let mut report = String::new();
        report.push_str("# Cross-Framework Validation Report\n\n");

        // Summary
        let total_frameworks = results.len();
        let passed_frameworks = results.values().filter(|r| r.passed).count();
        report.push_str("## Summary\n\n");
        report.push_str(&format!(
            "- **Total Frameworks Tested**: {}\n",
            total_frameworks
        ));
        report.push_str(&format!("- **Passed**: {}\n", passed_frameworks));
        report.push_str(&format!(
            "- **Failed**: {}\n",
            total_frameworks - passed_frameworks
        ));
        report.push_str(&format!(
            "- **Success Rate**: {:.1}%\n\n",
            100.0 * passed_frameworks as f64 / total_frameworks as f64
        ));

        // Detailed results
        report.push_str("## Detailed Results\n\n");
        for (framework, result) in results {
            report.push_str(&format!("### {}\n\n", framework.as_str()));
            report.push_str(&format!(
                "- **Status**: {}\n",
                if result.passed { "✅ PASSED" } else { "❌ FAILED" }
            ));
            report.push_str(&format!("- **Max Difference**: {:.2e}\n", result.max_diff));
            report.push_str(&format!(
                "- **Mean Difference**: {:.2e}\n",
                result.mean_diff
            ));
            report.push_str(&format!("- **Pass Rate**: {:.1}%\n", result.pass_rate()));
            report.push_str(&format!(
                "- **Execution Time**: {:.2}ms\n",
                result.execution_time_ms
            ));

            if !result.errors.is_empty() {
                report.push_str("- **Errors**:\n");
                for error in &result.errors {
                    report.push_str(&format!("  - {}\n", error));
                }
            }

            if !result.metrics.is_empty() {
                report.push_str("- **Metrics**:\n");
                for (name, value) in &result.metrics {
                    report.push_str(&format!("  - {}: {}\n", name, value));
                }
            }

            report.push('\n');
        }

        report
    }

    /// Get available frameworks
    pub fn available_frameworks(&self) -> &[Framework] {
        &self.available_frameworks
    }

    /// Get test cases
    pub fn test_cases(&self) -> &[ValidationTestCase] {
        &self.test_cases
    }
}

/// Helper struct for tensor comparison results
#[derive(Debug)]
struct TensorComparison {
    max_diff: f64,
    mean_diff: f64,
    mismatch_count: usize,
    total_elements: usize,
}

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

    #[test]
    fn test_framework_detection() {
        let mut validator = CrossFrameworkValidator::with_defaults();
        validator.detect_frameworks().expect("operation failed in test");

        // Should at least have TrustformeRS
        assert!(validator.available_frameworks().contains(&Framework::TrustformeRS));
    }

    #[test]
    fn test_tensor_comparison() {
        let validator = CrossFrameworkValidator::with_defaults();

        let tensor1 = Tensor::zeros(&[2, 2]).expect("Failed to create zero tensor");
        let tensor2 = Tensor::zeros(&[2, 2]).expect("Failed to create zero tensor");

        let result =
            validator.compare_tensors(&tensor1, &tensor2).expect("tensor operation failed");
        assert!(result.passed);
        assert_eq!(result.max_diff, 0.0);
    }

    #[test]
    fn test_validation_config() {
        let config = ValidationConfig {
            atol: 1e-6,
            rtol: 1e-5,
            max_errors: 5,
            validate_gradients: true,
            target_frameworks: vec![Framework::PyTorch],
            model_architecture: "gpt".to_string(),
            model_params: None,
            onnx_reference_model: None,
        };

        assert_eq!(config.atol, 1e-6);
        assert_eq!(config.target_frameworks.len(), 1);
    }

    /// Build a tiny reference ONNX model computing `output = input + bias`.
    fn write_reference_add_model(path: &std::path::Path) -> Result<()> {
        use crate::export::onnx::{
            ONNXDataType, ONNXDimension, ONNXGraph, ONNXModel, ONNXNode, ONNXOpsetImport,
            ONNXTensor, ONNXTensorShape, ONNXTensorType, ONNXTypeInfo, ONNXValueInfo,
        };
        use crate::export::onnx_proto::encode_model;

        let value_info = |name: &str, dims: &[i64]| ONNXValueInfo {
            name: name.to_string(),
            type_info: ONNXTypeInfo {
                tensor_type: ONNXTensorType {
                    elem_type: ONNXDataType::Float,
                    shape: ONNXTensorShape {
                        dims: dims.iter().map(|dim| ONNXDimension::Value(*dim)).collect(),
                    },
                },
            },
        };

        let bias_values: [f32; 4] = [0.5, -0.25, 1.0, 2.0];
        let mut raw_data = Vec::with_capacity(16);
        for value in bias_values {
            raw_data.extend_from_slice(&value.to_le_bytes());
        }

        let model = ONNXModel {
            graph: ONNXGraph {
                nodes: vec![ONNXNode {
                    op_type: "Add".to_string(),
                    inputs: vec!["input".to_string(), "bias".to_string()],
                    outputs: vec!["output".to_string()],
                    attributes: HashMap::new(),
                    name: "add".to_string(),
                }],
                inputs: vec![value_info("input", &[1, 4])],
                outputs: vec![value_info("output", &[1, 4])],
                initializers: vec![ONNXTensor {
                    name: "bias".to_string(),
                    data_type: ONNXDataType::Float,
                    dims: vec![1, 4],
                    raw_data,
                }],
                name: "reference".to_string(),
            },
            ir_version: 8,
            opset_imports: vec![ONNXOpsetImport {
                domain: String::new(),
                version: 13,
            }],
            producer_name: "trustformers-test".to_string(),
            producer_version: "0".to_string(),
            model_version: 1,
        };

        std::fs::write(path, encode_model(&model))
            .map_err(|error| TrustformersError::io_error(error.to_string()))?;
        Ok(())
    }

    fn temp_model_path(name: &str) -> std::path::PathBuf {
        std::env::temp_dir().join(format!(
            "trustformers_xframework_{}_{}.onnx",
            name,
            std::process::id()
        ))
    }

    /// Regression test: `validate_onnx` used to set `passed = true` with
    /// invented diff statistics whenever `ONNX_AVAILABLE` was set in the
    /// environment. It must now really execute the reference graph.
    #[test]
    fn test_onnx_validation_runs_the_reference_graph() -> Result<()> {
        let path = temp_model_path("match");
        write_reference_add_model(&path)?;

        let config = ValidationConfig::default().with_onnx_reference_model(path.clone());
        let mut validator = CrossFrameworkValidator::new(config);

        let input = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0], &[1, 4])?;
        // The graph adds [0.5, -0.25, 1.0, 2.0].
        let expected = Tensor::from_vec(vec![1.5, 1.75, 4.0, 6.0], &[1, 4])?;
        validator.add_test_case(
            ValidationTestCase::new("add".to_string(), vec![input])
                .with_input_names(vec!["input".to_string()])
                .with_expected_output("output", expected),
        );

        let result = validator.validate_framework(Framework::OnnxRuntime)?;
        assert!(
            result.errors.is_empty(),
            "unexpected errors: {:?}",
            result.errors
        );
        assert!(result.passed, "matching outputs must pass");
        assert_eq!(result.total_elements, 4, "four real elements were compared");
        assert_eq!(result.mismatch_count, 0);
        assert!(result.max_diff < 1e-6);
        // The old stub reported exactly these fabricated figures.
        assert_ne!(result.total_elements, 1000);
        assert!(!result.metrics.contains_key("onnx_version"));

        std::fs::remove_file(&path).ok();
        Ok(())
    }

    /// A genuinely wrong expected output must fail, not pass.
    #[test]
    fn test_onnx_validation_detects_a_real_mismatch() -> Result<()> {
        let path = temp_model_path("mismatch");
        write_reference_add_model(&path)?;

        let config = ValidationConfig::default().with_onnx_reference_model(path.clone());
        let mut validator = CrossFrameworkValidator::new(config);

        let input = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0], &[1, 4])?;
        let wrong = Tensor::from_vec(vec![9.0, 9.0, 9.0, 9.0], &[1, 4])?;
        validator.add_test_case(
            ValidationTestCase::new("add".to_string(), vec![input])
                .with_input_names(vec!["input".to_string()])
                .with_expected_output("output", wrong),
        );

        let result = validator.validate_framework(Framework::OnnxRuntime)?;
        assert!(!result.passed, "a real mismatch must not pass");
        assert_eq!(result.mismatch_count, 4);
        assert!(result.max_diff > 1.0);

        std::fs::remove_file(&path).ok();
        Ok(())
    }

    /// With no reference model configured, ONNX validation must report why,
    /// not pass.
    #[test]
    fn test_onnx_validation_without_a_reference_model() -> Result<()> {
        let validator = CrossFrameworkValidator::with_defaults();
        let result = validator.validate_framework(Framework::OnnxRuntime)?;
        assert!(!result.passed);
        assert!(result.errors.iter().any(|error| error.contains("reference model")));
        assert_eq!(result.total_elements, 0);
        Ok(())
    }

    /// Regression test: TF/JAX used to report a PASS with fabricated diff
    /// statistics whenever an environment variable was set.
    #[test]
    fn test_tensorflow_and_jax_are_never_reported_as_passing() -> Result<()> {
        let mut validator = CrossFrameworkValidator::with_defaults();
        validator.detect_frameworks()?;

        assert!(!validator.available_frameworks().contains(&Framework::TensorFlow));
        assert!(!validator.available_frameworks().contains(&Framework::Jax));
        assert!(!validator.available_frameworks().contains(&Framework::PyTorch));
        assert!(validator.available_frameworks().contains(&Framework::OnnxRuntime));

        for framework in [Framework::TensorFlow, Framework::Jax, Framework::PyTorch] {
            let result = validator.validate_framework(framework)?;
            assert!(!result.passed, "{:?} must never report a pass", framework);
            assert_eq!(result.max_diff, 0.0);
            assert_eq!(result.mean_diff, 0.0);
            assert_eq!(result.total_elements, 0);
            assert_eq!(result.mismatch_count, 0);
            assert!(
                result.metrics.is_empty(),
                "no version metric may be invented"
            );
            assert!(!result.errors.is_empty());
        }

        Ok(())
    }

    #[test]
    fn test_test_case_builder() {
        let inputs = vec![Tensor::zeros(&[2, 2]).expect("Failed to create zero tensor")];
        let test_case = ValidationTestCase::new("test".to_string(), inputs)
            .with_expected_shape(vec![2, 2])
            .with_model_config("layers".to_string(), serde_json::json!(12));

        assert_eq!(test_case.name, "test");
        assert_eq!(test_case.expected_shape, vec![2, 2]);
        assert!(test_case.model_config.contains_key("layers"));
    }

    #[test]
    fn test_validation_result() {
        let mut result = ValidationResult::new(Framework::PyTorch);
        result.total_elements = 100;
        result.mismatch_count = 5;

        assert_eq!(result.pass_rate(), 95.0);

        result.add_metric("version".to_string(), 2.1);
        assert!(result.metrics.contains_key("version"));
    }

    #[test]
    fn test_report_generation() {
        let mut results = HashMap::new();
        let mut result = ValidationResult::new(Framework::PyTorch);
        result.passed = true;
        result.max_diff = 1e-6;
        result.mean_diff = 1e-7;
        results.insert(Framework::PyTorch, result);

        let validator = CrossFrameworkValidator::with_defaults();
        let report = validator.generate_report(&results);

        assert!(report.contains("Cross-Framework Validation Report"));
        assert!(report.contains("PASSED"));
    }
}