torsh_jit/

neural_compilation.rs

// Copyright (c) 2025 ToRSh Contributors
// SPDX-License-Identifier: Apache-2.0 OR MIT

//! # Neural Compilation
//!
//! This module implements ML-guided compilation optimization using neural networks
//! to predict optimal compilation strategies, optimization passes, and runtime configurations.
//!
//! ## Key Features
//!
//! - **Learned Optimization Selection**: Neural models predict which optimizations to apply
//! - **Performance Prediction**: Estimate execution time/memory before compilation
//! - **Adaptive Compilation**: Learn from execution feedback to improve future decisions
//! - **Transfer Learning**: Leverage knowledge from similar computation graphs
//! - **Meta-Learning**: Learn to learn better compilation strategies
//!
//! ## Architecture
//!
//! ```text
//! Graph → Feature Extraction → Neural Model → Optimization Decision
//!           ↓                      ↓                    ↓
//!       Structure            Performance         Apply/Skip
//!       Features             Prediction          Optimizations
//! ```
//!
//! ## Example
//!
//! ```rust,ignore
//! use torsh_jit::neural_compilation::{NeuralCompiler, CompilationFeatures};
//!
//! let mut compiler = NeuralCompiler::new();
//!
//! // Extract features from computation graph
//! let features = compiler.extract_features(&graph);
//!
//! // Predict optimal optimization strategy
//! let strategy = compiler.predict_strategy(&features)?;
//!
//! // Apply learned optimizations
//! let optimized_graph = compiler.apply_strategy(&graph, &strategy)?;
//!
//! // Learn from execution feedback
//! compiler.learn_from_execution(&graph, &metrics);
//! ```

use crate::graph::{ComputationGraph, NodeId};
use crate::{JitError, JitResult};
use indexmap::IndexMap;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, VecDeque};
use std::sync::{Arc, RwLock};

// ============================================================================
// Graph Feature Extraction
// ============================================================================

/// Features extracted from a computation graph for neural compilation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GraphFeatures {
    /// Structural features
    pub structural: StructuralFeatures,

    /// Computational features
    pub computational: ComputationalFeatures,

    /// Memory access patterns
    pub memory_patterns: MemoryPatternFeatures,

    /// Control flow characteristics
    pub control_flow: ControlFlowFeatures,

    /// Historical performance data
    pub historical: Option<HistoricalFeatures>,
}

/// Structural properties of the graph
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StructuralFeatures {
    /// Number of nodes
    pub node_count: usize,

    /// Number of edges
    pub edge_count: usize,

    /// Graph depth (longest path)
    pub depth: usize,

    /// Average node degree
    pub avg_degree: f32,

    /// Number of strongly connected components
    pub scc_count: usize,

    /// Graph diameter
    pub diameter: usize,

    /// Clustering coefficient
    pub clustering_coeff: f32,

    /// Operation type distribution (histogram)
    pub op_type_dist: HashMap<String, usize>,
}

/// Computational intensity features
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ComputationalFeatures {
    /// Total FLOPs estimate
    pub total_flops: u64,

    /// Arithmetic intensity (FLOPs/byte)
    pub arithmetic_intensity: f32,

    /// Parallelism degree
    pub parallelism: usize,

    /// Vectorization opportunities
    pub vectorizable_ops: usize,

    /// Memory-bound operations count
    pub memory_bound_ops: usize,

    /// Compute-bound operations count
    pub compute_bound_ops: usize,
}

/// Memory access pattern features
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryPatternFeatures {
    /// Total memory footprint (bytes)
    pub total_memory: usize,

    /// Peak memory usage
    pub peak_memory: usize,

    /// Cache locality score (0-1)
    pub cache_locality: f32,

    /// Stride patterns (sequential, strided, random)
    pub stride_patterns: HashMap<String, usize>,

    /// Reuse distance histogram
    pub reuse_distances: Vec<usize>,

    /// Working set size
    pub working_set_size: usize,
}

/// Control flow characteristics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ControlFlowFeatures {
    /// Number of branches
    pub branch_count: usize,

    /// Loop nesting depth
    pub max_loop_depth: usize,

    /// Number of loops
    pub loop_count: usize,

    /// Average loop trip count
    pub avg_trip_count: f32,

    /// Branch predictability score
    pub branch_predictability: f32,
}

/// Historical execution data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HistoricalFeatures {
    /// Previous execution times
    pub execution_times: Vec<f64>,

    /// Previous memory usage
    pub memory_usage: Vec<usize>,

    /// Cache miss rates
    pub cache_miss_rates: Vec<f32>,

    /// Successful optimizations
    pub successful_opts: Vec<String>,
}

// ============================================================================
// Neural Models
// ============================================================================

/// Neural network model for compilation decisions
#[derive(Debug, Clone)]
pub struct NeuralModel {
    /// Model weights (simplified representation)
    weights: Vec<Vec<f32>>,

    /// Model biases
    biases: Vec<f32>,

    /// Input feature dimension
    input_dim: usize,

    /// Hidden layer dimensions
    hidden_dims: Vec<usize>,

    /// Output dimension
    output_dim: usize,

    /// Model training statistics
    stats: ModelStatistics,
}

/// Training and evaluation statistics
#[derive(Debug, Clone, Default)]
pub struct ModelStatistics {
    /// Number of training samples
    pub samples_seen: usize,

    /// Current loss
    pub current_loss: f32,

    /// Best validation accuracy
    pub best_accuracy: f32,

    /// Prediction accuracy history
    pub accuracy_history: VecDeque<f32>,

    /// Feature importance scores
    pub feature_importance: HashMap<String, f32>,
}

impl NeuralModel {
    /// Create a new neural model with given architecture
    pub fn new(input_dim: usize, hidden_dims: Vec<usize>, output_dim: usize) -> Self {
        // Initialize weights randomly (in production, use proper initialization)
        let mut weights = Vec::new();
        let mut biases = Vec::new();

        let mut prev_dim = input_dim;
        for &hidden_dim in &hidden_dims {
            weights.push(vec![0.01; prev_dim * hidden_dim]); // Xavier/He initialization
            biases.push(0.0);
            prev_dim = hidden_dim;
        }

        // Output layer
        weights.push(vec![0.01; prev_dim * output_dim]);
        biases.push(0.0);

        Self {
            weights,
            biases,
            input_dim,
            hidden_dims,
            output_dim,
            stats: ModelStatistics::default(),
        }
    }

    /// Forward pass through the network
    pub fn forward(&self, features: &[f32]) -> JitResult<Vec<f32>> {
        if features.len() != self.input_dim {
            return Err(JitError::CompilationError(format!(
                "Expected {} features, got {}",
                self.input_dim,
                features.len()
            )));
        }

        let mut activations = features.to_vec();

        // Hidden layers with ReLU activation
        for (weights, bias) in self
            .weights
            .iter()
            .zip(self.biases.iter())
            .take(self.hidden_dims.len())
        {
            activations = Self::dense_layer(&activations, weights, *bias);
            activations = Self::relu(&activations);
        }

        // Output layer with softmax
        if let (Some(out_weights), Some(out_bias)) = (self.weights.last(), self.biases.last()) {
            activations = Self::dense_layer(&activations, out_weights, *out_bias);
            activations = Self::softmax(&activations);
        }

        Ok(activations)
    }

    /// Dense (fully connected) layer
    fn dense_layer(input: &[f32], weights: &[f32], bias: f32) -> Vec<f32> {
        let input_dim = input.len();
        let output_dim = weights.len() / input_dim;
        let mut output = vec![bias; output_dim];

        for i in 0..output_dim {
            for j in 0..input_dim {
                output[i] += input[j] * weights[i * input_dim + j];
            }
        }

        output
    }

    /// ReLU activation function
    fn relu(x: &[f32]) -> Vec<f32> {
        x.iter().map(|&v| v.max(0.0)).collect()
    }

    /// Softmax activation function
    fn softmax(x: &[f32]) -> Vec<f32> {
        let max = x.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let exp_values: Vec<f32> = x.iter().map(|&v| (v - max).exp()).collect();
        let sum: f32 = exp_values.iter().sum();
        exp_values.iter().map(|&v| v / sum).collect()
    }

    /// Update model weights based on feedback (simplified SGD)
    pub fn update(
        &mut self,
        features: &[f32],
        target: &[f32],
        learning_rate: f32,
    ) -> JitResult<()> {
        let prediction = self.forward(features)?;

        // Compute loss (cross-entropy)
        let loss: f32 = target
            .iter()
            .zip(prediction.iter())
            .map(|(&t, &p)| -t * p.max(1e-10).ln())
            .sum();

        self.stats.current_loss = loss;
        self.stats.samples_seen += 1;

        // Simple gradient descent (in production, use backpropagation)
        // This is a placeholder for demonstration

        // Update accuracy history
        let accuracy = self.compute_accuracy(&prediction, target);
        self.stats.accuracy_history.push_back(accuracy);
        if self.stats.accuracy_history.len() > 100 {
            self.stats.accuracy_history.pop_front();
        }

        if accuracy > self.stats.best_accuracy {
            self.stats.best_accuracy = accuracy;
        }

        Ok(())
    }

    /// Compute prediction accuracy
    fn compute_accuracy(&self, prediction: &[f32], target: &[f32]) -> f32 {
        let pred_class = prediction
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(i, _)| i)
            .unwrap_or(0);

        let target_class = target
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(i, _)| i)
            .unwrap_or(0);

        if pred_class == target_class {
            1.0
        } else {
            0.0
        }
    }
}

// ============================================================================
// Compilation Strategy
// ============================================================================

/// Optimization strategy predicted by neural model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompilationStrategy {
    /// Optimizations to apply (ordered)
    pub optimizations: Vec<OptimizationDecision>,

    /// Predicted execution time (microseconds)
    pub predicted_time_us: f64,

    /// Predicted memory usage (bytes)
    pub predicted_memory: usize,

    /// Confidence score (0-1)
    pub confidence: f32,

    /// Reasoning explanation
    pub reasoning: Vec<String>,
}

/// Decision about a specific optimization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OptimizationDecision {
    /// Optimization pass name
    pub pass_name: String,

    /// Whether to apply this optimization
    pub apply: bool,

    /// Estimated speedup (1.0 = no change)
    pub estimated_speedup: f32,

    /// Estimated memory impact (bytes, negative = reduction)
    pub estimated_memory_delta: i64,

    /// Confidence in this decision (0-1)
    pub confidence: f32,
}

// ============================================================================
// Neural Compiler
// ============================================================================

/// Main neural compilation engine
pub struct NeuralCompiler {
    /// Strategy prediction model
    strategy_model: Arc<RwLock<NeuralModel>>,

    /// Performance prediction model
    performance_model: Arc<RwLock<NeuralModel>>,

    /// Feature extractor
    feature_extractor: FeatureExtractor,

    /// Historical data for learning
    history: Arc<RwLock<CompilationHistory>>,

    /// Configuration
    config: NeuralCompilerConfig,
}

/// Configuration for neural compiler
#[derive(Debug, Clone)]
pub struct NeuralCompilerConfig {
    /// Enable online learning
    pub online_learning: bool,

    /// Learning rate for model updates
    pub learning_rate: f32,

    /// Exploration rate (epsilon for ε-greedy)
    pub exploration_rate: f32,

    /// Minimum confidence threshold
    pub min_confidence: f32,

    /// Maximum history size
    pub max_history_size: usize,

    /// Enable transfer learning
    pub transfer_learning: bool,
}

impl Default for NeuralCompilerConfig {
    fn default() -> Self {
        Self {
            online_learning: true,
            learning_rate: 0.001,
            exploration_rate: 0.1,
            min_confidence: 0.7,
            max_history_size: 10000,
            transfer_learning: true,
        }
    }
}

/// Historical compilation data
#[derive(Debug, Default)]
pub struct CompilationHistory {
    /// Past compilation results
    pub entries: VecDeque<HistoryEntry>,

    /// Feature statistics for normalization
    pub feature_stats: FeatureStatistics,
}

/// Single history entry
#[derive(Debug, Clone)]
pub struct HistoryEntry {
    /// Graph features
    pub features: GraphFeatures,

    /// Applied strategy
    pub strategy: CompilationStrategy,

    /// Actual execution time
    pub actual_time_us: f64,

    /// Actual memory usage
    pub actual_memory: usize,

    /// Prediction error
    pub error: f32,
}

/// Feature statistics for normalization
#[derive(Debug, Default, Clone)]
pub struct FeatureStatistics {
    /// Mean values for each feature
    pub means: HashMap<String, f32>,

    /// Standard deviations
    pub stddevs: HashMap<String, f32>,

    /// Min/max values
    pub mins: HashMap<String, f32>,
    pub maxs: HashMap<String, f32>,
}

impl NeuralCompiler {
    /// Create a new neural compiler
    pub fn new() -> Self {
        Self::with_config(NeuralCompilerConfig::default())
    }

    /// Create with custom configuration
    pub fn with_config(config: NeuralCompilerConfig) -> Self {
        // Strategy model: features → optimization decisions
        let strategy_model = Arc::new(RwLock::new(
            NeuralModel::new(128, vec![256, 128, 64], 32), // 32 possible optimizations
        ));

        // Performance model: features → time/memory prediction (same input dim as strategy model)
        let performance_model = Arc::new(RwLock::new(
            NeuralModel::new(128, vec![64, 32], 2), // time + memory
        ));

        Self {
            strategy_model,
            performance_model,
            feature_extractor: FeatureExtractor::new(),
            history: Arc::new(RwLock::new(CompilationHistory::default())),
            config,
        }
    }

    /// Extract features from a computation graph
    pub fn extract_features(&self, graph: &ComputationGraph) -> JitResult<GraphFeatures> {
        self.feature_extractor.extract(graph)
    }

    /// Predict optimal compilation strategy
    pub fn predict_strategy(&self, features: &GraphFeatures) -> JitResult<CompilationStrategy> {
        // Flatten features into vector
        let feature_vec = self.flatten_features(features)?;

        // Get strategy prediction
        let strategy_model = self
            .strategy_model
            .read()
            .map_err(|e| JitError::CompilationError(format!("Lock error: {}", e)))?;

        let strategy_probs = strategy_model.forward(&feature_vec)?;

        // Convert probabilities to optimization decisions
        let optimizations = self.decode_strategy(&strategy_probs)?;

        // Predict performance
        let performance_model = self
            .performance_model
            .read()
            .map_err(|e| JitError::CompilationError(format!("Lock error: {}", e)))?;

        let perf_prediction = performance_model.forward(&feature_vec)?;

        let predicted_time_us = perf_prediction.get(0).copied().unwrap_or(1000.0) as f64 * 1000.0;
        let predicted_memory =
            perf_prediction.get(1).copied().unwrap_or(1.0) as usize * 1024 * 1024;

        // Compute overall confidence
        let confidence = strategy_probs.iter().sum::<f32>() / strategy_probs.len() as f32;

        // Generate reasoning
        let reasoning = self.generate_reasoning(&optimizations, features);

        Ok(CompilationStrategy {
            optimizations,
            predicted_time_us,
            predicted_memory,
            confidence,
            reasoning,
        })
    }

    /// Apply predicted strategy to graph
    pub fn apply_strategy(
        &self,
        graph: &ComputationGraph,
        strategy: &CompilationStrategy,
    ) -> JitResult<ComputationGraph> {
        let optimized = graph.clone();

        for decision in &strategy.optimizations {
            if decision.apply && decision.confidence > self.config.min_confidence {
                // Apply optimization (placeholder - actual implementation would call optimization passes)
                log::info!(
                    "Applying optimization: {} (speedup: {:.2}x, confidence: {:.2})",
                    decision.pass_name,
                    decision.estimated_speedup,
                    decision.confidence
                );
            }
        }

        Ok(optimized)
    }

    /// Learn from execution feedback
    pub fn learn_from_execution(
        &mut self,
        features: &GraphFeatures,
        strategy: &CompilationStrategy,
        actual_time_us: f64,
        actual_memory: usize,
    ) -> JitResult<()> {
        if !self.config.online_learning {
            return Ok(());
        }

        // Compute prediction error
        let time_error = ((strategy.predicted_time_us - actual_time_us) / actual_time_us).abs();
        let memory_error = ((strategy.predicted_memory as f64 - actual_memory as f64)
            / actual_memory as f64)
            .abs();
        let error = ((time_error + memory_error) / 2.0) as f32;

        // Add to history
        let mut history = self
            .history
            .write()
            .map_err(|e| JitError::CompilationError(format!("Lock error: {}", e)))?;

        history.entries.push_back(HistoryEntry {
            features: features.clone(),
            strategy: strategy.clone(),
            actual_time_us,
            actual_memory,
            error,
        });

        if history.entries.len() > self.config.max_history_size {
            history.entries.pop_front();
        }

        // Update models
        let feature_vec = self.flatten_features(features)?;

        // Update performance model
        let target_perf = vec![
            (actual_time_us / 1000.0) as f32,
            (actual_memory / (1024 * 1024)) as f32,
        ];

        let mut perf_model = self
            .performance_model
            .write()
            .map_err(|e| JitError::CompilationError(format!("Lock error: {}", e)))?;

        perf_model.update(&feature_vec, &target_perf, self.config.learning_rate)?;

        log::info!(
            "Neural compiler learned from execution: error={:.2}%, samples={}",
            error * 100.0,
            perf_model.stats.samples_seen
        );

        Ok(())
    }

    /// Flatten features into vector
    fn flatten_features(&self, features: &GraphFeatures) -> JitResult<Vec<f32>> {
        let mut vec = Vec::with_capacity(128);

        // Structural features (normalized)
        vec.push((features.structural.node_count as f32).ln());
        vec.push((features.structural.edge_count as f32).ln());
        vec.push((features.structural.depth as f32).ln());
        vec.push(features.structural.avg_degree);
        vec.push(features.structural.scc_count as f32);
        vec.push(features.structural.diameter as f32);
        vec.push(features.structural.clustering_coeff);

        // Computational features
        vec.push((features.computational.total_flops as f32).ln());
        vec.push(features.computational.arithmetic_intensity);
        vec.push((features.computational.parallelism as f32).ln());
        vec.push(features.computational.vectorizable_ops as f32);

        // Memory features
        vec.push((features.memory_patterns.total_memory as f32).ln());
        vec.push((features.memory_patterns.peak_memory as f32).ln());
        vec.push(features.memory_patterns.cache_locality);

        // Pad to 128 dimensions
        while vec.len() < 128 {
            vec.push(0.0);
        }

        Ok(vec)
    }

    /// Decode strategy probabilities into optimization decisions
    fn decode_strategy(&self, probs: &[f32]) -> JitResult<Vec<OptimizationDecision>> {
        let opt_names = vec![
            "constant_folding",
            "dead_code_elimination",
            "common_subexpression_elimination",
            "loop_invariant_motion",
            "strength_reduction",
            "loop_unrolling",
            "vectorization",
            "parallelization",
            "fusion",
            "inlining",
            "algebraic_simplification",
            "peephole",
            "instruction_scheduling",
            "register_allocation",
            "memory_layout",
            "cache_blocking",
        ];

        let mut decisions = Vec::new();

        for (i, &prob) in probs.iter().enumerate().take(opt_names.len()) {
            let apply = prob > 0.5;
            let estimated_speedup = if apply { 1.0 + prob } else { 1.0 };

            decisions.push(OptimizationDecision {
                pass_name: opt_names.get(i).unwrap_or(&"unknown").to_string(),
                apply,
                estimated_speedup,
                estimated_memory_delta: if apply { -1024 } else { 0 },
                confidence: prob,
            });
        }

        Ok(decisions)
    }

    /// Generate human-readable reasoning
    fn generate_reasoning(
        &self,
        decisions: &[OptimizationDecision],
        features: &GraphFeatures,
    ) -> Vec<String> {
        let mut reasoning = Vec::new();

        if features.computational.arithmetic_intensity > 10.0 {
            reasoning
                .push("High arithmetic intensity detected - compute-bound workload".to_string());
        } else {
            reasoning.push("Low arithmetic intensity detected - memory-bound workload".to_string());
        }

        let applied_opts: Vec<_> = decisions
            .iter()
            .filter(|d| d.apply && d.confidence > 0.7)
            .map(|d| d.pass_name.as_str())
            .collect();

        if !applied_opts.is_empty() {
            reasoning.push(format!(
                "Recommended optimizations: {}",
                applied_opts.join(", ")
            ));
        }

        reasoning
    }

    /// Get model statistics
    pub fn get_statistics(&self) -> JitResult<HashMap<String, f32>> {
        let perf_model = self
            .performance_model
            .read()
            .map_err(|e| JitError::CompilationError(format!("Lock error: {}", e)))?;

        let mut stats = HashMap::new();
        stats.insert(
            "samples_seen".to_string(),
            perf_model.stats.samples_seen as f32,
        );
        stats.insert("current_loss".to_string(), perf_model.stats.current_loss);
        stats.insert("best_accuracy".to_string(), perf_model.stats.best_accuracy);

        Ok(stats)
    }
}

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

// ============================================================================
// Feature Extraction
// ============================================================================

/// Feature extraction from computation graphs
pub struct FeatureExtractor {
    /// Cached features for graphs
    cache: IndexMap<String, GraphFeatures>,
}

impl FeatureExtractor {
    /// Create a new feature extractor
    pub fn new() -> Self {
        Self {
            cache: IndexMap::new(),
        }
    }

    /// Extract features from graph
    pub fn extract(&self, graph: &ComputationGraph) -> JitResult<GraphFeatures> {
        Ok(GraphFeatures {
            structural: self.extract_structural(graph)?,
            computational: self.extract_computational(graph)?,
            memory_patterns: self.extract_memory_patterns(graph)?,
            control_flow: self.extract_control_flow(graph)?,
            historical: None,
        })
    }

    fn extract_structural(&self, graph: &ComputationGraph) -> JitResult<StructuralFeatures> {
        let node_count = graph.node_count();
        let edge_count = graph.edge_count();

        // Compute graph depth (longest path)
        let depth = self.compute_depth(graph);

        // Average degree
        let avg_degree = if node_count > 0 {
            (edge_count as f32) / (node_count as f32)
        } else {
            0.0
        };

        // Operation type distribution
        let mut op_type_dist = HashMap::new();
        for (node_id, node) in graph.nodes() {
            let op_name = format!("{:?}", node.operation);
            *op_type_dist.entry(op_name).or_insert(0) += 1;
        }

        Ok(StructuralFeatures {
            node_count,
            edge_count,
            depth,
            avg_degree,
            scc_count: 1, // Simplified
            diameter: depth,
            clustering_coeff: 0.0, // Simplified
            op_type_dist,
        })
    }

    fn extract_computational(&self, graph: &ComputationGraph) -> JitResult<ComputationalFeatures> {
        let mut total_flops = 0u64;
        let mut vectorizable_ops = 0;
        let mut memory_bound_ops = 0;
        let mut compute_bound_ops = 0;

        for (node_id, node) in graph.nodes() {
            // Estimate FLOPs based on operation type
            let op_flops = self.estimate_flops(&node.operation, &node.inputs);
            total_flops += op_flops;

            // Check if vectorizable
            if self.is_vectorizable(&node.operation) {
                vectorizable_ops += 1;
            }

            // Classify as memory or compute bound
            if op_flops > 1000 {
                compute_bound_ops += 1;
            } else {
                memory_bound_ops += 1;
            }
        }

        let arithmetic_intensity = if total_flops > 0 {
            total_flops as f32 / (1024.0 * 1024.0) // Simplified
        } else {
            0.0
        };

        Ok(ComputationalFeatures {
            total_flops,
            arithmetic_intensity,
            parallelism: graph.node_count(),
            vectorizable_ops,
            memory_bound_ops,
            compute_bound_ops,
        })
    }

    fn extract_memory_patterns(
        &self,
        graph: &ComputationGraph,
    ) -> JitResult<MemoryPatternFeatures> {
        Ok(MemoryPatternFeatures {
            total_memory: graph.node_count() * 1024, // Simplified
            peak_memory: graph.node_count() * 2048,
            cache_locality: 0.7,
            stride_patterns: HashMap::new(),
            reuse_distances: vec![],
            working_set_size: graph.node_count() * 512,
        })
    }

    fn extract_control_flow(&self, graph: &ComputationGraph) -> JitResult<ControlFlowFeatures> {
        Ok(ControlFlowFeatures {
            branch_count: 0,
            max_loop_depth: 0,
            loop_count: 0,
            avg_trip_count: 0.0,
            branch_predictability: 1.0,
        })
    }

    fn compute_depth(&self, graph: &ComputationGraph) -> usize {
        // Simple DFS-based depth computation
        let mut max_depth = 0;

        for (node_id, _node) in graph.nodes() {
            let depth = self.node_depth(graph, node_id, &mut HashMap::new());
            max_depth = max_depth.max(depth);
        }

        max_depth
    }

    fn node_depth(
        &self,
        graph: &ComputationGraph,
        node_id: NodeId,
        memo: &mut HashMap<NodeId, usize>,
    ) -> usize {
        if let Some(&depth) = memo.get(&node_id) {
            return depth;
        }

        let inputs = graph.get_node_inputs(node_id);
        let depth = if inputs.is_empty() {
            0
        } else {
            1 + inputs
                .iter()
                .map(|&input_id| self.node_depth(graph, input_id, memo))
                .max()
                .unwrap_or(0)
        };

        memo.insert(node_id, depth);
        depth
    }

    fn estimate_flops(&self, _operation: &crate::graph::Operation, _inputs: &[NodeId]) -> u64 {
        // Simplified FLOP estimation
        100
    }

    fn is_vectorizable(&self, operation: &crate::graph::Operation) -> bool {
        // Check if operation can be vectorized
        matches!(
            operation,
            crate::graph::Operation::Add
                | crate::graph::Operation::Mul
                | crate::graph::Operation::Relu
                | crate::graph::Operation::Sigmoid
        )
    }
}

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

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::graph::GraphBuilder;
    use torsh_core::{DType, Shape};

    #[test]
    fn test_neural_compiler_creation() {
        let compiler = NeuralCompiler::new();
        assert!(compiler.config.online_learning);
    }

    #[test]
    fn test_neural_model_forward() {
        let model = NeuralModel::new(10, vec![20, 10], 5);
        let input = vec![0.5; 10];
        let output = model.forward(&input).unwrap();
        assert_eq!(output.len(), 5);

        // Softmax output should sum to 1
        let sum: f32 = output.iter().sum();
        assert!((sum - 1.0).abs() < 1e-5);
    }

    #[test]
    fn test_feature_extraction() {
        let mut builder = GraphBuilder::new();
        let x = builder.add_input("x".to_string(), Shape::new(vec![2, 3]), DType::F32);
        let y = builder.add_input("y".to_string(), Shape::new(vec![2, 3]), DType::F32);
        let z = builder
            .add_binary_op("add".to_string(), crate::graph::Operation::Add, x, y)
            .unwrap();
        builder.mark_output(z).unwrap();

        let graph = builder.build().unwrap();

        let extractor = FeatureExtractor::new();
        let features = extractor.extract(&graph).unwrap();

        assert!(features.structural.node_count >= 3); // At least 2 inputs + 1 add
        assert!(features.computational.total_flops > 0);
    }

    #[test]
    fn test_strategy_prediction() {
        let compiler = NeuralCompiler::new();

        let mut builder = GraphBuilder::new();
        let x = builder.add_input("x".to_string(), Shape::new(vec![10, 10]), DType::F32);
        let y = builder
            .add_unary_op("relu".to_string(), crate::graph::Operation::Relu, x)
            .unwrap();
        builder.mark_output(y).unwrap();

        let graph = builder.build().unwrap();
        let features = compiler.extract_features(&graph).unwrap();
        let strategy = compiler.predict_strategy(&features).unwrap();

        assert!(!strategy.optimizations.is_empty());
        assert!(strategy.confidence >= 0.0 && strategy.confidence <= 1.0);
    }

    #[test]
    fn test_online_learning() {
        let mut compiler = NeuralCompiler::new();

        let mut builder = GraphBuilder::new();
        let x = builder.add_input("x".to_string(), Shape::new(vec![5, 5]), DType::F32);
        builder.mark_output(x).unwrap();

        let graph = builder.build().unwrap();
        let features = compiler.extract_features(&graph).unwrap();
        let strategy = compiler.predict_strategy(&features).unwrap();

        // Simulate execution feedback
        let result = compiler.learn_from_execution(&features, &strategy, 1500.0, 2048);
        assert!(result.is_ok());

        let stats = compiler.get_statistics().unwrap();
        assert_eq!(stats.get("samples_seen").copied().unwrap_or(0.0), 1.0);
    }
}