sklears-core 0.1.2

/// Distributed computing infrastructure for sklears-core
///
/// This module provides comprehensive distributed computing capabilities for machine learning
/// workloads, including message-passing, cluster-aware estimators, distributed datasets,
/// and fault-tolerant training frameworks.
///
/// # Key Features
///
/// - **Message Passing**: Efficient communication primitives for cluster nodes
/// - **Distributed Estimators**: ML algorithms that scale across multiple nodes
/// - **Partitioned Datasets**: Data structures optimized for distributed processing
/// - **Fault Tolerance**: Automatic recovery and checkpoint management
/// - **Load Balancing**: Dynamic work distribution across cluster nodes
/// - **Consistency Models**: Eventual and strong consistency guarantees
///
/// # Architecture
///
/// The distributed computing system is built around several core abstractions:
///
/// ## Node Communication
/// ```rust,ignore
/// use sklears_core::distributed::{MessagePassing, ClusterNode, NodeId};
///
/// // Basic message passing between cluster nodes
/// async fn example_communication(node: &dyn ClusterNode) -> Result<(), Box<dyn std::error::Error>> {
///     let target_node = NodeId::new("worker-01");
///     let message = b"training_data_chunk_1";
///
///     node.send_message(target_node, message).await?;
///     let response = node.receive_message().await?;
///
///     Ok(())
/// }
/// ```
///
/// ## Distributed Training
/// ```rust,ignore
/// use sklears_core::distributed::{DistributedEstimator, ParameterServer};
///
/// // Distributed machine learning with parameter server architecture
/// async fn example_distributed_training() -> Result<(), Box<dyn std::error::Error>> {
///     let cluster = DistributedCluster::new()
///         .with_nodes(4)
///         .with_parameter_server()
///         .build().await?;
///
///     let model = DistributedLinearRegression::new()
///         .with_cluster(cluster)
///         .with_fault_tolerance(true)
///         .build();
///
///     // Training automatically distributes across cluster
///     model.fit_distributed(&X_train, &y_train).await?;
///
///     Ok(())
/// }
/// ```
use crate::error::{Result, SklearsError};
use futures_core::future::BoxFuture;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::{Duration, SystemTime};

// =============================================================================
// Core Distributed Computing Traits
// =============================================================================

/// Unique identifier for cluster nodes
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct NodeId(pub String);

impl NodeId {
    /// Create a new node identifier
    pub fn new(id: impl Into<String>) -> Self {
        Self(id.into())
    }

    /// Get the string representation of the node ID
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl std::fmt::Display for NodeId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// Message envelope for inter-node communication
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DistributedMessage {
    /// Unique message identifier
    pub id: String,
    /// Source node identifier
    pub sender: NodeId,
    /// Target node identifier
    pub receiver: NodeId,
    /// Message type classification
    pub message_type: MessageType,
    /// Actual message payload
    pub payload: Vec<u8>,
    /// Message timestamp
    pub timestamp: SystemTime,
    /// Message priority level
    pub priority: MessagePriority,
    /// Retry count for fault tolerance
    pub retry_count: u32,
}

/// Classification of distributed messages
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum MessageType {
    /// Data transfer between nodes
    DataTransfer,
    /// Model parameter synchronization
    ParameterSync,
    /// Gradient aggregation
    GradientAggregation,
    /// Cluster coordination
    Coordination,
    /// Health check and monitoring
    HealthCheck,
    /// Fault recovery
    FaultRecovery,
    /// Load balancing
    LoadBalance,
    /// Custom application-specific messages
    Custom(String),
}

/// Message priority levels for scheduling
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub enum MessagePriority {
    /// Low priority background tasks
    Low = 0,
    /// Normal operation messages
    Normal = 1,
    /// High priority coordination
    High = 2,
    /// Critical system messages
    Critical = 3,
}

/// Core trait for message-passing communication in distributed systems
pub trait MessagePassing: Send + Sync {
    /// Send a message to a specific node
    fn send_message(
        &self,
        target: NodeId,
        message: DistributedMessage,
    ) -> BoxFuture<'_, Result<()>>;

    /// Receive the next available message
    fn receive_message(&self) -> BoxFuture<'_, Result<DistributedMessage>>;

    /// Broadcast a message to all nodes in the cluster
    fn broadcast_message(&self, message: DistributedMessage) -> BoxFuture<'_, Result<()>>;

    /// Send a message and wait for a response
    fn send_and_receive(
        &self,
        target: NodeId,
        message: DistributedMessage,
    ) -> BoxFuture<'_, Result<DistributedMessage>>;

    /// Check if any messages are available
    fn has_pending_messages(&self) -> BoxFuture<'_, Result<bool>>;

    /// Get the number of pending messages
    fn pending_message_count(&self) -> BoxFuture<'_, Result<usize>>;

    /// Flush all pending outgoing messages
    fn flush_outgoing(&self) -> BoxFuture<'_, Result<()>>;
}

/// Cluster node abstraction for distributed computing
pub trait ClusterNode: MessagePassing + Send + Sync {
    /// Get the unique identifier for this node
    fn node_id(&self) -> &NodeId;

    /// Get the current cluster membership
    fn cluster_nodes(&self) -> BoxFuture<'_, Result<Vec<NodeId>>>;

    /// Check if this node is the cluster coordinator
    fn is_coordinator(&self) -> bool;

    /// Get current node health status
    fn health_status(&self) -> BoxFuture<'_, Result<NodeHealth>>;

    /// Get node computational resources
    fn resources(&self) -> BoxFuture<'_, Result<NodeResources>>;

    /// Join a cluster
    fn join_cluster(&mut self, coordinator: NodeId) -> BoxFuture<'_, Result<()>>;

    /// Leave the current cluster
    fn leave_cluster(&mut self) -> BoxFuture<'_, Result<()>>;

    /// Handle node failure detection
    fn handle_node_failure(&mut self, failed_node: NodeId) -> BoxFuture<'_, Result<()>>;
}

/// Node health status information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeHealth {
    /// Overall health score (0.0 to 1.0)
    pub health_score: f64,
    /// CPU utilization percentage
    pub cpu_usage: f64,
    /// Memory utilization percentage
    pub memory_usage: f64,
    /// Network latency to coordinator (ms)
    pub network_latency: Duration,
    /// Last heartbeat timestamp
    pub last_heartbeat: SystemTime,
    /// Error count in last hour
    pub recent_errors: u32,
    /// Node uptime
    pub uptime: Duration,
}

/// Node computational resources
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeResources {
    /// Number of CPU cores
    pub cpu_cores: u32,
    /// Total memory in bytes
    pub total_memory: u64,
    /// Available memory in bytes
    pub available_memory: u64,
    /// GPU devices available
    pub gpu_devices: Vec<GpuDevice>,
    /// Network bandwidth (bytes/sec)
    pub network_bandwidth: u64,
    /// Storage capacity in bytes
    pub storage_capacity: u64,
    /// Custom resource tags
    pub tags: HashMap<String, String>,
}

/// GPU device information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuDevice {
    /// Device identifier
    pub device_id: u32,
    /// Device name/model
    pub name: String,
    /// Total VRAM in bytes
    pub total_memory: u64,
    /// Available VRAM in bytes
    pub available_memory: u64,
    /// Compute capability
    pub compute_capability: String,
}

// =============================================================================
// Distributed Estimator Framework
// =============================================================================

/// Core trait for distributed machine learning estimators
pub trait DistributedEstimator: Send + Sync {
    /// Associated type for training data
    type TrainingData;

    /// Associated type for prediction input
    type PredictionInput;

    /// Associated type for prediction output
    type PredictionOutput;

    /// Associated type for model parameters
    type Parameters: Serialize + for<'de> Deserialize<'de>;

    /// Fit the model using distributed training
    fn fit_distributed<'a>(
        &'a mut self,
        cluster: &'a dyn DistributedCluster,
        training_data: &Self::TrainingData,
    ) -> BoxFuture<'a, Result<()>>;

    /// Make predictions using the distributed model
    fn predict_distributed<'a>(
        &'a self,
        cluster: &dyn DistributedCluster,
        input: &'a Self::PredictionInput,
    ) -> BoxFuture<'a, Result<Self::PredictionOutput>>;

    /// Get current model parameters
    fn get_parameters(&self) -> Result<Self::Parameters>;

    /// Set model parameters
    fn set_parameters(&mut self, params: Self::Parameters) -> Result<()>;

    /// Synchronize parameters across cluster nodes
    fn sync_parameters(&mut self, cluster: &dyn DistributedCluster) -> BoxFuture<'_, Result<()>>;

    /// Get training progress information
    fn training_progress(&self) -> DistributedTrainingProgress;
}

/// Progress tracking for distributed training
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DistributedTrainingProgress {
    /// Current epoch number
    pub epoch: u32,
    /// Total epochs planned
    pub total_epochs: u32,
    /// Training loss value
    pub training_loss: f64,
    /// Validation loss value
    pub validation_loss: Option<f64>,
    /// Number of samples processed
    pub samples_processed: u64,
    /// Training start time
    pub start_time: SystemTime,
    /// Estimated completion time
    pub estimated_completion: Option<SystemTime>,
    /// Active cluster nodes
    pub active_nodes: Vec<NodeId>,
    /// Per-node training statistics
    pub node_statistics: HashMap<NodeId, NodeTrainingStats>,
}

/// Training statistics for individual nodes
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeTrainingStats {
    /// Samples processed by this node
    pub samples_processed: u64,
    /// Processing rate (samples/sec)
    pub processing_rate: f64,
    /// Current loss value
    pub current_loss: f64,
    /// Memory usage during training
    pub memory_usage: u64,
    /// CPU utilization during training
    pub cpu_utilization: f64,
}

/// Distributed cluster management interface
pub trait DistributedCluster: Send + Sync {
    /// Get all active nodes in the cluster
    fn active_nodes(&self) -> BoxFuture<'_, Result<Vec<NodeId>>>;

    /// Get the cluster coordinator node
    fn coordinator(&self) -> &NodeId;

    /// Get cluster configuration
    fn configuration(&self) -> &ClusterConfiguration;

    /// Add a new node to the cluster
    fn add_node(&mut self, node: NodeId) -> BoxFuture<'_, Result<()>>;

    /// Remove a node from the cluster
    fn remove_node(&mut self, node: NodeId) -> BoxFuture<'_, Result<()>>;

    /// Redistribute work across cluster nodes
    fn rebalance_load(&mut self) -> BoxFuture<'_, Result<()>>;

    /// Get cluster health status
    fn cluster_health(&self) -> BoxFuture<'_, Result<ClusterHealth>>;

    /// Create a checkpoint of cluster state
    fn create_checkpoint(&self) -> BoxFuture<'_, Result<ClusterCheckpoint>>;

    /// Restore from a checkpoint
    fn restore_checkpoint(&mut self, checkpoint: ClusterCheckpoint) -> BoxFuture<'_, Result<()>>;
}

/// Cluster configuration parameters
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClusterConfiguration {
    /// Maximum number of nodes
    pub max_nodes: u32,
    /// Heartbeat interval
    pub heartbeat_interval: Duration,
    /// Node failure timeout
    pub failure_timeout: Duration,
    /// Message retry limit
    pub max_retries: u32,
    /// Load balancing strategy
    pub load_balancing: LoadBalancingStrategy,
    /// Fault tolerance mode
    pub fault_tolerance: FaultToleranceMode,
    /// Consistency requirements
    pub consistency_level: ConsistencyLevel,
}

/// Load balancing strategies
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum LoadBalancingStrategy {
    /// Round-robin assignment
    RoundRobin,
    /// Assign based on node resources
    ResourceBased,
    /// Assign based on current load
    LoadBased,
    /// Assign based on data locality
    LocalityAware,
    /// Custom balancing strategy
    Custom(String),
}

/// Fault tolerance modes
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum FaultToleranceMode {
    /// No fault tolerance
    None,
    /// Basic retry mechanisms
    BasicRetry,
    /// Checkpoint-based recovery
    CheckpointRecovery,
    /// Redundant computation
    RedundantComputation,
    /// Byzantine fault tolerance
    Byzantine,
}

/// Consistency levels for distributed operations
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ConsistencyLevel {
    /// No consistency guarantees
    None,
    /// Eventually consistent
    Eventual,
    /// Strong consistency
    Strong,
    /// Causal consistency
    Causal,
    /// Sequential consistency
    Sequential,
}

/// Overall cluster health information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClusterHealth {
    /// Overall cluster health score
    pub overall_health: f64,
    /// Number of healthy nodes
    pub healthy_nodes: u32,
    /// Number of failed nodes
    pub failed_nodes: u32,
    /// Average node response time
    pub average_response_time: Duration,
    /// Total cluster throughput
    pub total_throughput: f64,
    /// Resource utilization across cluster
    pub resource_utilization: ClusterResourceUtilization,
}

/// Cluster-wide resource utilization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClusterResourceUtilization {
    /// Average CPU utilization
    pub cpu_utilization: f64,
    /// Average memory utilization
    pub memory_utilization: f64,
    /// Network utilization
    pub network_utilization: f64,
    /// Storage utilization
    pub storage_utilization: f64,
}

/// Cluster state checkpoint for fault recovery
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClusterCheckpoint {
    /// Checkpoint identifier
    pub checkpoint_id: String,
    /// Checkpoint timestamp
    pub timestamp: SystemTime,
    /// Cluster configuration at checkpoint time
    pub configuration: ClusterConfiguration,
    /// Node states at checkpoint time
    pub node_states: HashMap<NodeId, NodeCheckpoint>,
    /// Global cluster state
    pub cluster_state: Vec<u8>,
}

/// Individual node checkpoint data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeCheckpoint {
    /// Node identifier
    pub node_id: NodeId,
    /// Node state data
    pub state_data: Vec<u8>,
    /// Node health at checkpoint time
    pub health: NodeHealth,
    /// Node resources at checkpoint time
    pub resources: NodeResources,
}

// =============================================================================
// Distributed Dataset Abstractions
// =============================================================================

/// Trait for datasets that can be distributed across cluster nodes
pub trait DistributedDataset: Send + Sync {
    /// Associated type for data items
    type Item;

    /// Associated type for partitioning strategy
    type PartitionStrategy;

    /// Get the total size of the dataset
    fn size(&self) -> u64;

    /// Get the number of partitions
    fn partition_count(&self) -> u32;

    /// Partition the dataset across cluster nodes
    fn partition<'a>(
        &'a mut self,
        cluster: &'a dyn DistributedCluster,
        strategy: Self::PartitionStrategy,
    ) -> BoxFuture<'a, Result<Vec<DistributedPartition<Self::Item>>>>;

    /// Get a specific partition
    fn get_partition(
        &self,
        partition_id: u32,
    ) -> BoxFuture<'_, Result<DistributedPartition<Self::Item>>>;

    /// Repartition the dataset with a new strategy
    fn repartition<'a>(
        &'a mut self,
        cluster: &'a dyn DistributedCluster,
        new_strategy: Self::PartitionStrategy,
    ) -> BoxFuture<'a, Result<()>>;

    /// Collect all partitions back to coordinator
    fn collect(&self, cluster: &dyn DistributedCluster) -> BoxFuture<'_, Result<Vec<Self::Item>>>;

    /// Get partition assignment for nodes
    fn partition_assignment(&self) -> HashMap<NodeId, Vec<u32>>;
}

/// A partition of a distributed dataset
#[derive(Debug, Clone)]
pub struct DistributedPartition<T> {
    /// Partition identifier
    pub partition_id: u32,
    /// Node holding this partition
    pub node_id: NodeId,
    /// Partition data
    pub data: Vec<T>,
    /// Partition metadata
    pub metadata: PartitionMetadata,
}

/// Metadata about a data partition
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PartitionMetadata {
    /// Number of items in partition
    pub item_count: u64,
    /// Partition size in bytes
    pub size_bytes: u64,
    /// Data schema information
    pub schema: Option<String>,
    /// Partition creation timestamp
    pub created_at: SystemTime,
    /// Last modification timestamp
    pub modified_at: SystemTime,
    /// Checksum for integrity verification
    pub checksum: String,
}

/// Partitioning strategies for distributed datasets
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum PartitioningStrategy {
    /// Split data evenly across nodes
    EvenSplit,
    /// Partition based on data hash
    HashBased(u32),
    /// Partition based on data ranges
    RangeBased,
    /// Random partitioning
    Random,
    /// Stratified partitioning (for classification)
    Stratified,
    /// Custom partitioning function
    Custom(String),
}

// =============================================================================
// Parameter Server Architecture
// =============================================================================

/// Parameter server for coordinating distributed machine learning
pub trait ParameterServer: Send + Sync {
    /// Associated type for parameters
    type Parameters: Serialize + for<'de> Deserialize<'de>;

    /// Initialize the parameter server
    fn initialize(&mut self, initial_params: Self::Parameters) -> BoxFuture<'_, Result<()>>;

    /// Get current parameters
    fn get_parameters(&self) -> BoxFuture<'_, Result<Self::Parameters>>;

    /// Update parameters with gradients
    fn update_parameters(&mut self, gradients: Vec<Self::Parameters>) -> BoxFuture<'_, Result<()>>;

    /// Push parameters to all worker nodes
    fn push_parameters(&self, cluster: &dyn DistributedCluster) -> BoxFuture<'_, Result<()>>;

    /// Pull parameters from worker nodes
    fn pull_parameters(&mut self, cluster: &dyn DistributedCluster) -> BoxFuture<'_, Result<()>>;

    /// Aggregate gradients from worker nodes
    fn aggregate_gradients(
        &mut self,
        gradients: Vec<Self::Parameters>,
    ) -> BoxFuture<'_, Result<Self::Parameters>>;

    /// Apply learning rate and optimization
    fn apply_optimization(
        &mut self,
        aggregated_gradients: Self::Parameters,
    ) -> BoxFuture<'_, Result<()>>;
}

/// Gradient aggregation strategies
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum GradientAggregation {
    /// Simple averaging
    Average,
    /// Weighted averaging by node resources
    WeightedAverage,
    /// Federated averaging with decay
    FederatedAveraging,
    /// Byzantine-robust aggregation
    ByzantineRobust,
    /// Compression-based aggregation
    Compressed,
}

// =============================================================================
// Fault Tolerance Framework
// =============================================================================

/// Comprehensive fault tolerance system for distributed training
pub trait FaultTolerance: Send + Sync {
    /// Detect when a node has failed
    fn detect_failure(
        &self,
        cluster: &dyn DistributedCluster,
    ) -> BoxFuture<'_, Result<Vec<NodeId>>>;

    /// Recover from node failures
    fn recover_from_failure(
        &mut self,
        cluster: &mut dyn DistributedCluster,
        failed_nodes: Vec<NodeId>,
    ) -> BoxFuture<'_, Result<()>>;

    /// Create a checkpoint for recovery
    fn create_checkpoint(
        &self,
        cluster: &dyn DistributedCluster,
    ) -> BoxFuture<'_, Result<FaultToleranceCheckpoint>>;

    /// Restore from a checkpoint
    fn restore_checkpoint(
        &mut self,
        cluster: &mut dyn DistributedCluster,
        checkpoint: FaultToleranceCheckpoint,
    ) -> BoxFuture<'_, Result<()>>;

    /// Replicate critical data across nodes
    fn replicate_data(
        &self,
        cluster: &dyn DistributedCluster,
        data: Vec<u8>,
    ) -> BoxFuture<'_, Result<()>>;

    /// Validate cluster integrity
    fn validate_integrity(
        &self,
        cluster: &dyn DistributedCluster,
    ) -> BoxFuture<'_, Result<IntegrityReport>>;
}

/// Checkpoint data for fault tolerance
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FaultToleranceCheckpoint {
    /// Checkpoint identifier
    pub id: String,
    /// Checkpoint timestamp
    pub timestamp: SystemTime,
    /// Training state at checkpoint
    pub training_state: Vec<u8>,
    /// Model parameters at checkpoint
    pub model_parameters: Vec<u8>,
    /// Node assignments at checkpoint
    pub node_assignments: HashMap<NodeId, Vec<u32>>,
    /// Replication information
    pub replication_map: HashMap<String, Vec<NodeId>>,
}

/// Cluster integrity validation report
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IntegrityReport {
    /// Overall integrity score
    pub integrity_score: f64,
    /// Data consistency validation
    pub data_consistency: bool,
    /// Parameter synchronization status
    pub parameter_sync: bool,
    /// Replication health
    pub replication_health: f64,
    /// Detected inconsistencies
    pub inconsistencies: Vec<String>,
    /// Recommended actions
    pub recommendations: Vec<String>,
}

// =============================================================================
// Concrete Implementations
// =============================================================================

/// Default implementation of a distributed cluster
pub struct DefaultDistributedCluster {
    /// Cluster configuration
    configuration: ClusterConfiguration,
    /// Coordinator node
    coordinator: NodeId,
    /// Active cluster nodes
    nodes: Arc<RwLock<HashMap<NodeId, Arc<dyn ClusterNode>>>>,
    /// Cluster health monitoring
    health_monitor: Arc<RwLock<ClusterHealth>>,
}

impl std::fmt::Debug for DefaultDistributedCluster {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DefaultDistributedCluster")
            .field("configuration", &self.configuration)
            .field("coordinator", &self.coordinator)
            .field("nodes", &"<HashMap<NodeId, Arc<dyn ClusterNode>>>")
            .field("health_monitor", &self.health_monitor)
            .finish()
    }
}

impl DefaultDistributedCluster {
    /// Create a new distributed cluster
    pub fn new(coordinator: NodeId, configuration: ClusterConfiguration) -> Self {
        Self {
            configuration,
            coordinator,
            nodes: Arc::new(RwLock::new(HashMap::new())),
            health_monitor: Arc::new(RwLock::new(ClusterHealth {
                overall_health: 1.0,
                healthy_nodes: 0,
                failed_nodes: 0,
                average_response_time: Duration::from_millis(10),
                total_throughput: 0.0,
                resource_utilization: ClusterResourceUtilization {
                    cpu_utilization: 0.0,
                    memory_utilization: 0.0,
                    network_utilization: 0.0,
                    storage_utilization: 0.0,
                },
            })),
        }
    }
}

impl DistributedCluster for DefaultDistributedCluster {
    fn active_nodes(&self) -> BoxFuture<'_, Result<Vec<NodeId>>> {
        Box::pin(async move {
            let nodes = self.nodes.read().map_err(|_| {
                SklearsError::InvalidOperation("Failed to acquire read lock on nodes".to_string())
            })?;
            Ok(nodes.keys().cloned().collect())
        })
    }

    fn coordinator(&self) -> &NodeId {
        &self.coordinator
    }

    fn configuration(&self) -> &ClusterConfiguration {
        &self.configuration
    }

    fn add_node(&mut self, _node_id: NodeId) -> BoxFuture<'_, Result<()>> {
        Box::pin(async move {
            // Implementation would add the node to the cluster
            // For now, this is a placeholder
            Ok(())
        })
    }

    fn remove_node(&mut self, node_id: NodeId) -> BoxFuture<'_, Result<()>> {
        Box::pin(async move {
            let mut nodes = self.nodes.write().map_err(|_| {
                SklearsError::InvalidOperation("Failed to acquire write lock on nodes".to_string())
            })?;
            nodes.remove(&node_id);
            Ok(())
        })
    }

    fn rebalance_load(&mut self) -> BoxFuture<'_, Result<()>> {
        Box::pin(async move {
            // Implementation would redistribute work based on current load
            Ok(())
        })
    }

    fn cluster_health(&self) -> BoxFuture<'_, Result<ClusterHealth>> {
        Box::pin(async move {
            let health = self.health_monitor.read().map_err(|_| {
                SklearsError::InvalidOperation(
                    "Failed to acquire read lock on health monitor".to_string(),
                )
            })?;
            Ok(health.clone())
        })
    }

    fn create_checkpoint(&self) -> BoxFuture<'_, Result<ClusterCheckpoint>> {
        Box::pin(async move {
            let checkpoint = ClusterCheckpoint {
                checkpoint_id: format!("checkpoint_{}", chrono::Utc::now().timestamp()),
                timestamp: SystemTime::now(),
                configuration: self.configuration.clone(),
                node_states: HashMap::new(), // Would collect actual node states
                cluster_state: Vec::new(),   // Would serialize cluster state
            };
            Ok(checkpoint)
        })
    }

    fn restore_checkpoint(&mut self, _checkpoint: ClusterCheckpoint) -> BoxFuture<'_, Result<()>> {
        Box::pin(async move {
            // Implementation would restore cluster state from checkpoint
            Ok(())
        })
    }
}

impl Default for ClusterConfiguration {
    fn default() -> Self {
        Self {
            max_nodes: 64,
            heartbeat_interval: Duration::from_secs(30),
            failure_timeout: Duration::from_secs(120),
            max_retries: 3,
            load_balancing: LoadBalancingStrategy::ResourceBased,
            fault_tolerance: FaultToleranceMode::CheckpointRecovery,
            consistency_level: ConsistencyLevel::Eventual,
        }
    }
}

// =============================================================================
// Example Distributed Estimator Implementation
// =============================================================================

/// Example distributed linear regression implementation
#[derive(Debug)]
pub struct DistributedLinearRegression {
    /// Model parameters (weights and bias)
    parameters: Option<Vec<f64>>,
    /// Training configuration
    config: DistributedTrainingConfig,
    /// Training progress
    progress: DistributedTrainingProgress,
}

/// Configuration for distributed training
#[derive(Debug, Clone)]
pub struct DistributedTrainingConfig {
    /// Learning rate
    pub learning_rate: f64,
    /// Number of epochs
    pub epochs: u32,
    /// Batch size per node
    pub batch_size: u32,
    /// Gradient aggregation strategy
    pub aggregation: GradientAggregation,
    /// Checkpoint frequency
    pub checkpoint_frequency: u32,
}

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

impl DistributedLinearRegression {
    /// Create a new distributed linear regression model
    pub fn new() -> Self {
        Self {
            parameters: None,
            config: DistributedTrainingConfig::default(),
            progress: DistributedTrainingProgress {
                epoch: 0,
                total_epochs: 0,
                training_loss: 0.0,
                validation_loss: None,
                samples_processed: 0,
                start_time: SystemTime::now(),
                estimated_completion: None,
                active_nodes: Vec::new(),
                node_statistics: HashMap::new(),
            },
        }
    }

    /// Configure the distributed training parameters
    pub fn with_config(mut self, config: DistributedTrainingConfig) -> Self {
        self.config = config;
        self
    }
}

impl Default for DistributedTrainingConfig {
    fn default() -> Self {
        Self {
            learning_rate: 0.01,
            epochs: 100,
            batch_size: 32,
            aggregation: GradientAggregation::Average,
            checkpoint_frequency: 10,
        }
    }
}

impl DistributedEstimator for DistributedLinearRegression {
    type TrainingData = (Vec<Vec<f64>>, Vec<f64>); // (X, y)
    type PredictionInput = Vec<Vec<f64>>;
    type PredictionOutput = Vec<f64>;
    type Parameters = Vec<f64>;

    fn fit_distributed<'a>(
        &'a mut self,
        _cluster: &'a dyn DistributedCluster,
        training_data: &Self::TrainingData,
    ) -> BoxFuture<'a, Result<()>> {
        let training_data = training_data.clone();
        Box::pin(async move {
            let (x, _y) = &training_data;

            // Initialize parameters if needed
            if self.parameters.is_none() {
                let feature_count = x.first().map(|row| row.len()).unwrap_or(0);
                self.parameters = Some(vec![0.0; feature_count + 1]); // +1 for bias
            }

            // Set up training progress
            self.progress.total_epochs = self.config.epochs;
            self.progress.start_time = SystemTime::now();
            self.progress.active_nodes = vec![]; // Simplified for now

            // Simulate distributed training process
            for epoch in 0..self.config.epochs {
                self.progress.epoch = epoch;

                // In a real implementation, this would:
                // 1. Distribute data across nodes
                // 2. Compute gradients on each node
                // 3. Aggregate gradients using parameter server
                // 4. Update parameters
                // 5. Synchronize across cluster

                // Placeholder implementation
                if let Some(ref mut params) = self.parameters {
                    // Simulate gradient descent step
                    for param in params.iter_mut() {
                        *param += self.config.learning_rate * 0.1; // Dummy gradient
                    }
                }

                // Update progress
                self.progress.samples_processed += x.len() as u64;
                self.progress.training_loss = (epoch as f64 * 0.1).exp().recip(); // Decreasing loss

                // Create checkpoint if needed
                if epoch % self.config.checkpoint_frequency == 0 {
                    // Simplified: Would create checkpoint in real implementation
                    // let _checkpoint = cluster.create_checkpoint().await?;
                }
            }

            Ok(())
        })
    }

    fn predict_distributed<'a>(
        &'a self,
        _cluster: &dyn DistributedCluster,
        input: &'a Self::PredictionInput,
    ) -> BoxFuture<'a, Result<Self::PredictionOutput>> {
        Box::pin(async move {
            let Some(ref params) = self.parameters else {
                return Err(SklearsError::InvalidOperation(
                    "Model not trained. Call fit_distributed first.".to_string(),
                ));
            };

            // Simple linear prediction: X * weights + bias
            let predictions = input
                .iter()
                .map(|features| {
                    let mut prediction = *params.last().unwrap_or(&0.0); // bias term
                    for (feature, weight) in features.iter().zip(params.iter()) {
                        prediction += feature * weight;
                    }
                    prediction
                })
                .collect();

            Ok(predictions)
        })
    }

    fn get_parameters(&self) -> Result<Self::Parameters> {
        self.parameters
            .clone()
            .ok_or_else(|| SklearsError::InvalidOperation("Model not trained".to_string()))
    }

    fn set_parameters(&mut self, params: Self::Parameters) -> Result<()> {
        self.parameters = Some(params);
        Ok(())
    }

    fn sync_parameters(&mut self, _cluster: &dyn DistributedCluster) -> BoxFuture<'_, Result<()>> {
        Box::pin(async move {
            // Implementation would synchronize parameters across all cluster nodes
            Ok(())
        })
    }

    fn training_progress(&self) -> DistributedTrainingProgress {
        self.progress.clone()
    }
}

// =============================================================================
// Distributed Dataset Implementation
// =============================================================================

/// Example distributed dataset implementation for numerical data
#[derive(Debug)]
pub struct DistributedNumericalDataset {
    /// Raw data
    data: Vec<Vec<f64>>,
    /// Current partitions
    partitions: Vec<DistributedPartition<Vec<f64>>>,
    /// Partition assignment map
    assignment: HashMap<NodeId, Vec<u32>>,
}

impl DistributedNumericalDataset {
    /// Create a new distributed numerical dataset
    pub fn new(data: Vec<Vec<f64>>) -> Self {
        Self {
            data,
            partitions: Vec::new(),
            assignment: HashMap::new(),
        }
    }

    /// Build one partition per node from pre-bucketed rows.
    ///
    /// Empty buckets are skipped, and partition identifiers are assigned
    /// sequentially so that each identifier always matches the partition's
    /// position in [`Self::partitions`] (required by [`Self::get_partition`]).
    /// The node-to-partition assignment map is updated for every partition that
    /// is created.
    fn flush_buckets_into_partitions(&mut self, nodes: &[NodeId], buckets: Vec<Vec<Vec<f64>>>) {
        for (node_id, partition_data) in nodes.iter().zip(buckets) {
            if partition_data.is_empty() {
                continue;
            }

            let partition_id = self.partitions.len() as u32;
            self.partitions.push(build_numerical_partition(
                partition_id,
                node_id,
                partition_data,
            ));
            self.assignment
                .entry(node_id.clone())
                .or_default()
                .push(partition_id);
        }
    }
}

/// Mix a byte slice into a running 64-bit FNV-1a hash state.
#[inline]
fn fnv1a_mix(hash: &mut u64, bytes: &[u8]) {
    const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
    for &byte in bytes {
        *hash ^= u64::from(byte);
        *hash = hash.wrapping_mul(FNV_PRIME);
    }
}

/// Compute a deterministic, content-derived checksum for a set of numerical
/// partition rows using the 64-bit FNV-1a hash over the IEEE-754 bit patterns.
///
/// The checksum is reproducible for identical data and changes whenever any
/// value, row length, or row count changes. The row count and per-row length
/// are mixed in so that different groupings of the same values yield different
/// checksums. Negative zero is normalized to positive zero so that
/// numerically-equal data produces an identical checksum.
fn numerical_partition_checksum(rows: &[Vec<f64>]) -> String {
    const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;

    let mut hash = FNV_OFFSET_BASIS;
    fnv1a_mix(&mut hash, &(rows.len() as u64).to_le_bytes());
    for row in rows {
        fnv1a_mix(&mut hash, &(row.len() as u64).to_le_bytes());
        for &value in row {
            let normalized = if value == 0.0 { 0.0_f64 } else { value };
            fnv1a_mix(&mut hash, &normalized.to_bits().to_le_bytes());
        }
    }
    format!("{hash:016x}")
}

/// Deterministically hash a single numerical row, perturbed by `seed`, for
/// hash-based partitioning. Equal rows with equal seeds always hash equally,
/// so the resulting node assignment is reproducible.
fn hash_numerical_row(row: &[f64], seed: u32) -> u64 {
    const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;

    let mut hash = FNV_OFFSET_BASIS;
    fnv1a_mix(&mut hash, &u64::from(seed).to_le_bytes());
    for &value in row {
        let normalized = if value == 0.0 { 0.0_f64 } else { value };
        fnv1a_mix(&mut hash, &normalized.to_bits().to_le_bytes());
    }
    hash
}

/// Construct a [`DistributedPartition`] from numerical rows, attaching a real,
/// content-derived checksum and accurate item/size metadata.
fn build_numerical_partition(
    partition_id: u32,
    node_id: &NodeId,
    partition_data: Vec<Vec<f64>>,
) -> DistributedPartition<Vec<f64>> {
    let item_count = partition_data.len() as u64;
    let size_bytes = item_count * std::mem::size_of::<f64>() as u64;
    let checksum = numerical_partition_checksum(&partition_data);
    let now = SystemTime::now();

    DistributedPartition {
        partition_id,
        node_id: node_id.clone(),
        metadata: PartitionMetadata {
            item_count,
            size_bytes,
            schema: Some("numerical_array".to_string()),
            created_at: now,
            modified_at: now,
            checksum,
        },
        data: partition_data,
    }
}

impl DistributedDataset for DistributedNumericalDataset {
    type Item = Vec<f64>;
    type PartitionStrategy = PartitioningStrategy;

    fn size(&self) -> u64 {
        self.data.len() as u64
    }

    fn partition_count(&self) -> u32 {
        self.partitions.len() as u32
    }

    fn partition<'a>(
        &'a mut self,
        cluster: &'a dyn DistributedCluster,
        strategy: Self::PartitionStrategy,
    ) -> BoxFuture<'a, Result<Vec<DistributedPartition<Self::Item>>>> {
        Box::pin(async move {
            let nodes = cluster.active_nodes().await?;
            let num_nodes = nodes.len();

            if num_nodes == 0 {
                return Err(SklearsError::InvalidOperation(
                    "No active nodes in cluster".to_string(),
                ));
            }

            self.partitions.clear();
            self.assignment.clear();

            match strategy {
                PartitioningStrategy::EvenSplit => {
                    // Contiguous, equal-sized blocks assigned by item position:
                    // node `i` receives the `i`-th block of the input order.
                    let chunk_size = self.data.len().div_ceil(num_nodes);

                    for (i, node_id) in nodes.iter().enumerate() {
                        let start = i * chunk_size;
                        let end = std::cmp::min(start + chunk_size, self.data.len());

                        if start < self.data.len() {
                            let partition_data = self.data[start..end].to_vec();
                            self.partitions.push(build_numerical_partition(
                                i as u32,
                                node_id,
                                partition_data,
                            ));
                            self.assignment
                                .entry(node_id.clone())
                                .or_default()
                                .push(i as u32);
                        }
                    }
                }
                PartitioningStrategy::RangeBased => {
                    // Order items by a representative key (the leading feature,
                    // or 0.0 for empty rows) using a total order so the result
                    // is deterministic even with NaN values, then give each node
                    // a contiguous slice of the value-ordered data. Partitions
                    // therefore correspond to disjoint value ranges rather than
                    // raw input positions.
                    let mut order: Vec<usize> = (0..self.data.len()).collect();
                    order.sort_by(|&a, &b| {
                        let key_a = self.data[a].first().copied().unwrap_or(0.0);
                        let key_b = self.data[b].first().copied().unwrap_or(0.0);
                        key_a.total_cmp(&key_b).then_with(|| a.cmp(&b))
                    });

                    let chunk_size = order.len().div_ceil(num_nodes);
                    for (i, node_id) in nodes.iter().enumerate() {
                        let start = i * chunk_size;
                        let end = std::cmp::min(start + chunk_size, order.len());

                        if start < order.len() {
                            let partition_data: Vec<Vec<f64>> = order[start..end]
                                .iter()
                                .map(|&idx| self.data[idx].clone())
                                .collect();
                            self.partitions.push(build_numerical_partition(
                                i as u32,
                                node_id,
                                partition_data,
                            ));
                            self.assignment
                                .entry(node_id.clone())
                                .or_default()
                                .push(i as u32);
                        }
                    }
                }
                PartitioningStrategy::HashBased(seed) => {
                    // Deterministically assign each item to node
                    // `hash(item, seed) % num_nodes`. Identical data and seed
                    // always produce the same assignment.
                    let mut buckets: Vec<Vec<Vec<f64>>> = vec![Vec::new(); num_nodes];
                    for row in &self.data {
                        let node_index =
                            (hash_numerical_row(row, seed) % num_nodes as u64) as usize;
                        buckets[node_index].push(row.clone());
                    }
                    self.flush_buckets_into_partitions(&nodes, buckets);
                }
                PartitioningStrategy::Random => {
                    // The cluster configuration carries no seed, so this
                    // assignment is non-deterministic across runs. Every item is
                    // still assigned to exactly one node, so the union of all
                    // partitions always reproduces the full dataset.
                    use scirs2_core::random::thread_rng;

                    let mut rng = thread_rng();
                    let mut buckets: Vec<Vec<Vec<f64>>> = vec![Vec::new(); num_nodes];
                    for row in &self.data {
                        let node_index = rng.gen_range(0..num_nodes);
                        buckets[node_index].push(row.clone());
                    }
                    self.flush_buckets_into_partitions(&nodes, buckets);
                }
                PartitioningStrategy::Stratified => {
                    // Group rows by their class label (the trailing feature) and
                    // spread each stratum round-robin across the nodes, so every
                    // partition keeps a proportional share of every class. Strata
                    // are visited in a deterministic key order, and the
                    // round-robin cursor carries across strata to keep overall
                    // load balanced.
                    use std::collections::BTreeMap;

                    let mut strata: BTreeMap<u64, Vec<usize>> = BTreeMap::new();
                    for (idx, row) in self.data.iter().enumerate() {
                        let label = row.last().copied().unwrap_or(0.0);
                        let normalized = if label == 0.0 { 0.0_f64 } else { label };
                        strata.entry(normalized.to_bits()).or_default().push(idx);
                    }

                    let mut buckets: Vec<Vec<Vec<f64>>> = vec![Vec::new(); num_nodes];
                    let mut cursor = 0_usize;
                    for indices in strata.values() {
                        for &idx in indices {
                            buckets[cursor % num_nodes].push(self.data[idx].clone());
                            cursor += 1;
                        }
                    }
                    self.flush_buckets_into_partitions(&nodes, buckets);
                }
                PartitioningStrategy::Custom(name) => match name.as_str() {
                    "round_robin" | "roundrobin" | "round-robin" => {
                        // Round-robin: assign item `i` to node `i % num_nodes`.
                        let mut buckets: Vec<Vec<Vec<f64>>> = vec![Vec::new(); num_nodes];
                        for (idx, row) in self.data.iter().enumerate() {
                            buckets[idx % num_nodes].push(row.clone());
                        }
                        self.flush_buckets_into_partitions(&nodes, buckets);
                    }
                    other => {
                        return Err(SklearsError::InvalidOperation(format!(
                            "Custom partitioning strategy '{other}' is not registered"
                        )));
                    }
                },
            }

            Ok(self.partitions.clone())
        })
    }

    fn get_partition(
        &self,
        partition_id: u32,
    ) -> BoxFuture<'_, Result<DistributedPartition<Self::Item>>> {
        Box::pin(async move {
            self.partitions
                .get(partition_id as usize)
                .cloned()
                .ok_or_else(|| {
                    SklearsError::InvalidOperation(format!("Partition {} not found", partition_id))
                })
        })
    }

    fn repartition<'a>(
        &'a mut self,
        cluster: &'a dyn DistributedCluster,
        new_strategy: Self::PartitionStrategy,
    ) -> BoxFuture<'a, Result<()>> {
        Box::pin(async move {
            // Collect all data back first
            let collected_data = self.collect(cluster).await?;
            self.data = collected_data;

            // Repartition with new strategy
            self.partition(cluster, new_strategy).await?;

            Ok(())
        })
    }

    fn collect(&self, _cluster: &dyn DistributedCluster) -> BoxFuture<'_, Result<Vec<Self::Item>>> {
        Box::pin(async move {
            let mut collected = Vec::new();
            for partition in &self.partitions {
                collected.extend(partition.data.clone());
            }
            Ok(collected)
        })
    }

    fn partition_assignment(&self) -> HashMap<NodeId, Vec<u32>> {
        self.assignment.clone()
    }
}

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

    #[test]
    fn test_node_id_creation() {
        let node_id = NodeId::new("worker-01");
        assert_eq!(node_id.as_str(), "worker-01");
        assert_eq!(node_id.to_string(), "worker-01");
    }

    #[test]
    fn test_message_priority_ordering() {
        assert!(MessagePriority::Critical > MessagePriority::High);
        assert!(MessagePriority::High > MessagePriority::Normal);
        assert!(MessagePriority::Normal > MessagePriority::Low);
    }

    #[test]
    fn test_cluster_configuration_default() {
        let config = ClusterConfiguration::default();
        assert_eq!(config.max_nodes, 64);
        assert_eq!(config.load_balancing, LoadBalancingStrategy::ResourceBased);
        assert_eq!(
            config.fault_tolerance,
            FaultToleranceMode::CheckpointRecovery
        );
    }

    #[test]
    fn test_distributed_linear_regression_creation() {
        let model = DistributedLinearRegression::new();
        assert!(model.parameters.is_none());
        assert_eq!(model.progress.epoch, 0);
    }

    #[test]
    fn test_distributed_dataset_size() {
        let data = vec![vec![1.0, 2.0], vec![3.0, 4.0], vec![5.0, 6.0]];
        let dataset = DistributedNumericalDataset::new(data);
        assert_eq!(dataset.size(), 3);
        assert_eq!(dataset.partition_count(), 0); // No partitions initially
    }

    #[test]
    fn test_message_type_serialization() {
        let msg_type = MessageType::ParameterSync;
        let serialized = serde_json::to_string(&msg_type).unwrap_or_default();
        let deserialized: MessageType =
            serde_json::from_str(&serialized).expect("valid JSON operation");
        assert_eq!(msg_type, deserialized);
    }

    #[test]
    fn test_partitioning_strategy_variants() {
        let strategies = vec![
            PartitioningStrategy::EvenSplit,
            PartitioningStrategy::HashBased(4),
            PartitioningStrategy::RangeBased,
            PartitioningStrategy::Random,
            PartitioningStrategy::Stratified,
            PartitioningStrategy::Custom("custom_strategy".to_string()),
        ];

        for strategy in strategies {
            let serialized = serde_json::to_string(&strategy).unwrap_or_default();
            let _deserialized: PartitioningStrategy =
                serde_json::from_str(&serialized).expect("valid JSON operation");
        }
    }

    #[test]
    fn test_distributed_training_config() {
        let config = DistributedTrainingConfig::default();
        assert_eq!(config.learning_rate, 0.01);
        assert_eq!(config.epochs, 100);
        assert_eq!(config.batch_size, 32);
    }

    #[cfg(feature = "async_support")]
    #[tokio::test]
    async fn test_default_cluster_operations() {
        let coordinator = NodeId::new("coordinator");
        let config = ClusterConfiguration::default();
        let cluster = DefaultDistributedCluster::new(coordinator.clone(), config);

        assert_eq!(cluster.coordinator(), &coordinator);

        let nodes = cluster.active_nodes().await.expect("expected valid value");
        assert!(nodes.is_empty()); // No nodes initially

        let health = cluster
            .cluster_health()
            .await
            .expect("expected valid value");
        assert_eq!(health.overall_health, 1.0);
    }

    #[test]
    fn test_numerical_partition_checksum_is_deterministic() {
        let data = vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]];
        let first = numerical_partition_checksum(&data);
        let second = numerical_partition_checksum(&data);
        assert_eq!(
            first, second,
            "identical data must produce identical checksums"
        );
    }

    #[test]
    fn test_numerical_partition_checksum_changes_with_data() {
        let data_a = vec![vec![1.0, 2.0, 3.0]];
        let data_b = vec![vec![1.0, 2.0, 4.0]];
        assert_ne!(
            numerical_partition_checksum(&data_a),
            numerical_partition_checksum(&data_b),
            "different data must produce different checksums"
        );
    }

    #[test]
    fn test_numerical_partition_checksum_respects_row_structure() {
        // The same values grouped into different rows must hash differently.
        let grouped_a = vec![vec![1.0, 2.0], vec![3.0]];
        let grouped_b = vec![vec![1.0], vec![2.0, 3.0]];
        assert_ne!(
            numerical_partition_checksum(&grouped_a),
            numerical_partition_checksum(&grouped_b)
        );
    }

    #[test]
    fn test_numerical_partition_checksum_is_not_index_placeholder() {
        // Guard against regressing to the fabricated `checksum_{index}` form.
        let checksum = numerical_partition_checksum(&[vec![1.0, 2.0]]);
        assert!(!checksum.starts_with("checksum_"));
        assert_eq!(checksum.len(), 16);
    }

    #[test]
    fn test_hash_numerical_row_determinism_and_seed_sensitivity() {
        let row = vec![1.5, -2.25, 3.0];
        assert_eq!(hash_numerical_row(&row, 7), hash_numerical_row(&row, 7));
        assert_ne!(hash_numerical_row(&row, 7), hash_numerical_row(&row, 8));
    }

    #[cfg(feature = "async_support")]
    struct TestCluster {
        coordinator: NodeId,
        configuration: ClusterConfiguration,
        nodes: Vec<NodeId>,
    }

    #[cfg(feature = "async_support")]
    impl TestCluster {
        fn with_nodes(count: usize) -> Self {
            let nodes = (0..count)
                .map(|i| NodeId::new(format!("node-{i}")))
                .collect();
            Self {
                coordinator: NodeId::new("coordinator"),
                configuration: ClusterConfiguration::default(),
                nodes,
            }
        }
    }

    #[cfg(feature = "async_support")]
    impl DistributedCluster for TestCluster {
        fn active_nodes(&self) -> BoxFuture<'_, Result<Vec<NodeId>>> {
            let nodes = self.nodes.clone();
            Box::pin(async move { Ok(nodes) })
        }

        fn coordinator(&self) -> &NodeId {
            &self.coordinator
        }

        fn configuration(&self) -> &ClusterConfiguration {
            &self.configuration
        }

        fn add_node(&mut self, node: NodeId) -> BoxFuture<'_, Result<()>> {
            self.nodes.push(node);
            Box::pin(async move { Ok(()) })
        }

        fn remove_node(&mut self, node: NodeId) -> BoxFuture<'_, Result<()>> {
            self.nodes.retain(|existing| existing != &node);
            Box::pin(async move { Ok(()) })
        }

        fn rebalance_load(&mut self) -> BoxFuture<'_, Result<()>> {
            Box::pin(async move { Ok(()) })
        }

        fn cluster_health(&self) -> BoxFuture<'_, Result<ClusterHealth>> {
            Box::pin(async move {
                Err(SklearsError::InvalidOperation(
                    "cluster_health is unused in partition tests".to_string(),
                ))
            })
        }

        fn create_checkpoint(&self) -> BoxFuture<'_, Result<ClusterCheckpoint>> {
            Box::pin(async move {
                Err(SklearsError::InvalidOperation(
                    "create_checkpoint is unused in partition tests".to_string(),
                ))
            })
        }

        fn restore_checkpoint(
            &mut self,
            _checkpoint: ClusterCheckpoint,
        ) -> BoxFuture<'_, Result<()>> {
            Box::pin(async move { Ok(()) })
        }
    }

    #[cfg(feature = "async_support")]
    fn parse_node_index(node_id: &NodeId) -> usize {
        node_id
            .as_str()
            .strip_prefix("node-")
            .and_then(|suffix| suffix.parse().ok())
            .expect("test node ids follow the node-<index> convention")
    }

    #[cfg(feature = "async_support")]
    #[tokio::test]
    async fn test_partition_without_nodes_errors() {
        let mut dataset = DistributedNumericalDataset::new(vec![vec![1.0]]);
        let cluster = TestCluster::with_nodes(0);
        let result = dataset
            .partition(&cluster, PartitioningStrategy::EvenSplit)
            .await;
        assert!(result.is_err(), "partitioning with no nodes must error");
    }

    #[cfg(feature = "async_support")]
    #[tokio::test]
    async fn test_even_split_covers_all_items() {
        let data: Vec<Vec<f64>> = (0..10).map(|i| vec![i as f64, (i * 2) as f64]).collect();
        let mut dataset = DistributedNumericalDataset::new(data.clone());
        let cluster = TestCluster::with_nodes(3);

        let partitions = dataset
            .partition(&cluster, PartitioningStrategy::EvenSplit)
            .await
            .expect("even split must succeed");

        let mut collected: Vec<Vec<f64>> = partitions.iter().flat_map(|p| p.data.clone()).collect();
        collected.sort_by(|a, b| a[0].total_cmp(&b[0]));
        assert_eq!(collected, data, "union of partitions must cover all items");
    }

    #[cfg(feature = "async_support")]
    #[tokio::test]
    async fn test_round_robin_custom_strategy_assignment() {
        let data: Vec<Vec<f64>> = (0..9).map(|i| vec![i as f64]).collect();
        let mut dataset = DistributedNumericalDataset::new(data);
        let cluster = TestCluster::with_nodes(3);

        let partitions = dataset
            .partition(
                &cluster,
                PartitioningStrategy::Custom("round_robin".to_string()),
            )
            .await
            .expect("round-robin custom strategy must succeed");

        assert_eq!(partitions.len(), 3);
        // Item `i` must land on node `i % num_nodes`: every value in a node's
        // partition shares the same `value % num_nodes` class.
        for partition in &partitions {
            let node_index = parse_node_index(&partition.node_id);
            for row in &partition.data {
                assert_eq!((row[0] as usize) % 3, node_index);
            }
            assert_eq!(partition.data.len(), 3);
        }
    }

    #[cfg(feature = "async_support")]
    #[tokio::test]
    async fn test_hash_based_partitioning_is_deterministic() {
        let data: Vec<Vec<f64>> = (0..20).map(|i| vec![i as f64, (i % 4) as f64]).collect();
        let cluster = TestCluster::with_nodes(4);

        let mut first_dataset = DistributedNumericalDataset::new(data.clone());
        let mut second_dataset = DistributedNumericalDataset::new(data.clone());

        let first = first_dataset
            .partition(&cluster, PartitioningStrategy::HashBased(13))
            .await
            .expect("hash partitioning must succeed");
        let second = second_dataset
            .partition(&cluster, PartitioningStrategy::HashBased(13))
            .await
            .expect("hash partitioning must succeed");

        let summarize = |partitions: &[DistributedPartition<Vec<f64>>]| {
            partitions
                .iter()
                .map(|p| (p.partition_id, p.data.clone()))
                .collect::<Vec<_>>()
        };
        assert_eq!(
            summarize(&first),
            summarize(&second),
            "hash partitioning must be deterministic for identical data and seed"
        );

        // Every item must reside on node `hash(item, seed) % num_nodes`.
        for partition in &first {
            let node_index = parse_node_index(&partition.node_id);
            for row in &partition.data {
                assert_eq!((hash_numerical_row(row, 13) % 4) as usize, node_index);
            }
        }

        // Union of partitions reproduces the entire dataset.
        let mut all: Vec<Vec<f64>> = first.iter().flat_map(|p| p.data.clone()).collect();
        all.sort_by(|a, b| a[0].total_cmp(&b[0]));
        assert_eq!(all, data);
    }

    #[cfg(feature = "async_support")]
    #[tokio::test]
    async fn test_range_based_partitioning_orders_by_value() {
        let data: Vec<Vec<f64>> = vec![
            vec![9.0],
            vec![1.0],
            vec![7.0],
            vec![3.0],
            vec![5.0],
            vec![2.0],
        ];
        let mut dataset = DistributedNumericalDataset::new(data.clone());
        let cluster = TestCluster::with_nodes(3);

        let partitions = dataset
            .partition(&cluster, PartitioningStrategy::RangeBased)
            .await
            .expect("range partitioning must succeed");

        // Flattening in partition order must yield globally value-ordered keys.
        let ordered_keys: Vec<f64> = partitions
            .iter()
            .flat_map(|p| p.data.iter().map(|row| row[0]))
            .collect();
        let mut expected: Vec<f64> = data.iter().map(|row| row[0]).collect();
        expected.sort_by(|a, b| a.total_cmp(b));
        assert_eq!(
            ordered_keys, expected,
            "range partitions must be value-ordered"
        );

        // 6 items across 3 nodes => 2 contiguous items each.
        assert_eq!(partitions.len(), 3);
        for partition in &partitions {
            assert_eq!(partition.data.len(), 2);
        }
    }

    #[cfg(feature = "async_support")]
    #[tokio::test]
    async fn test_stratified_partitioning_balances_classes() {
        // The trailing column is the class label; 3 classes with 6 items each.
        let mut data: Vec<Vec<f64>> = Vec::new();
        for class in 0..3 {
            for k in 0..6 {
                data.push(vec![k as f64, class as f64]);
            }
        }
        let mut dataset = DistributedNumericalDataset::new(data.clone());
        let cluster = TestCluster::with_nodes(3);

        let partitions = dataset
            .partition(&cluster, PartitioningStrategy::Stratified)
            .await
            .expect("stratified partitioning must succeed");

        let total: usize = partitions.iter().map(|p| p.data.len()).sum();
        assert_eq!(total, data.len(), "stratified must not drop items");

        // 18 items / 3 nodes => 6 per node; 3 classes => 2 of each class per node.
        for partition in &partitions {
            let mut class_counts: HashMap<i64, usize> = HashMap::new();
            for row in &partition.data {
                *class_counts.entry(row[1] as i64).or_insert(0) += 1;
            }
            assert_eq!(
                class_counts.len(),
                3,
                "each partition must hold all classes"
            );
            for count in class_counts.values() {
                assert_eq!(*count, 2, "each class must be evenly represented");
            }
        }
    }

    #[cfg(feature = "async_support")]
    #[tokio::test]
    async fn test_random_partitioning_covers_all_items() {
        let data: Vec<Vec<f64>> = (0..30).map(|i| vec![i as f64]).collect();
        let mut dataset = DistributedNumericalDataset::new(data);
        let cluster = TestCluster::with_nodes(4);

        let partitions = dataset
            .partition(&cluster, PartitioningStrategy::Random)
            .await
            .expect("random partitioning must succeed");

        let mut all: Vec<f64> = partitions
            .iter()
            .flat_map(|p| p.data.iter().map(|row| row[0]))
            .collect();
        all.sort_by(|a, b| a.total_cmp(b));
        let expected: Vec<f64> = (0..30).map(|i| i as f64).collect();
        assert_eq!(
            all, expected,
            "random partitioning must not lose or duplicate items"
        );
    }

    #[cfg(feature = "async_support")]
    #[tokio::test]
    async fn test_unknown_custom_strategy_errors() {
        let mut dataset = DistributedNumericalDataset::new(vec![vec![1.0], vec![2.0]]);
        let cluster = TestCluster::with_nodes(2);

        let result = dataset
            .partition(
                &cluster,
                PartitioningStrategy::Custom("totally_unknown".to_string()),
            )
            .await;
        assert!(
            result.is_err(),
            "unknown custom strategy must return an honest error"
        );
    }

    #[cfg(feature = "async_support")]
    #[tokio::test]
    async fn test_partition_checksums_reflect_content() {
        let cluster = TestCluster::with_nodes(1);

        let mut dataset_a = DistributedNumericalDataset::new(vec![vec![1.0, 2.0], vec![3.0, 4.0]]);
        let mut dataset_b = DistributedNumericalDataset::new(vec![vec![1.0, 2.0], vec![3.0, 5.0]]);
        let mut dataset_a_again =
            DistributedNumericalDataset::new(vec![vec![1.0, 2.0], vec![3.0, 4.0]]);

        let partitions_a = dataset_a
            .partition(&cluster, PartitioningStrategy::EvenSplit)
            .await
            .expect("partition a");
        let partitions_b = dataset_b
            .partition(&cluster, PartitioningStrategy::EvenSplit)
            .await
            .expect("partition b");
        let partitions_a_again = dataset_a_again
            .partition(&cluster, PartitioningStrategy::EvenSplit)
            .await
            .expect("partition a again");

        assert_eq!(partitions_a.len(), 1);
        assert_eq!(partitions_b.len(), 1);
        assert_ne!(
            partitions_a[0].metadata.checksum, partitions_b[0].metadata.checksum,
            "different data must yield different checksums"
        );
        assert_eq!(
            partitions_a[0].metadata.checksum, partitions_a_again[0].metadata.checksum,
            "identical data must yield identical checksums"
        );
        assert!(!partitions_a[0].metadata.checksum.starts_with("checksum_"));
    }
}