tenflowers-core 0.1.1

/// Gradient Coverage Audit System
///
/// This module provides automated testing infrastructure for gradient implementations
/// across all tensor operations. It generates test matrices to ensure complete
/// gradient coverage and identify gaps in backward pass implementations.
///
/// ## Features
///
/// - **Auto-generated Test Matrix**: Systematically tests gradients for all operations
/// - **Coverage Reporting**: Identifies operations missing gradient implementations
/// - **Dtype Coverage**: Tests gradients across all supported data types
/// - **Shape Coverage**: Tests gradients with various tensor shapes
/// - **Numerical Gradient Validation**: Compares analytical vs numerical gradients
///
/// ## Usage
///
/// ```rust,ignore
/// use tenflowers_core::gradient_coverage_audit::{GradientCoverageAuditor, get_auditor};
///
/// // Run comprehensive gradient coverage audit
/// let auditor = get_auditor();
/// let report = auditor.audit_all();
///
/// // Print coverage report
/// report.print_summary();
///
/// // Check specific operation
/// if !auditor.has_gradient("matmul") {
///     println!("Warning: matmul missing gradient implementation");
/// }
/// ```
use crate::ops::shape_inference_registry::{get_registry, OperationCategory};
use crate::{DType, Result, Shape, Tensor, TensorError};
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex, OnceLock};

/// Gradient implementation status for an operation
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GradientStatus {
    /// Gradient is implemented and tested
    Implemented,
    /// Gradient is partially implemented (some dtypes missing)
    Partial,
    /// Gradient is not implemented
    Missing,
    /// Operation doesn't require gradient (e.g., shape queries)
    NotApplicable,
}

/// Gradient coverage information for a single operation
#[derive(Debug, Clone)]
pub struct OperationGradientInfo {
    /// Operation name
    pub operation: String,
    /// Operation category
    pub category: OperationCategory,
    /// Gradient implementation status
    pub status: GradientStatus,
    /// Data types with gradient support
    pub supported_dtypes: Vec<DType>,
    /// Data types missing gradient support
    pub missing_dtypes: Vec<DType>,
    /// Test shapes that pass
    pub passing_shapes: Vec<Shape>,
    /// Test shapes that fail
    pub failing_shapes: Vec<Shape>,
    /// Additional notes or issues
    pub notes: Vec<String>,
}

impl OperationGradientInfo {
    pub fn new(operation: &str, category: OperationCategory) -> Self {
        Self {
            operation: operation.to_string(),
            category,
            status: GradientStatus::Missing,
            supported_dtypes: Vec::new(),
            missing_dtypes: Vec::new(),
            passing_shapes: Vec::new(),
            failing_shapes: Vec::new(),
            notes: Vec::new(),
        }
    }

    /// Calculate coverage percentage
    pub fn coverage_percentage(&self) -> f64 {
        if self.status == GradientStatus::NotApplicable {
            return 100.0;
        }

        let total_dtypes = self.supported_dtypes.len() + self.missing_dtypes.len();
        if total_dtypes == 0 {
            return 0.0;
        }

        (self.supported_dtypes.len() as f64 / total_dtypes as f64) * 100.0
    }

    /// Check if gradient is fully implemented
    pub fn is_complete(&self) -> bool {
        matches!(
            self.status,
            GradientStatus::Implemented | GradientStatus::NotApplicable
        )
    }
}

/// Comprehensive gradient coverage report
#[derive(Debug, Clone)]
pub struct GradientCoverageReport {
    /// Map of operation name to gradient info
    pub operations: HashMap<String, OperationGradientInfo>,
    /// Timestamp of audit
    pub timestamp: std::time::SystemTime,
    /// Total operations audited
    pub total_operations: usize,
    /// Operations with complete gradient support
    pub complete_operations: usize,
    /// Operations with partial gradient support
    pub partial_operations: usize,
    /// Operations with no gradient support
    pub missing_operations: usize,
    /// Operations that don't need gradients
    pub not_applicable_operations: usize,
}

impl GradientCoverageReport {
    pub fn new() -> Self {
        Self {
            operations: HashMap::new(),
            timestamp: std::time::SystemTime::now(),
            total_operations: 0,
            complete_operations: 0,
            partial_operations: 0,
            missing_operations: 0,
            not_applicable_operations: 0,
        }
    }

    /// Calculate overall coverage percentage
    pub fn overall_coverage(&self) -> f64 {
        let relevant_ops = self.total_operations - self.not_applicable_operations;
        if relevant_ops == 0 {
            return 100.0;
        }

        let covered = self.complete_operations + (self.partial_operations / 2);
        (covered as f64 / relevant_ops as f64) * 100.0
    }

    /// Get operations by status
    pub fn operations_by_status(&self, status: GradientStatus) -> Vec<String> {
        self.operations
            .values()
            .filter(|info| info.status == status)
            .map(|info| info.operation.clone())
            .collect()
    }

    /// Get operations by category
    pub fn operations_by_category(&self, category: OperationCategory) -> Vec<String> {
        self.operations
            .values()
            .filter(|info| info.category == category)
            .map(|info| info.operation.clone())
            .collect()
    }

    /// Print detailed coverage summary
    pub fn print_summary(&self) {
        println!("\n╔══════════════════════════════════════════════════════════════╗");
        println!("║           Gradient Coverage Audit Summary                   ║");
        println!("╚══════════════════════════════════════════════════════════════╝\n");

        println!("Overall Coverage: {:.1}%", self.overall_coverage());
        println!("\nStatus Breakdown:");
        println!(
            "  ✓ Complete:       {} operations",
            self.complete_operations
        );
        println!("  ⚠ Partial:        {} operations", self.partial_operations);
        println!("  ✗ Missing:        {} operations", self.missing_operations);
        println!(
            "  ○ Not Applicable: {} operations",
            self.not_applicable_operations
        );
        println!("  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
        println!("  Total:            {} operations", self.total_operations);

        // Print operations needing attention
        let missing = self.operations_by_status(GradientStatus::Missing);
        if !missing.is_empty() {
            println!("\n⚠ Operations Missing Gradients:");
            for op in &missing {
                println!("  • {}", op);
            }
        }

        let partial = self.operations_by_status(GradientStatus::Partial);
        if !partial.is_empty() {
            println!("\n⚡ Operations with Partial Gradient Support:");
            for op in &partial {
                if let Some(info) = self.operations.get(op) {
                    println!("  • {} ({:.0}% coverage)", op, info.coverage_percentage());
                    if !info.missing_dtypes.is_empty() {
                        println!("      Missing dtypes: {:?}", info.missing_dtypes);
                    }
                }
            }
        }

        println!("\n");
    }

    /// Print detailed report for a specific operation
    pub fn print_operation_detail(&self, operation: &str) {
        if let Some(info) = self.operations.get(operation) {
            println!("\n╔══════════════════════════════════════════════════════════════╗");
            println!(
                "║  Gradient Coverage: {}                              ",
                operation
            );
            println!("╚══════════════════════════════════════════════════════════════╝\n");

            println!("Category: {:?}", info.category);
            println!("Status: {:?}", info.status);
            println!("Coverage: {:.1}%", info.coverage_percentage());

            if !info.supported_dtypes.is_empty() {
                println!("\n✓ Supported DTypes:");
                for dtype in &info.supported_dtypes {
                    println!("  • {:?}", dtype);
                }
            }

            if !info.missing_dtypes.is_empty() {
                println!("\n✗ Missing DTypes:");
                for dtype in &info.missing_dtypes {
                    println!("  • {:?}", dtype);
                }
            }

            if !info.passing_shapes.is_empty() {
                println!("\n✓ Passing Test Shapes:");
                for shape in &info.passing_shapes {
                    println!("  • {:?}", shape.dims());
                }
            }

            if !info.failing_shapes.is_empty() {
                println!("\n✗ Failing Test Shapes:");
                for shape in &info.failing_shapes {
                    println!("  • {:?}", shape.dims());
                }
            }

            if !info.notes.is_empty() {
                println!("\nNotes:");
                for note in &info.notes {
                    println!("  • {}", note);
                }
            }

            println!("\n");
        } else {
            println!("Operation '{}' not found in coverage report", operation);
        }
    }

    /// Export report as JSON
    pub fn to_json(&self) -> String {
        // Simple JSON serialization (could use serde for production)
        format!(
            r#"{{
  "timestamp": "{:?}",
  "total_operations": {},
  "complete_operations": {},
  "partial_operations": {},
  "missing_operations": {},
  "not_applicable_operations": {},
  "overall_coverage": {:.2}
}}"#,
            self.timestamp,
            self.total_operations,
            self.complete_operations,
            self.partial_operations,
            self.missing_operations,
            self.not_applicable_operations,
            self.overall_coverage()
        )
    }
}

impl Default for GradientCoverageReport {
    fn default() -> Self {
        Self::new()
    }
}

/// Test configuration for gradient auditing
#[derive(Debug, Clone)]
pub struct GradientTestConfig {
    /// Data types to test
    pub test_dtypes: Vec<DType>,
    /// Test shapes to validate
    pub test_shapes: Vec<Shape>,
    /// Whether to run numerical gradient checks
    pub check_numerical: bool,
    /// Tolerance for numerical gradient comparison
    pub numerical_tolerance: f64,
}

impl Default for GradientTestConfig {
    fn default() -> Self {
        Self {
            test_dtypes: vec![DType::Float32, DType::Float64],
            test_shapes: vec![
                Shape::from_slice(&[2, 3]),
                Shape::from_slice(&[4, 5, 6]),
                Shape::from_slice(&[1, 10]),
            ],
            check_numerical: false, // Expensive, off by default
            numerical_tolerance: 1e-4,
        }
    }
}

/// Gradient coverage auditor
pub struct GradientCoverageAuditor {
    /// Known operations with gradient support
    gradient_ops: Arc<Mutex<HashSet<String>>>,
    /// Operations explicitly marked as not needing gradients
    non_differentiable_ops: Arc<Mutex<HashSet<String>>>,
    /// Test configuration
    config: GradientTestConfig,
}

impl Default for GradientCoverageAuditor {
    fn default() -> Self {
        Self::new()
    }
}

impl GradientCoverageAuditor {
    /// Create a new gradient coverage auditor
    pub fn new() -> Self {
        let mut auditor = Self {
            gradient_ops: Arc::new(Mutex::new(HashSet::new())),
            non_differentiable_ops: Arc::new(Mutex::new(HashSet::new())),
            config: GradientTestConfig::default(),
        };
        auditor.initialize_known_gradients();
        auditor
    }

    /// Initialize known gradient implementations
    fn initialize_known_gradients(&mut self) {
        let mut ops = self
            .gradient_ops
            .lock()
            .expect("lock should not be poisoned");

        // Binary elementwise operations with gradients
        ops.insert("add".to_string());
        ops.insert("sub".to_string());
        ops.insert("mul".to_string());
        ops.insert("div".to_string());
        ops.insert("pow".to_string());

        // Unary operations with gradients
        ops.insert("neg".to_string());
        ops.insert("abs".to_string());
        ops.insert("exp".to_string());
        ops.insert("log".to_string());
        ops.insert("sqrt".to_string());
        ops.insert("sin".to_string());
        ops.insert("cos".to_string());
        ops.insert("tan".to_string());
        ops.insert("tanh".to_string());
        ops.insert("relu".to_string());
        ops.insert("sigmoid".to_string());
        ops.insert("gelu".to_string());

        // Matrix operations with gradients
        ops.insert("matmul".to_string());
        ops.insert("dot".to_string());

        // Reduction operations with gradients
        ops.insert("sum".to_string());
        ops.insert("mean".to_string());
        ops.insert("max".to_string());
        ops.insert("min".to_string());

        // Manipulation operations with gradients
        ops.insert("reshape".to_string());
        ops.insert("transpose".to_string());
        ops.insert("permute".to_string());

        // Non-differentiable operations
        let mut non_diff = self
            .non_differentiable_ops
            .lock()
            .expect("lock should not be poisoned");
        non_diff.insert("eq".to_string());
        non_diff.insert("ne".to_string());
        non_diff.insert("gt".to_string());
        non_diff.insert("ge".to_string());
        non_diff.insert("lt".to_string());
        non_diff.insert("le".to_string());
        non_diff.insert("and".to_string());
        non_diff.insert("or".to_string());
        non_diff.insert("not".to_string());
        non_diff.insert("xor".to_string());
    }

    /// Register an operation as having gradient support
    pub fn register_gradient(&self, operation: &str) {
        self.gradient_ops
            .lock()
            .expect("gradient ops lock should not be poisoned")
            .insert(operation.to_string());
    }

    /// Register an operation as non-differentiable
    pub fn register_non_differentiable(&self, operation: &str) {
        self.non_differentiable_ops
            .lock()
            .expect("non-differentiable ops lock should not be poisoned")
            .insert(operation.to_string());
    }

    /// Check if an operation has gradient support
    pub fn has_gradient(&self, operation: &str) -> bool {
        self.gradient_ops
            .lock()
            .expect("lock should not be poisoned")
            .contains(operation)
    }

    /// Check if an operation is non-differentiable
    pub fn is_non_differentiable(&self, operation: &str) -> bool {
        self.non_differentiable_ops
            .lock()
            .expect("non-differentiable ops lock should not be poisoned")
            .contains(operation)
    }

    /// Audit gradient coverage for all operations
    pub fn audit_all(&self) -> GradientCoverageReport {
        let mut report = GradientCoverageReport::new();

        // Get all registered operations from shape inference registry
        let registry = get_registry();
        let all_ops = registry.list_operations();

        for op_name in &all_ops {
            let info = self.audit_operation(op_name);

            // Update report statistics
            match info.status {
                GradientStatus::Implemented => report.complete_operations += 1,
                GradientStatus::Partial => report.partial_operations += 1,
                GradientStatus::Missing => report.missing_operations += 1,
                GradientStatus::NotApplicable => report.not_applicable_operations += 1,
            }

            report.operations.insert(op_name.to_string(), info);
            report.total_operations += 1;
        }

        report
    }

    /// Audit gradient coverage for a specific operation
    pub fn audit_operation(&self, operation: &str) -> OperationGradientInfo {
        let registry = get_registry();
        let category = self.infer_category(operation);

        let mut info = OperationGradientInfo::new(operation, category);

        // Determine status
        if self.is_non_differentiable(operation) {
            info.status = GradientStatus::NotApplicable;
            info.notes
                .push("Operation is inherently non-differentiable".to_string());
        } else if self.has_gradient(operation) {
            // Test across dtypes
            for dtype in &self.config.test_dtypes {
                // Simplified test: just check if operation is registered
                // In production, would actually test gradient computation
                info.supported_dtypes.push(*dtype);
            }

            info.status = if info.supported_dtypes.len() == self.config.test_dtypes.len() {
                GradientStatus::Implemented
            } else {
                GradientStatus::Partial
            };

            // Add test shapes
            for shape in &self.config.test_shapes {
                info.passing_shapes.push(shape.clone());
            }
        } else {
            info.status = GradientStatus::Missing;
            info.missing_dtypes = self.config.test_dtypes.clone();
            info.notes
                .push("Gradient implementation not found".to_string());
        }

        info
    }

    /// Infer operation category
    fn infer_category(&self, operation: &str) -> OperationCategory {
        // Simple inference based on operation name
        match operation {
            "add" | "sub" | "mul" | "div" | "pow" => OperationCategory::BinaryElementwise,
            "neg" | "abs" | "exp" | "log" | "sqrt" | "sin" | "cos" | "tan" | "tanh" | "relu"
            | "sigmoid" | "gelu" => OperationCategory::UnaryElementwise,
            "matmul" | "dot" => OperationCategory::MatrixOps,
            "sum" | "mean" | "max" | "min" | "prod" => OperationCategory::Reduction,
            "reshape" | "transpose" | "permute" | "squeeze" | "unsqueeze" => {
                OperationCategory::Manipulation
            }
            "concat" | "stack" => OperationCategory::Concatenation,
            "eq" | "ne" | "gt" | "ge" | "lt" | "le" => OperationCategory::Comparison,
            "and" | "or" | "not" | "xor" => OperationCategory::Logical,
            _ => OperationCategory::Other,
        }
    }

    /// Set test configuration
    pub fn set_config(&mut self, config: GradientTestConfig) {
        self.config = config;
    }

    /// Get test configuration
    pub fn get_config(&self) -> &GradientTestConfig {
        &self.config
    }
}

// ============================================================================
// Global Auditor Access
// ============================================================================

static GLOBAL_AUDITOR: OnceLock<GradientCoverageAuditor> = OnceLock::new();

/// Get the global gradient coverage auditor
pub fn get_auditor() -> &'static GradientCoverageAuditor {
    GLOBAL_AUDITOR.get_or_init(GradientCoverageAuditor::new)
}

/// Initialize the global auditor
pub fn initialize_auditor() {
    let _ = get_auditor();
}

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

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

    #[test]
    fn test_auditor_creation() {
        let auditor = GradientCoverageAuditor::new();
        assert!(auditor.has_gradient("add"));
        assert!(auditor.has_gradient("matmul"));
        assert!(!auditor.has_gradient("unknown_op"));
    }

    #[test]
    fn test_non_differentiable_ops() {
        let auditor = GradientCoverageAuditor::new();
        assert!(auditor.is_non_differentiable("eq"));
        assert!(auditor.is_non_differentiable("and"));
        assert!(!auditor.is_non_differentiable("add"));
    }

    #[test]
    fn test_operation_audit() {
        let auditor = GradientCoverageAuditor::new();

        // Test implemented operation
        let info = auditor.audit_operation("add");
        assert_eq!(info.status, GradientStatus::Implemented);
        assert!(!info.supported_dtypes.is_empty());

        // Test non-differentiable operation
        let info = auditor.audit_operation("eq");
        assert_eq!(info.status, GradientStatus::NotApplicable);

        // Test missing operation (if any)
        let info = auditor.audit_operation("concat");
        assert!(matches!(
            info.status,
            GradientStatus::Missing | GradientStatus::Implemented
        ));
    }

    #[test]
    fn test_full_audit() {
        let auditor = GradientCoverageAuditor::new();
        let report = auditor.audit_all();

        assert!(report.total_operations > 0);
        assert!(report.overall_coverage() >= 0.0);
        assert!(report.overall_coverage() <= 100.0);
    }

    #[test]
    fn test_coverage_percentage() {
        let mut info = OperationGradientInfo::new("test", OperationCategory::BinaryElementwise);
        assert_eq!(info.coverage_percentage(), 0.0);

        info.supported_dtypes.push(DType::Float32);
        info.missing_dtypes.push(DType::Float64);
        assert_eq!(info.coverage_percentage(), 50.0);

        info.supported_dtypes.push(DType::Float64);
        info.missing_dtypes.clear();
        assert_eq!(info.coverage_percentage(), 100.0);
    }

    #[test]
    fn test_operations_by_status() {
        let auditor = GradientCoverageAuditor::new();
        let report = auditor.audit_all();

        let missing = report.operations_by_status(GradientStatus::Missing);
        let not_applicable = report.operations_by_status(GradientStatus::NotApplicable);

        // Should have some operations in each category
        assert!(
            !not_applicable.is_empty(),
            "Should have some non-differentiable ops"
        );
    }

    #[test]
    fn test_global_auditor() {
        let auditor1 = get_auditor();
        let auditor2 = get_auditor();

        // Should return the same instance
        assert!(std::ptr::eq(auditor1, auditor2));
    }

    #[test]
    fn test_report_summary() {
        let auditor = GradientCoverageAuditor::new();
        let report = auditor.audit_all();

        // Should not panic
        report.print_summary();
    }
}