optirs-tpu 0.3.2

use std::fmt::Debug;
// Execution scheduling optimization for XLA computations
//
// This module implements various scheduling strategies for XLA operations,
// including dependency-aware scheduling, resource-aware scheduling,
// latency hiding, and multi-core parallelization.

use scirs2_core::numeric::Float;
use std::cmp::Ordering;
use std::collections::{BTreeMap, BinaryHeap, HashMap, HashSet, VecDeque};

use super::super::frontend::{OperandId, OperationId, OperationType, XLAComputation, XLAOperation};
use super::{HardwareTarget, OptimizationPipelineConfig};
use crate::error::Result;

/// Execution scheduler for XLA computations
pub struct ExecutionScheduler<T: Float + Debug + Send + Sync + 'static> {
    /// Scheduling configuration
    config: SchedulingConfig,

    /// Dependency analyzer
    dependency_analyzer: DependencyAnalyzer<T>,

    /// Resource manager
    resource_manager: ResourceManager,

    /// Performance predictor
    performance_predictor: PerformancePredictor<T>,

    /// Scheduling statistics
    scheduling_stats: SchedulingStatistics,
}

/// Scheduling configuration
#[derive(Debug, Clone)]
pub struct SchedulingConfig {
    /// Scheduling strategy
    pub strategy: SchedulingStrategy,

    /// Enable resource-aware scheduling
    pub enable_resource_aware: bool,

    /// Enable latency hiding
    pub enable_latency_hiding: bool,

    /// Enable multi-core parallelization
    pub enable_parallelization: bool,

    /// Maximum scheduling lookahead
    pub max_lookahead: usize,

    /// Resource utilization target
    pub resource_utilization_target: f64,

    /// Critical path optimization priority
    pub critical_path_priority: f64,

    /// Memory bandwidth consideration
    pub memory_bandwidth_weight: f64,
}

/// Scheduling strategies
#[derive(Debug, Clone, PartialEq)]
pub enum SchedulingStrategy {
    /// Topological ordering (dependency-based)
    Topological,

    /// Critical path first
    CriticalPath,

    /// Resource-aware scheduling
    ResourceAware,

    /// Load balancing
    LoadBalancing,

    /// Latency-optimized scheduling
    LatencyOptimized,

    /// Memory bandwidth optimized
    MemoryBandwidthOptimized,

    /// Custom scheduling algorithm
    Custom(String),
}

/// Dependency analyzer for operations
pub struct DependencyAnalyzer<T: Float + Debug + Send + Sync + 'static> {
    /// Dependency graph
    dependency_graph: DependencyGraph,

    /// Critical path analysis
    critical_path_analyzer: CriticalPathAnalyzer<T>,
}

/// Dependency graph representation
#[derive(Debug)]
pub struct DependencyGraph {
    /// Adjacency list (operation -> dependencies)
    pub dependencies: HashMap<OperationId, Vec<OperationId>>,

    /// Reverse adjacency list (operation -> dependents)
    pub dependents: HashMap<OperationId, Vec<OperationId>>,

    /// In-degree count for topological sorting
    pub in_degrees: HashMap<OperationId, usize>,

    /// Strongly connected components
    pub scc: Vec<Vec<OperationId>>,
}

/// Critical path analyzer
pub struct CriticalPathAnalyzer<T: Float + Debug + Send + Sync + 'static> {
    /// Critical path length for each operation
    critical_path_lengths: HashMap<OperationId, f64>,

    /// Critical path operations
    critical_operations: Vec<OperationId>,

    /// Path analysis cache
    _phantom: std::marker::PhantomData<T>,
}

/// Path analysis results
#[derive(Debug, Clone)]
pub struct PathAnalysis {
    /// Path length
    pub length: f64,

    /// Operations on path
    pub operations: Vec<OperationId>,

    /// Bottleneck operations
    pub bottlenecks: Vec<OperationId>,
}

/// Data flow pattern
#[derive(Debug, Clone)]
pub enum DataFlowPattern {
    /// Pipeline pattern
    Pipeline,

    /// Fan-out pattern
    FanOut,

    /// Fan-in pattern
    FanIn,

    /// Scatter-gather pattern
    ScatterGather,

    /// Reduction pattern
    Reduction,

    /// Broadcast pattern
    Broadcast,
}

/// Memory access pattern for scheduling
#[derive(Debug, Clone)]
pub struct MemoryAccessPattern {
    /// Access type
    pub access_type: MemoryAccessType,

    /// Access frequency
    pub frequency: f64,

    /// Access size
    pub size: usize,

    /// Access locality
    pub locality: LocalityType,
}

/// Memory access types for scheduling
#[derive(Debug, Clone)]
pub enum MemoryAccessType {
    /// Sequential access
    Sequential,

    /// Random access
    Random,

    /// Strided access
    Strided { stride: usize },

    /// Gather access
    Gather,

    /// Scatter access
    Scatter,
}

/// Memory locality types
#[derive(Debug, Clone)]
pub enum LocalityType {
    /// Temporal locality
    Temporal,

    /// Spatial locality
    Spatial,

    /// No locality
    None,
}

/// Resource manager for scheduling
/// Resource manager for the scheduler.
///
/// It holds no resource inventory of its own: the compute/memory resource
/// lists, the per-operation allocation map and the utilization timeline that
/// used to be declared here were all built empty and never read.
/// `assign_resources` derives assignments from the hardware target it is given,
/// which is the only resource description that was ever consulted.
pub struct ResourceManager;

/// Compute resource information
#[derive(Debug, Clone)]
pub struct ComputeResource {
    /// Resource identifier
    pub id: String,

    /// Resource type
    pub resource_type: ComputeResourceType,

    /// Capacity (operations per second)
    pub capacity: f64,

    /// Current utilization (0.0-1.0)
    pub utilization: f64,

    /// Power consumption (watts)
    pub power: f64,
}

/// Types of compute resources
#[derive(Debug, Clone)]
pub enum ComputeResourceType {
    /// Matrix multiplication unit
    MatrixUnit,

    /// Vector processing unit
    VectorUnit,

    /// Scalar processing unit
    ScalarUnit,

    /// Memory controller
    MemoryController,

    /// Special function unit
    SpecialFunction,
}

/// Memory resource information
#[derive(Debug, Clone)]
pub struct MemoryResource {
    /// Resource identifier
    pub id: String,

    /// Memory level
    pub level: MemoryLevel,

    /// Capacity (bytes)
    pub capacity: usize,

    /// Bandwidth (bytes/second)
    pub bandwidth: f64,

    /// Current usage
    pub usage: usize,
}

/// Memory levels for resource management
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum MemoryLevel {
    /// L1 cache
    L1Cache,

    /// L2 cache
    L2Cache,

    /// High bandwidth memory
    HBM,

    /// Host memory
    HostMemory,
}

/// Resource allocation for operation
#[derive(Debug, Clone)]
pub struct ResourceAllocation {
    /// Allocated compute resources
    pub compute_resources: Vec<String>,

    /// Allocated memory resources
    pub memory_resources: Vec<String>,

    /// Allocation time range
    pub time_range: (u64, u64),

    /// Resource requirements
    pub requirements: ResourceRequirements,
}

/// Resource requirements for operations
#[derive(Debug, Clone, Default)]
pub struct ResourceRequirements {
    /// Compute requirements (FLOPS)
    pub compute_flops: f64,

    /// Memory requirements (bytes)
    pub memory_bytes: usize,

    /// Memory bandwidth (bytes/second)
    pub memory_bandwidth: f64,

    /// Execution time estimate (microseconds)
    pub execution_time_us: u64,

    /// Parallelization factor
    pub parallelization_factor: f64,
}

/// Resource utilization snapshot
#[derive(Debug, Clone, Default)]
pub struct ResourceUtilization {
    /// Compute utilization by resource type
    pub compute_utilization: HashMap<ComputeResourceType, f64>,

    /// Memory utilization by level
    pub memory_utilization: HashMap<MemoryLevel, f64>,

    /// Overall utilization
    pub overall_utilization: f64,
}

/// Latency hiding strategies
#[derive(Debug, Clone)]
pub enum LatencyHidingStrategy {
    /// Overlap computation with communication
    ComputeCommsOverlap,

    /// Prefetch data
    Prefetching,

    /// Pipeline execution
    Pipelining,

    /// Speculative execution
    SpeculativeExecution,
}

/// Prefetch opportunity
#[derive(Debug)]
pub struct PrefetchOpportunity {
    /// Operation to prefetch for
    pub target_operation: OperationId,

    /// Data to prefetch
    pub data_operand: OperandId,

    /// Prefetch distance (operations ahead)
    pub prefetch_distance: usize,

    /// Expected benefit
    pub benefit: f64,
}

/// Parallelization strategies
#[derive(Debug, Clone)]
pub enum ParallelizationStrategy {
    /// Data parallelism
    DataParallel,

    /// Model parallelism
    ModelParallel,

    /// Pipeline parallelism
    PipelineParallel,

    /// Task parallelism
    TaskParallel,
}

/// Parallel execution block
#[derive(Debug)]
pub struct ParallelBlock {
    /// Block identifier
    pub id: String,

    /// Operations in block
    pub operations: Vec<OperationId>,

    /// Parallelization type
    pub parallelization_type: ParallelizationStrategy,

    /// Target resources
    pub target_resources: Vec<String>,
}

/// Synchronization point
#[derive(Debug)]
pub struct SynchronizationPoint {
    /// Synchronization identifier
    pub id: String,

    /// Operations to synchronize
    pub operations: Vec<OperationId>,

    /// Synchronization type
    pub sync_type: SynchronizationType,
}

/// Types of synchronization
#[derive(Debug)]
pub enum SynchronizationType {
    /// Barrier synchronization
    Barrier,

    /// Point-to-point synchronization
    PointToPoint,

    /// Collective synchronization
    Collective,
}

/// Load balancing strategies
#[derive(Debug, Clone)]
pub enum LoadBalancingStrategy {
    /// Round-robin distribution
    RoundRobin,

    /// Weighted distribution
    Weighted,

    /// Dynamic load balancing
    Dynamic,

    /// Work stealing
    WorkStealing,
}

/// Performance metric
#[derive(Debug, Clone)]
pub struct PerformanceMetric {
    /// Metric name
    pub name: String,

    /// Current value
    pub value: f64,

    /// Target value
    pub target: f64,

    /// Trend direction
    pub trend: TrendDirection,
}

/// Performance trend direction
#[derive(Debug, Clone)]
pub enum TrendDirection {
    /// Improving
    Improving,

    /// Stable
    Stable,

    /// Degrading
    Degrading,
}

/// Performance snapshot
#[derive(Debug)]
pub struct PerformanceSnapshot {
    /// Timestamp
    pub timestamp: u64,

    /// Resource utilization
    pub resource_utilization: ResourceUtilization,

    /// Throughput
    pub throughput: f64,

    /// Latency
    pub latency: f64,
}

/// Performance predictor
/// Performance predictor for a scheduled operation list.
///
/// It carries no model registry, historical-sample buffer or per-model accuracy
/// table: all three were empty and unread, and `predict_performance` has no
/// trained model behind it -- see its own documentation.
pub struct PerformancePredictor<T: Float + Debug + Send + Sync + 'static> {
    _phantom: std::marker::PhantomData<T>,
}

/// Prediction model
#[derive(Debug)]
pub struct PredictionModel {
    /// Model type
    pub model_type: ModelType,

    /// Model parameters
    pub parameters: HashMap<String, f64>,

    /// Model accuracy
    pub accuracy: f64,
}

/// Types of prediction models
#[derive(Debug)]
pub enum ModelType {
    /// Linear regression
    Linear,

    /// Neural network
    NeuralNetwork,

    /// Decision tree
    DecisionTree,

    /// Performance counters based
    PerformanceCounters,
}

/// Performance data point
#[derive(Debug)]
pub struct PerformanceDataPoint {
    /// Operation characteristics
    pub operation_chars: OperationCharacteristics,

    /// Resource state
    pub resource_state: ResourceState,

    /// Actual performance
    pub actual_performance: f64,

    /// Timestamp
    pub timestamp: u64,
}

/// Operation characteristics for prediction
#[derive(Debug, Clone)]
pub struct OperationCharacteristics {
    /// Operation type
    pub op_type: OperationType,

    /// Input sizes
    pub input_sizes: Vec<usize>,

    /// Output size
    pub output_size: usize,

    /// Compute intensity
    pub compute_intensity: f64,

    /// Memory intensity
    pub memory_intensity: f64,
}

/// Resource state for prediction
#[derive(Debug, Clone)]
pub struct ResourceState {
    /// Current resource utilization
    pub utilization: ResourceUtilization,

    /// Available bandwidth
    pub available_bandwidth: f64,

    /// Memory pressure
    pub memory_pressure: f64,
}

/// Scheduling statistics
#[derive(Debug, Default)]
pub struct SchedulingStatistics {
    /// Total operations scheduled
    pub operations_scheduled: usize,

    /// Average scheduling time
    pub avg_scheduling_time: f64,

    /// Resource utilization achieved
    pub resource_utilization: f64,

    /// Critical path length
    pub critical_path_length: f64,

    /// Parallelization efficiency
    pub parallelization_efficiency: f64,

    /// Scheduling overhead
    pub scheduling_overhead: f64,
}

/// Scheduled execution plan
#[derive(Debug)]
pub struct ExecutionPlan<T: Float + Debug + Send + Sync + 'static> {
    /// Scheduled operations in execution order
    pub scheduled_operations: Vec<ScheduledOperation>,

    /// Resource assignments
    pub resource_assignments: HashMap<OperationId, ResourceAllocation>,

    /// Execution timeline
    pub timeline: ExecutionTimeline,

    /// Performance predictions
    pub performance_predictions: PerformancePredictions,

    /// Synchronization plan
    pub synchronization_plan: SynchronizationPlan,

    /// Phantom data for type parameter
    _phantom: std::marker::PhantomData<T>,
}

/// Scheduled operation with timing
#[derive(Debug)]
pub struct ScheduledOperation {
    /// Operation ID
    pub operation_id: OperationId,

    /// Scheduled start time
    pub start_time: u64,

    /// Scheduled end time
    pub end_time: u64,

    /// Assigned resources
    pub assigned_resources: Vec<String>,

    /// Priority level
    pub priority: u32,
}

/// Execution timeline
#[derive(Debug)]
pub struct ExecutionTimeline {
    /// Timeline events
    pub events: Vec<TimelineEvent>,

    /// Total execution time
    pub total_time: u64,

    /// Resource utilization over time
    pub utilization_timeline: BTreeMap<u64, ResourceUtilization>,
}

/// Timeline event
#[derive(Debug)]
pub struct TimelineEvent {
    /// Event timestamp
    pub timestamp: u64,

    /// Event type
    pub event_type: EventType,

    /// Associated operation
    pub operation_id: Option<OperationId>,

    /// Event details
    pub details: String,
}

/// Types of timeline events
#[derive(Debug)]
pub enum EventType {
    /// Operation start
    OperationStart,

    /// Operation end
    OperationEnd,

    /// Resource allocation
    ResourceAllocation,

    /// Resource deallocation
    ResourceDeallocation,

    /// Synchronization
    Synchronization,

    /// Communication
    Communication,
}

/// Performance predictions
#[derive(Debug, Default)]
pub struct PerformancePredictions {
    /// Predicted execution time
    pub execution_time: u64,

    /// Predicted throughput
    pub throughput: f64,

    /// Predicted resource utilization
    pub resource_utilization: f64,

    /// Predicted energy consumption
    pub energy_consumption: f64,

    /// Confidence intervals
    pub confidence_intervals: HashMap<String, (f64, f64)>,
}

/// Synchronization plan
#[derive(Debug)]
pub struct SynchronizationPlan {
    /// Synchronization points
    pub sync_points: Vec<SynchronizationPoint>,

    /// Communication schedule
    pub communication_schedule: Vec<CommunicationEvent>,

    /// Barrier operations
    pub barriers: Vec<BarrierOperation>,
}

/// Communication event
#[derive(Debug)]
pub struct CommunicationEvent {
    /// Source operation
    pub source: OperationId,

    /// Target operation
    pub target: OperationId,

    /// Data size
    pub data_size: usize,

    /// Scheduled time
    pub scheduled_time: u64,

    /// Communication type
    pub comm_type: CommunicationType,
}

/// Types of communication
#[derive(Debug)]
pub enum CommunicationType {
    /// Point-to-point
    PointToPoint,

    /// Broadcast
    Broadcast,

    /// All-reduce
    AllReduce,

    /// All-gather
    AllGather,

    /// All-to-all
    AllToAll,
}

/// Barrier operation
#[derive(Debug)]
pub struct BarrierOperation {
    /// Barrier ID
    pub id: String,

    /// Participating operations
    pub participants: Vec<OperationId>,

    /// Barrier time
    pub barrier_time: u64,
}

impl<T: Float + Debug + Default + std::fmt::Debug + Clone + Send + Sync> ExecutionScheduler<T> {
    /// Create new execution scheduler
    pub fn new(pipeline_config: &OptimizationPipelineConfig) -> Self {
        let config = SchedulingConfig {
            strategy: if pipeline_config.aggressive_mode {
                SchedulingStrategy::CriticalPath
            } else {
                SchedulingStrategy::ResourceAware
            },
            enable_resource_aware: true,
            enable_latency_hiding: true,
            enable_parallelization: pipeline_config.enable_graph_optimization,
            max_lookahead: 10,
            resource_utilization_target: 0.85,
            critical_path_priority: 1.5,
            memory_bandwidth_weight: 0.3,
        };

        Self {
            config: config.clone(),
            dependency_analyzer: DependencyAnalyzer::new(),
            resource_manager: ResourceManager::new(&pipeline_config.target_hardware),
            performance_predictor: PerformancePredictor::new(),
            scheduling_stats: SchedulingStatistics::default(),
        }
    }

    /// Optimize execution schedule for computation
    pub fn optimize_schedule(
        &mut self,
        computation: XLAComputation<T>,
    ) -> Result<XLAComputation<T>> {
        let started = std::time::Instant::now();

        // Analyze dependencies
        let dependency_graph = self
            .dependency_analyzer
            .analyze_dependencies(&computation)?;

        // Create execution plan
        let execution_plan = self.create_execution_plan(&computation, &dependency_graph)?;

        // Record what the schedule actually achieved. Previously
        // `scheduling_stats` was default-constructed and never updated, so
        // `scheduling_statistics()` reported zeros after every schedule.
        let operations = execution_plan.scheduled_operations.len();
        let makespan = execution_plan
            .scheduled_operations
            .iter()
            .map(|operation| operation.end_time)
            .fold(0u64, u64::max);
        let busy: u64 = execution_plan
            .scheduled_operations
            .iter()
            .map(|operation| operation.end_time.saturating_sub(operation.start_time))
            .sum();
        self.scheduling_stats.operations_scheduled += operations;
        self.scheduling_stats.avg_scheduling_time = started.elapsed().as_secs_f64();
        self.scheduling_stats.critical_path_length = makespan as f64;
        self.scheduling_stats.resource_utilization = if makespan == 0 {
            0.0
        } else {
            (busy as f64 / makespan as f64).min(1.0)
        };
        // A perfectly serial schedule has efficiency 0; one whose makespan is
        // much shorter than the summed work has exploited real parallelism.
        self.scheduling_stats.parallelization_efficiency = if busy == 0 {
            0.0
        } else {
            1.0 - (makespan as f64 / busy as f64).min(1.0)
        };
        self.scheduling_stats.scheduling_overhead = started.elapsed().as_secs_f64();

        // Apply schedule optimizations
        let optimized_computation =
            self.apply_schedule_optimizations(computation, &execution_plan)?;

        Ok(optimized_computation)
    }

    /// Statistics from every [`Self::optimize_schedule`] call so far.
    pub fn scheduling_statistics(&self) -> &SchedulingStatistics {
        &self.scheduling_stats
    }

    /// Create execution plan for computation
    fn create_execution_plan(
        &mut self,
        computation: &XLAComputation<T>,
        dependency_graph: &DependencyGraph,
    ) -> Result<ExecutionPlan<T>> {
        // Schedule operations based on strategy
        let scheduled_operations = match self.config.strategy {
            SchedulingStrategy::Topological => {
                self.schedule_topological(computation, dependency_graph)?
            }
            SchedulingStrategy::CriticalPath => {
                self.schedule_critical_path(computation, dependency_graph)?
            }
            SchedulingStrategy::ResourceAware => {
                self.schedule_resource_aware(computation, dependency_graph)?
            }
            _ => self.schedule_default(computation, dependency_graph)?,
        };

        // Assign resources
        let resource_assignments = self
            .resource_manager
            .assign_resources(&scheduled_operations)?;

        // Create timeline
        let timeline = self.create_execution_timeline(&scheduled_operations)?;

        // Predict performance
        let performance_predictions = self
            .performance_predictor
            .predict_performance(&scheduled_operations)?;

        // Create synchronization plan
        let synchronization_plan = self.create_synchronization_plan(&scheduled_operations)?;

        Ok(ExecutionPlan {
            scheduled_operations,
            resource_assignments,
            timeline,
            performance_predictions,
            synchronization_plan,
            _phantom: std::marker::PhantomData,
        })
    }

    /// Schedule operations using topological ordering
    fn schedule_topological(
        &self,
        computation: &XLAComputation<T>,
        dependency_graph: &DependencyGraph,
    ) -> Result<Vec<ScheduledOperation>> {
        let mut scheduled = Vec::new();
        let mut current_time = 0u64;
        let mut in_degrees = dependency_graph.in_degrees.clone();
        let mut queue = VecDeque::new();

        // Find operations with no dependencies
        for (&op_id, &degree) in &in_degrees {
            if degree == 0 {
                queue.push_back(op_id);
            }
        }

        while let Some(op_id) = queue.pop_front() {
            if let Some(operation) = computation.operations.iter().find(|op| op.id == op_id) {
                let execution_time = self.estimate_execution_time(operation);

                scheduled.push(ScheduledOperation {
                    operation_id: op_id,
                    start_time: current_time,
                    end_time: current_time + execution_time,
                    assigned_resources: vec!["default".to_string()],
                    priority: 0,
                });

                current_time += execution_time;

                // Update dependencies
                if let Some(dependents) = dependency_graph.dependents.get(&op_id) {
                    for &dependent_id in dependents {
                        if let Some(degree) = in_degrees.get_mut(&dependent_id) {
                            *degree -= 1;
                            if *degree == 0 {
                                queue.push_back(dependent_id);
                            }
                        }
                    }
                }
            }
        }

        Ok(scheduled)
    }

    /// Schedule operations using critical path priority
    fn schedule_critical_path(
        &mut self,
        computation: &XLAComputation<T>,
        dependency_graph: &DependencyGraph,
    ) -> Result<Vec<ScheduledOperation>> {
        // Compute critical path lengths
        let critical_paths = self
            .dependency_analyzer
            .critical_path_analyzer
            .compute_critical_paths(computation, dependency_graph)?;

        // Priority queue based on critical path length
        let mut priority_queue = BinaryHeap::new();
        let mut scheduled = Vec::new();
        let mut current_time = 0u64;
        let mut in_degrees = dependency_graph.in_degrees.clone();

        // Add ready operations to priority queue
        for (&op_id, &degree) in &in_degrees {
            if degree == 0 {
                let priority = critical_paths.get(&op_id).unwrap_or(&0.0);
                priority_queue.push(CriticalPathItem {
                    operation_id: op_id,
                    critical_path_length: *priority,
                });
            }
        }

        while let Some(item) = priority_queue.pop() {
            if let Some(operation) = computation
                .operations
                .iter()
                .find(|op| op.id == item.operation_id)
            {
                let execution_time = self.estimate_execution_time(operation);

                scheduled.push(ScheduledOperation {
                    operation_id: item.operation_id,
                    start_time: current_time,
                    end_time: current_time + execution_time,
                    assigned_resources: vec!["default".to_string()],
                    priority: (item.critical_path_length * 100.0) as u32,
                });

                current_time += execution_time;

                // Update dependencies and add newly ready operations
                if let Some(dependents) = dependency_graph.dependents.get(&item.operation_id) {
                    for &dependent_id in dependents {
                        if let Some(degree) = in_degrees.get_mut(&dependent_id) {
                            *degree -= 1;
                            if *degree == 0 {
                                let priority = critical_paths.get(&dependent_id).unwrap_or(&0.0);
                                priority_queue.push(CriticalPathItem {
                                    operation_id: dependent_id,
                                    critical_path_length: *priority,
                                });
                            }
                        }
                    }
                }
            }
        }

        Ok(scheduled)
    }

    /// Schedule operations with resource awareness
    fn schedule_resource_aware(
        &self,
        computation: &XLAComputation<T>,
        dependency_graph: &DependencyGraph,
    ) -> Result<Vec<ScheduledOperation>> {
        // Simplified resource-aware scheduling
        self.schedule_topological(computation, dependency_graph)
    }

    /// Default scheduling strategy
    fn schedule_default(
        &self,
        computation: &XLAComputation<T>,
        dependency_graph: &DependencyGraph,
    ) -> Result<Vec<ScheduledOperation>> {
        self.schedule_topological(computation, dependency_graph)
    }

    /// Estimate execution time for operation
    fn estimate_execution_time(&self, operation: &XLAOperation<T>) -> u64 {
        // Use performance characteristics if available
        if operation.performance.execution_time_us > 0 {
            operation.performance.execution_time_us
        } else {
            // Default estimates based on operation type
            match &operation.op_type {
                OperationType::Add | OperationType::Multiply | OperationType::Subtract => 10,
                OperationType::Dot | OperationType::DotGeneral => 100,
                OperationType::Convolution(_) => 500,
                OperationType::Reduce(_) => 50,
                _ => 20,
            }
        }
    }

    /// Apply schedule optimizations to computation
    fn apply_schedule_optimizations(
        &self,
        mut computation: XLAComputation<T>,
        execution_plan: &ExecutionPlan<T>,
    ) -> Result<XLAComputation<T>> {
        // Reorder operations according to schedule
        let operation_order: HashMap<OperationId, usize> = execution_plan
            .scheduled_operations
            .iter()
            .enumerate()
            .map(|(i, sched_op)| (sched_op.operation_id, i))
            .collect();

        computation
            .operations
            .sort_by_key(|op| operation_order.get(&op.id).unwrap_or(&usize::MAX));

        Ok(computation)
    }

    /// Create execution timeline
    fn create_execution_timeline(
        &self,
        scheduled_operations: &[ScheduledOperation],
    ) -> Result<ExecutionTimeline> {
        let mut events = Vec::new();
        let mut total_time = 0u64;

        for scheduled_op in scheduled_operations {
            events.push(TimelineEvent {
                timestamp: scheduled_op.start_time,
                event_type: EventType::OperationStart,
                operation_id: Some(scheduled_op.operation_id),
                details: "Operation started".to_string(),
            });

            events.push(TimelineEvent {
                timestamp: scheduled_op.end_time,
                event_type: EventType::OperationEnd,
                operation_id: Some(scheduled_op.operation_id),
                details: "Operation completed".to_string(),
            });

            total_time = total_time.max(scheduled_op.end_time);
        }

        Ok(ExecutionTimeline {
            events,
            total_time,
            utilization_timeline: BTreeMap::new(),
        })
    }

    /// Create synchronization plan
    fn create_synchronization_plan(
        &self,
        _scheduled_operations: &[ScheduledOperation],
    ) -> Result<SynchronizationPlan> {
        Ok(SynchronizationPlan {
            sync_points: vec![],
            communication_schedule: vec![],
            barriers: vec![],
        })
    }
}

/// Critical path item for priority queue
#[derive(Debug)]
struct CriticalPathItem {
    operation_id: OperationId,
    critical_path_length: f64,
}

impl PartialEq for CriticalPathItem {
    fn eq(&self, other: &Self) -> bool {
        self.critical_path_length == other.critical_path_length
    }
}

impl Eq for CriticalPathItem {}

impl Ord for CriticalPathItem {
    fn cmp(&self, other: &Self) -> Ordering {
        self.critical_path_length
            .partial_cmp(&other.critical_path_length)
            .unwrap_or(Ordering::Equal)
    }
}

impl PartialOrd for CriticalPathItem {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<T: Float + Debug + Default + std::fmt::Debug + Clone + Send + Sync> Default
    for DependencyAnalyzer<T>
{
    fn default() -> Self {
        Self::new()
    }
}

impl<T: Float + Debug + Default + std::fmt::Debug + Clone + Send + Sync> DependencyAnalyzer<T> {
    pub fn new() -> Self {
        Self {
            dependency_graph: DependencyGraph::new(),
            critical_path_analyzer: CriticalPathAnalyzer::new(),
        }
    }

    pub fn analyze_dependencies(
        &mut self,
        computation: &XLAComputation<T>,
    ) -> Result<DependencyGraph> {
        self.dependency_graph.build_from_computation(computation)?;
        Ok(self.dependency_graph.clone())
    }
}

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

impl DependencyGraph {
    pub fn new() -> Self {
        Self {
            dependencies: HashMap::new(),
            dependents: HashMap::new(),
            in_degrees: HashMap::new(),
            scc: vec![],
        }
    }

    pub fn build_from_computation<T: Float + Debug + Send + Sync + 'static>(
        &mut self,
        computation: &XLAComputation<T>,
    ) -> Result<()> {
        // Build dependency relationships from computation
        for operation in &computation.operations {
            self.in_degrees.insert(operation.id, 0);
            self.dependencies.insert(operation.id, vec![]);
            self.dependents.insert(operation.id, vec![]);
        }

        // Add dependencies based on operand producers
        for operation in &computation.operations {
            for &input_operand in &operation.inputs {
                // Find producer of this operand
                if let Some(producer_op) = computation
                    .operations
                    .iter()
                    .find(|op| op.output == input_operand)
                {
                    // All three maps are pre-populated for every operation id
                    // in the loop above, so these entries are always present;
                    // guard with `if let` to avoid panicking regardless.
                    if let Some(deps) = self.dependencies.get_mut(&operation.id) {
                        deps.push(producer_op.id);
                    }
                    if let Some(dependents) = self.dependents.get_mut(&producer_op.id) {
                        dependents.push(operation.id);
                    }
                    if let Some(in_degree) = self.in_degrees.get_mut(&operation.id) {
                        *in_degree += 1;
                    }
                }
            }
        }

        Ok(())
    }
}

impl Clone for DependencyGraph {
    fn clone(&self) -> Self {
        Self {
            dependencies: self.dependencies.clone(),
            dependents: self.dependents.clone(),
            in_degrees: self.in_degrees.clone(),
            scc: self.scc.clone(),
        }
    }
}

impl<T: Float + Debug + Default + std::fmt::Debug + Clone + Send + Sync> Default
    for CriticalPathAnalyzer<T>
{
    fn default() -> Self {
        Self::new()
    }
}

impl<T: Float + Debug + Default + std::fmt::Debug + Clone + Send + Sync> CriticalPathAnalyzer<T> {
    pub fn new() -> Self {
        Self {
            critical_path_lengths: HashMap::new(),
            critical_operations: vec![],
            _phantom: std::marker::PhantomData,
        }
    }

    /// Per-operation latency proxy used to weight the critical path.
    ///
    /// Uses the measured `execution_time_us` when available, otherwise a
    /// type-based estimate consistent with the scheduler's own cost model.
    fn operation_cost(operation: &XLAOperation<T>) -> f64 {
        if operation.performance.execution_time_us > 0 {
            operation.performance.execution_time_us as f64
        } else {
            match &operation.op_type {
                OperationType::Add | OperationType::Multiply | OperationType::Subtract => 10.0,
                OperationType::Dot | OperationType::DotGeneral => 100.0,
                OperationType::Convolution(_) => 500.0,
                OperationType::Reduce(_) => 50.0,
                _ => 20.0,
            }
        }
    }

    /// Compute, for every operation, the length of the longest (most costly)
    /// dependency path from that operation to a sink — its critical-path
    /// length. Higher values must be scheduled earlier because they gate the
    /// largest amount of remaining work.
    ///
    /// This is a real DAG longest-path: operations are ordered topologically
    /// (Kahn's algorithm over the actual dependency graph) and relaxed in
    /// reverse, so `lp[node] = cost(node) + max(lp[successor])`.
    pub fn compute_critical_paths(
        &mut self,
        computation: &XLAComputation<T>,
        dependency_graph: &DependencyGraph,
    ) -> Result<HashMap<OperationId, f64>> {
        // Per-operation cost.
        let mut cost: HashMap<OperationId, f64> =
            HashMap::with_capacity(computation.operations.len());
        for op in &computation.operations {
            cost.insert(op.id, Self::operation_cost(op));
        }

        // Topological order (sources first) via Kahn's algorithm.
        let mut in_degrees = dependency_graph.in_degrees.clone();
        let mut ready: VecDeque<OperationId> = in_degrees
            .iter()
            .filter(|(_, &d)| d == 0)
            .map(|(&id, _)| id)
            .collect();
        let mut topo_order: Vec<OperationId> = Vec::with_capacity(in_degrees.len());
        while let Some(node) = ready.pop_front() {
            topo_order.push(node);
            if let Some(succs) = dependency_graph.dependents.get(&node) {
                for &succ in succs {
                    if let Some(d) = in_degrees.get_mut(&succ) {
                        *d = d.saturating_sub(1);
                        if *d == 0 {
                            ready.push_back(succ);
                        }
                    }
                }
            }
        }

        // Relax in reverse topological order: longest path to a sink.
        let mut lp: HashMap<OperationId, f64> = HashMap::with_capacity(cost.len());
        for &node in topo_order.iter().rev() {
            let node_cost = *cost.get(&node).unwrap_or(&0.0);
            let mut best_succ = 0.0f64;
            if let Some(succs) = dependency_graph.dependents.get(&node) {
                for &succ in succs {
                    if let Some(&v) = lp.get(&succ) {
                        best_succ = best_succ.max(v);
                    }
                }
            }
            lp.insert(node, node_cost + best_succ);
        }

        // Nodes excluded from the topological order (e.g. inside a cycle) still
        // receive their own cost as an honest lower bound.
        for op in &computation.operations {
            lp.entry(op.id)
                .or_insert_with(|| *cost.get(&op.id).unwrap_or(&0.0));
        }

        // Trace the actual critical path: start at the source with the largest
        // length, then follow the most-costly successor at each step.
        let mut critical_ops: Vec<OperationId> = Vec::new();
        let start = lp
            .iter()
            .filter(|(id, _)| dependency_graph.in_degrees.get(id).copied().unwrap_or(0) == 0)
            .max_by(|a, b| a.1.total_cmp(b.1))
            .map(|(&id, _)| id);
        if let Some(start) = start {
            let mut visited: HashSet<OperationId> = HashSet::new();
            let mut current = start;
            while visited.insert(current) {
                critical_ops.push(current);
                let next = dependency_graph.dependents.get(&current).and_then(|succs| {
                    succs
                        .iter()
                        .filter(|s| lp.contains_key(*s))
                        .max_by(|a, b| {
                            lp.get(*a)
                                .unwrap_or(&0.0)
                                .total_cmp(lp.get(*b).unwrap_or(&0.0))
                        })
                        .copied()
                });
                match next {
                    Some(n) => current = n,
                    None => break,
                }
            }
        }

        self.critical_operations = critical_ops;
        self.critical_path_lengths = lp.clone();
        Ok(lp)
    }
}

impl ResourceManager {
    pub fn new(_target_hardware: &HardwareTarget) -> Self {
        Self
    }

    pub fn assign_resources(
        &mut self,
        scheduled_operations: &[ScheduledOperation],
    ) -> Result<HashMap<OperationId, ResourceAllocation>> {
        let mut assignments = HashMap::new();

        for scheduled_op in scheduled_operations {
            let allocation = ResourceAllocation {
                compute_resources: vec!["matrix_unit_0".to_string()],
                memory_resources: vec!["hbm_0".to_string()],
                time_range: (scheduled_op.start_time, scheduled_op.end_time),
                requirements: ResourceRequirements::default(),
            };

            assignments.insert(scheduled_op.operation_id, allocation);
        }

        Ok(assignments)
    }
}

// NOTE: `LatencyOptimizer`, `LatencyModel`, `ParallelizationEngine`,
// `ParallelExecutionGraph`, this module's `LoadBalancer`, `WorkDistribution` and
// this module's `PerformanceMonitor` used to live here. Every one of them was
// constructed by `ExecutionScheduler::new` and then never read: no latency was
// ever hidden, no work was ever distributed and no metric was ever recorded
// through them. Latency-aware ordering that *is* real lives in
// `xla::backend::code_generation`'s instruction scheduler (critical path over a
// modelled machine), load balancing in `tpu_backend::DeviceManager` and
// `coordination::LoadBalancer`, and per-execution monitoring in
// `tpu_backend::PerformanceMonitor`.

impl<T: Float + Debug + Default + std::fmt::Debug + Clone + Send + Sync> Default
    for PerformancePredictor<T>
{
    fn default() -> Self {
        Self::new()
    }
}

impl<T: Float + Debug + Default + std::fmt::Debug + Clone + Send + Sync> PerformancePredictor<T> {
    pub fn new() -> Self {
        Self {
            _phantom: std::marker::PhantomData,
        }
    }

    /// Predicted performance for a schedule.
    ///
    /// There is no trained model here, so this reports the default (all-zero)
    /// predictions rather than inventing plausible numbers. Real derived
    /// estimates are available from
    /// [`super::PerformanceAnalyzer::analyze`], which computes FLOPs and
    /// timing from the graph and the memory plan.
    pub fn predict_performance(
        &mut self,
        _scheduled_operations: &[ScheduledOperation],
    ) -> Result<PerformancePredictions> {
        Ok(PerformancePredictions::default())
    }
}

#[cfg(test)]
mod tests {
    use super::super::{ComputeCapability, HardwareTarget};
    use super::*;

    #[test]
    fn test_execution_scheduler_creation() {
        let config = OptimizationPipelineConfig {
            optimization_level: crate::main_types::XLAOptimizationLevel::Standard,
            enable_graph_optimization: true,
            enable_kernel_fusion: true,
            enable_memory_optimization: true,
            enable_scheduling_optimization: true,
            max_optimization_time: 300,
            target_hardware: HardwareTarget {
                tpu_version: "v4".to_string(),
                num_cores: 4,
                memory_capacity: 1024 * 1024 * 1024,
                memory_bandwidth: 1600.0,
                compute_capability: ComputeCapability {
                    matrix_unit_dims: (128, 128),
                    vector_unit_width: 256,
                    supported_dtypes: vec!["F32".to_string()],
                    special_instructions: vec![],
                },
            },
            custom_passes: vec![],
            aggressive_mode: false,
            debug_mode: false,
        };

        let scheduler: ExecutionScheduler<f32> = ExecutionScheduler::new(&config);
        assert_eq!(scheduler.config.strategy, SchedulingStrategy::ResourceAware);
        assert!(scheduler.config.enable_resource_aware);
    }

    #[test]
    fn test_dependency_graph_creation() {
        let graph = DependencyGraph::new();
        assert!(graph.dependencies.is_empty());
        assert!(graph.dependents.is_empty());
        assert!(graph.in_degrees.is_empty());
    }
}