//! Model Sharding Utilities
//!
//! This module provides automatic model partitioning across devices:
//! - Balanced memory and compute distribution
//! - Layer-wise and tensor-wise sharding strategies
//! - Automatic shard placement optimization
//! - Cross-shard communication management
//! - Support for various model architectures
//!
//! Model sharding enables training models larger than single GPU memory
//! by distributing parameters, gradients, and optimizer states.

use crate::error::Result;
use parking_lot::{Mutex, RwLock};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

// =============================================================================
// Configuration
// =============================================================================

/// Sharding strategy for model partitioning
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ShardingStrategy {
    /// Shard by layers (pipeline-style)
    LayerWise,
    /// Shard tensors across devices (tensor parallelism)
    TensorWise,
    /// Shard by rows (for embeddings, linear layers)
    RowWise,
    /// Shard by columns (for linear layers)
    ColumnWise,
    /// Replicate across all devices
    Replicated,
    /// Custom sharding with explicit placement
    Custom,
}

/// Dimension along which to shard
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ShardDimension {
    /// No sharding (replicated)
    None,
    /// Shard along first dimension
    Dim0,
    /// Shard along second dimension
    Dim1,
    /// Shard along last dimension
    DimLast,
    /// Shard along specific dimension
    Dim(usize),
}

/// Configuration for model sharding
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelShardingConfig {
    /// Number of devices to shard across
    pub num_devices: usize,
    /// Default sharding strategy
    pub default_strategy: ShardingStrategy,
    /// Memory limit per device (bytes, 0 = unlimited)
    pub memory_limit_per_device: usize,
    /// Balance memory across devices
    pub balance_memory: bool,
    /// Balance compute across devices
    pub balance_compute: bool,
    /// Minimum shard size (bytes)
    pub min_shard_size: usize,
    /// Enable automatic resharding on imbalance
    pub auto_reshard: bool,
    /// Threshold for resharding (imbalance ratio)
    pub reshard_threshold: f32,
}

impl Default for ModelShardingConfig {
    fn default() -> Self {
        Self {
            num_devices: 1,
            default_strategy: ShardingStrategy::TensorWise,
            memory_limit_per_device: 0,
            balance_memory: true,
            balance_compute: true,
            min_shard_size: 1024,
            auto_reshard: false,
            reshard_threshold: 0.2,
        }
    }
}

// =============================================================================
// Parameter Info
// =============================================================================

/// Information about a model parameter
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ParameterInfo {
    /// Parameter name
    pub name: String,
    /// Shape of the parameter
    pub shape: Vec<usize>,
    /// Data type (e.g., "float32", "float16")
    pub dtype: String,
    /// Size in bytes
    pub size_bytes: usize,
    /// Estimated compute cost (FLOPs for forward pass)
    pub compute_cost: u64,
    /// Layer index (for layer-wise sharding)
    pub layer_index: Option<usize>,
    /// Parameter type (weight, bias, etc.)
    pub param_type: ParameterType,
    /// Whether parameter requires gradient
    pub requires_grad: bool,
}

impl ParameterInfo {
    /// Create a new parameter info
    pub fn new(name: String, shape: Vec<usize>, dtype: String) -> Self {
        let elem_size = match dtype.as_str() {
            "float32" | "f32" => 4,
            "float16" | "f16" | "half" => 2,
            "bfloat16" | "bf16" => 2,
            "float64" | "f64" => 8,
            "int32" | "i32" => 4,
            "int64" | "i64" => 8,
            _ => 4,
        };
        let num_elements: usize = shape.iter().product();
        let size_bytes = num_elements * elem_size;

        Self {
            name,
            shape,
            dtype,
            size_bytes,
            compute_cost: num_elements as u64 * 2, // Simple estimate
            layer_index: None,
            param_type: ParameterType::Weight,
            requires_grad: true,
        }
    }

    /// Get number of elements
    pub fn num_elements(&self) -> usize {
        self.shape.iter().product()
    }

    /// Check if parameter can be sharded along a dimension
    pub fn can_shard_dim(&self, dim: usize) -> bool {
        dim < self.shape.len() && self.shape[dim] > 1
    }
}

/// Type of parameter
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ParameterType {
    /// Weight matrix
    Weight,
    /// Bias vector
    Bias,
    /// Embedding table
    Embedding,
    /// Normalization scale
    Scale,
    /// Normalization shift
    Shift,
    /// Other parameter
    Other,
}

// =============================================================================
// Shard Specification
// =============================================================================

/// Specification for how a parameter is sharded
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ShardSpec {
    /// Parameter name
    pub param_name: String,
    /// Sharding strategy
    pub strategy: ShardingStrategy,
    /// Dimension to shard along
    pub shard_dim: ShardDimension,
    /// Device placements for each shard
    pub placements: Vec<DevicePlacement>,
    /// Original shape
    pub original_shape: Vec<usize>,
    /// Shard shapes
    pub shard_shapes: Vec<Vec<usize>>,
}

impl ShardSpec {
    /// Create a replicated shard spec
    pub fn replicated(param_name: String, shape: Vec<usize>, devices: Vec<i32>) -> Self {
        let placements = devices.iter().map(|&d| DevicePlacement::new(d)).collect();
        let shard_shapes = vec![shape.clone(); devices.len()];

        Self {
            param_name,
            strategy: ShardingStrategy::Replicated,
            shard_dim: ShardDimension::None,
            placements,
            original_shape: shape,
            shard_shapes,
        }
    }

    /// Create a tensor-wise sharded spec
    pub fn tensor_sharded(
        param_name: String,
        shape: Vec<usize>,
        shard_dim: usize,
        devices: Vec<i32>,
    ) -> Self {
        let num_shards = devices.len();
        let placements = devices.iter().map(|&d| DevicePlacement::new(d)).collect();

        // Calculate shard shapes
        let mut shard_shapes = Vec::new();
        let dim_size = shape.get(shard_dim).copied().unwrap_or(1);
        let base_size = dim_size / num_shards;
        let remainder = dim_size % num_shards;

        for i in 0..num_shards {
            let mut shard_shape = shape.clone();
            let shard_size = base_size + usize::from(i < remainder);
            if shard_dim < shard_shape.len() {
                shard_shape[shard_dim] = shard_size;
            }
            shard_shapes.push(shard_shape);
        }

        Self {
            param_name,
            strategy: ShardingStrategy::TensorWise,
            shard_dim: ShardDimension::Dim(shard_dim),
            placements,
            original_shape: shape,
            shard_shapes,
        }
    }

    /// Get number of shards
    pub fn num_shards(&self) -> usize {
        self.placements.len()
    }

    /// Get total memory for all shards
    pub fn total_memory(&self) -> usize {
        self.shard_shapes
            .iter()
            .map(|s| s.iter().product::<usize>() * 4) // Assume float32
            .sum()
    }
}

/// Device placement for a shard
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DevicePlacement {
    /// Device ID
    pub device_id: i32,
    /// Memory offset on device
    pub memory_offset: usize,
    /// Priority (lower = preferred)
    pub priority: i32,
}

impl DevicePlacement {
    /// Create a new device placement
    pub fn new(device_id: i32) -> Self {
        Self {
            device_id,
            memory_offset: 0,
            priority: 0,
        }
    }
}

// =============================================================================
// Sharding Plan
// =============================================================================

/// A complete sharding plan for a model
#[derive(Debug, Clone)]
pub struct ShardingPlan {
    /// Shard specifications per parameter
    pub specs: HashMap<String, ShardSpec>,
    /// Device memory usage
    pub device_memory: HashMap<i32, usize>,
    /// Device compute load
    pub device_compute: HashMap<i32, u64>,
    /// Total model size
    pub total_size: usize,
    /// Total compute cost
    pub total_compute: u64,
    /// Communication cost estimate
    pub comm_cost: u64,
}

impl ShardingPlan {
    /// Create a new empty sharding plan
    pub fn new() -> Self {
        Self {
            specs: HashMap::new(),
            device_memory: HashMap::new(),
            device_compute: HashMap::new(),
            total_size: 0,
            total_compute: 0,
            comm_cost: 0,
        }
    }

    /// Add a shard spec to the plan
    pub fn add_spec(&mut self, spec: ShardSpec) {
        let param_name = spec.param_name.clone();

        // Update device memory usage
        for (i, placement) in spec.placements.iter().enumerate() {
            let shard_size = spec.shard_shapes[i].iter().product::<usize>() * 4;
            *self.device_memory.entry(placement.device_id).or_insert(0) += shard_size;
        }

        self.specs.insert(param_name, spec);
    }

    /// Get memory balance ratio (0 = perfect, 1 = all on one device)
    pub fn memory_balance(&self) -> f32 {
        if self.device_memory.is_empty() {
            return 0.0;
        }

        let values: Vec<_> = self.device_memory.values().copied().collect();
        let max = *values.iter().max().unwrap_or(&0) as f32;
        let min = *values.iter().min().unwrap_or(&0) as f32;
        let avg = values.iter().sum::<usize>() as f32 / values.len() as f32;

        if avg == 0.0 { 0.0 } else { (max - min) / avg }
    }

    /// Get compute balance ratio
    pub fn compute_balance(&self) -> f32 {
        if self.device_compute.is_empty() {
            return 0.0;
        }

        let values: Vec<_> = self.device_compute.values().copied().collect();
        let max = *values.iter().max().unwrap_or(&0) as f32;
        let min = *values.iter().min().unwrap_or(&0) as f32;
        let avg = values.iter().sum::<u64>() as f32 / values.len() as f32;

        if avg == 0.0 { 0.0 } else { (max - min) / avg }
    }

    /// Get number of parameters
    pub fn num_parameters(&self) -> usize {
        self.specs.len()
    }

    /// Get spec for a parameter
    pub fn get_spec(&self, param_name: &str) -> Option<&ShardSpec> {
        self.specs.get(param_name)
    }
}

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

// =============================================================================
// Model Sharder
// =============================================================================

/// Automatic model sharding utility
pub struct ModelSharder {
    /// Configuration
    config: ModelShardingConfig,
    /// Available devices
    devices: Vec<i32>,
    /// Device memory capacities
    device_capacities: HashMap<i32, usize>,
    /// Current sharding plan
    plan: RwLock<ShardingPlan>,
    /// Parameter registry
    parameters: RwLock<HashMap<String, ParameterInfo>>,
    /// Statistics
    stats: RwLock<ShardingStats>,
}

impl ModelSharder {
    /// Create a new model sharder
    pub fn new(config: ModelShardingConfig, devices: Vec<i32>) -> Self {
        let device_capacities = devices
            .iter()
            .map(|&d| (d, config.memory_limit_per_device))
            .collect();

        Self {
            config,
            devices,
            device_capacities,
            plan: RwLock::new(ShardingPlan::new()),
            parameters: RwLock::new(HashMap::new()),
            stats: RwLock::new(ShardingStats::default()),
        }
    }

    /// Register a parameter for sharding
    pub fn register_parameter(&self, info: ParameterInfo) {
        let mut params = self.parameters.write();
        params.insert(info.name.clone(), info);
    }

    /// Register multiple parameters
    pub fn register_parameters(&self, infos: Vec<ParameterInfo>) {
        let mut params = self.parameters.write();
        for info in infos {
            params.insert(info.name.clone(), info);
        }
    }

    /// Generate a sharding plan for all registered parameters
    pub fn generate_plan(&self) -> Result<ShardingPlan> {
        let params = self.parameters.read();
        let mut plan = ShardingPlan::new();

        // Sort parameters by size for better load balancing
        let mut sorted_params: Vec<_> = params.values().collect();
        sorted_params.sort_by(|a, b| b.size_bytes.cmp(&a.size_bytes));

        for param in sorted_params {
            let spec = self.create_shard_spec(param)?;
            plan.add_spec(spec);
        }

        // Calculate totals
        plan.total_size = params.values().map(|p| p.size_bytes).sum();
        plan.total_compute = params.values().map(|p| p.compute_cost).sum();

        // Estimate communication cost
        plan.comm_cost = self.estimate_comm_cost(&plan);

        // Update stats
        {
            let mut stats = self.stats.write();
            stats.plans_generated += 1;
        }

        *self.plan.write() = plan.clone();
        Ok(plan)
    }

    /// Create shard spec for a parameter
    fn create_shard_spec(&self, param: &ParameterInfo) -> Result<ShardSpec> {
        match self.config.default_strategy {
            ShardingStrategy::Replicated => Ok(ShardSpec::replicated(
                param.name.clone(),
                param.shape.clone(),
                self.devices.clone(),
            )),
            ShardingStrategy::TensorWise | ShardingStrategy::RowWise => {
                // Shard along first dimension
                let shard_dim = 0;
                if param.can_shard_dim(shard_dim) && param.shape[shard_dim] >= self.devices.len() {
                    Ok(ShardSpec::tensor_sharded(
                        param.name.clone(),
                        param.shape.clone(),
                        shard_dim,
                        self.devices.clone(),
                    ))
                } else {
                    // Fall back to replication for small tensors
                    Ok(ShardSpec::replicated(
                        param.name.clone(),
                        param.shape.clone(),
                        self.devices.clone(),
                    ))
                }
            }
            ShardingStrategy::ColumnWise => {
                // Shard along last dimension
                let shard_dim = param.shape.len().saturating_sub(1);
                if param.can_shard_dim(shard_dim) && param.shape[shard_dim] >= self.devices.len() {
                    Ok(ShardSpec::tensor_sharded(
                        param.name.clone(),
                        param.shape.clone(),
                        shard_dim,
                        self.devices.clone(),
                    ))
                } else {
                    Ok(ShardSpec::replicated(
                        param.name.clone(),
                        param.shape.clone(),
                        self.devices.clone(),
                    ))
                }
            }
            ShardingStrategy::LayerWise => {
                // Assign entire layer to one device based on layer index
                let device_idx = param.layer_index.unwrap_or(0) % self.devices.len();
                Ok(ShardSpec::replicated(
                    param.name.clone(),
                    param.shape.clone(),
                    vec![self.devices[device_idx]],
                ))
            }
            ShardingStrategy::Custom => {
                // Use default tensor-wise for custom
                self.create_shard_spec(&ParameterInfo {
                    param_type: param.param_type,
                    ..param.clone()
                })
            }
        }
    }

    /// Estimate communication cost for a plan
    fn estimate_comm_cost(&self, plan: &ShardingPlan) -> u64 {
        let mut cost = 0u64;

        for spec in plan.specs.values() {
            match spec.strategy {
                ShardingStrategy::TensorWise
                | ShardingStrategy::RowWise
                | ShardingStrategy::ColumnWise => {
                    // AllGather or ReduceScatter needed
                    let shard_size: usize = spec
                        .shard_shapes
                        .iter()
                        .map(|s| s.iter().product::<usize>())
                        .sum();
                    cost += shard_size as u64 * 2; // Forward + backward
                }
                ShardingStrategy::Replicated => {
                    // AllReduce needed for gradients
                    let size: usize = spec.original_shape.iter().product();
                    cost += size as u64;
                }
                _ => {}
            }
        }

        cost
    }

    /// Get current sharding plan
    pub fn current_plan(&self) -> ShardingPlan {
        self.plan.read().clone()
    }

    /// Get parameter info
    pub fn get_parameter(&self, name: &str) -> Option<ParameterInfo> {
        self.parameters.read().get(name).cloned()
    }

    /// Get all parameter names
    pub fn parameter_names(&self) -> Vec<String> {
        self.parameters.read().keys().cloned().collect()
    }

    /// Get statistics
    pub fn stats(&self) -> ShardingStats {
        self.stats.read().clone()
    }

    /// Get number of devices
    pub fn num_devices(&self) -> usize {
        self.devices.len()
    }

    /// Get device list
    pub fn devices(&self) -> Vec<i32> {
        self.devices.clone()
    }
}

// =============================================================================
// Sharding Statistics
// =============================================================================

/// Statistics for model sharding
#[derive(Debug, Default, Clone)]
pub struct ShardingStats {
    /// Number of plans generated
    pub plans_generated: usize,
    /// Number of resharding operations
    pub reshard_count: usize,
    /// Total parameters sharded
    pub parameters_sharded: usize,
    /// Total bytes sharded
    pub bytes_sharded: usize,
}

// =============================================================================
// Shard Manager
// =============================================================================

/// Manages sharded parameter storage and communication
pub struct ShardManager {
    /// Sharder reference
    sharder: Arc<ModelSharder>,
    /// Shard data storage (device_id -> param_name -> data)
    shard_data: RwLock<HashMap<i32, HashMap<String, Vec<u8>>>>,
    /// Communication pending
    pending_comms: Mutex<Vec<PendingComm>>,
    /// Next comm ID
    next_comm_id: AtomicU64,
}

/// A pending communication operation
#[derive(Debug, Clone)]
pub struct PendingComm {
    /// Communication ID
    pub id: u64,
    /// Type of communication
    pub comm_type: CommType,
    /// Source devices
    pub src_devices: Vec<i32>,
    /// Destination devices
    pub dst_devices: Vec<i32>,
    /// Parameter name
    pub param_name: String,
    /// Size in bytes
    pub size_bytes: usize,
}

/// Type of communication
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CommType {
    /// AllGather shards
    AllGather,
    /// ReduceScatter gradients
    ReduceScatter,
    /// AllReduce gradients
    AllReduce,
    /// Point-to-point send
    Send,
    /// Point-to-point receive
    Recv,
}

impl ShardManager {
    /// Create a new shard manager
    pub fn new(sharder: Arc<ModelSharder>) -> Self {
        Self {
            sharder,
            shard_data: RwLock::new(HashMap::new()),
            pending_comms: Mutex::new(Vec::new()),
            next_comm_id: AtomicU64::new(1),
        }
    }

    /// Store shard data
    pub fn store_shard(&self, device_id: i32, param_name: String, data: Vec<u8>) {
        let mut storage = self.shard_data.write();
        storage
            .entry(device_id)
            .or_default()
            .insert(param_name, data);
    }

    /// Get shard data
    pub fn get_shard(&self, device_id: i32, param_name: &str) -> Option<Vec<u8>> {
        let storage = self.shard_data.read();
        storage.get(&device_id)?.get(param_name).cloned()
    }

    /// Schedule an AllGather operation
    pub fn schedule_all_gather(&self, param_name: String) -> u64 {
        let plan = self.sharder.current_plan();
        let spec = plan.get_spec(&param_name);

        let devices: Vec<i32> = spec
            .map(|s| s.placements.iter().map(|p| p.device_id).collect())
            .unwrap_or_default();

        let size = spec.map_or(0, ShardSpec::total_memory);

        let id = self.next_comm_id.fetch_add(1, Ordering::SeqCst);
        let comm = PendingComm {
            id,
            comm_type: CommType::AllGather,
            src_devices: devices.clone(),
            dst_devices: devices,
            param_name,
            size_bytes: size,
        };

        self.pending_comms.lock().push(comm);
        id
    }

    /// Schedule a ReduceScatter operation
    pub fn schedule_reduce_scatter(&self, param_name: String) -> u64 {
        let plan = self.sharder.current_plan();
        let spec = plan.get_spec(&param_name);

        let devices: Vec<i32> = spec
            .map(|s| s.placements.iter().map(|p| p.device_id).collect())
            .unwrap_or_default();

        let size = spec.map_or(0, ShardSpec::total_memory);

        let id = self.next_comm_id.fetch_add(1, Ordering::SeqCst);
        let comm = PendingComm {
            id,
            comm_type: CommType::ReduceScatter,
            src_devices: devices.clone(),
            dst_devices: devices,
            param_name,
            size_bytes: size,
        };

        self.pending_comms.lock().push(comm);
        id
    }

    /// Get pending communication count
    pub fn pending_count(&self) -> usize {
        self.pending_comms.lock().len()
    }

    /// Clear pending communications
    pub fn clear_pending(&self) {
        self.pending_comms.lock().clear();
    }

    /// Get sharder reference
    pub fn sharder(&self) -> Arc<ModelSharder> {
        self.sharder.clone()
    }
}

// =============================================================================
// Thread-Safe Wrappers
// =============================================================================

/// Thread-safe shared model sharder
pub type SharedModelSharder = Arc<ModelSharder>;

/// Create a shared model sharder
pub fn shared_model_sharder(config: ModelShardingConfig, devices: Vec<i32>) -> SharedModelSharder {
    Arc::new(ModelSharder::new(config, devices))
}

/// Thread-safe shared shard manager
pub type SharedShardManager = Arc<ShardManager>;

/// Create a shared shard manager
pub fn shared_shard_manager(sharder: SharedModelSharder) -> SharedShardManager {
    Arc::new(ShardManager::new(sharder))
}

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

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

    #[test]
    fn test_model_sharding_config_default() {
        let config = ModelShardingConfig::default();
        assert_eq!(config.num_devices, 1);
        assert_eq!(config.default_strategy, ShardingStrategy::TensorWise);
        assert!(config.balance_memory);
    }

    #[test]
    fn test_parameter_info_creation() {
        let param = ParameterInfo::new(
            "layer1.weight".to_string(),
            vec![1024, 512],
            "float32".to_string(),
        );

        assert_eq!(param.num_elements(), 1024 * 512);
        assert_eq!(param.size_bytes, 1024 * 512 * 4);
        assert!(param.can_shard_dim(0));
        assert!(param.can_shard_dim(1));
        assert!(!param.can_shard_dim(2));
    }

    #[test]
    fn test_parameter_info_dtypes() {
        let f32_param = ParameterInfo::new("p".to_string(), vec![100], "float32".to_string());
        assert_eq!(f32_param.size_bytes, 400);

        let f16_param = ParameterInfo::new("p".to_string(), vec![100], "float16".to_string());
        assert_eq!(f16_param.size_bytes, 200);

        let f64_param = ParameterInfo::new("p".to_string(), vec![100], "float64".to_string());
        assert_eq!(f64_param.size_bytes, 800);
    }

    #[test]
    fn test_shard_spec_replicated() {
        let spec = ShardSpec::replicated("weight".to_string(), vec![1024, 512], vec![0, 1, 2, 3]);

        assert_eq!(spec.num_shards(), 4);
        assert_eq!(spec.strategy, ShardingStrategy::Replicated);
        assert_eq!(spec.shard_shapes.len(), 4);

        // All shards should have same shape
        for shape in &spec.shard_shapes {
            assert_eq!(shape, &vec![1024, 512]);
        }
    }

    #[test]
    fn test_shard_spec_tensor_sharded() {
        let spec = ShardSpec::tensor_sharded(
            "weight".to_string(),
            vec![1024, 512],
            0, // Shard along dim 0
            vec![0, 1, 2, 3],
        );

        assert_eq!(spec.num_shards(), 4);
        assert_eq!(spec.strategy, ShardingStrategy::TensorWise);

        // Each shard should have 1/4 of dim 0
        for shape in &spec.shard_shapes {
            assert_eq!(shape[0], 256);
            assert_eq!(shape[1], 512);
        }
    }

    #[test]
    fn test_shard_spec_uneven_split() {
        let spec =
            ShardSpec::tensor_sharded("weight".to_string(), vec![10, 512], 0, vec![0, 1, 2, 3]);

        // 10 / 4 = 2 with remainder 2
        // First 2 shards get 3, last 2 get 2
        let sizes: Vec<_> = spec.shard_shapes.iter().map(|s| s[0]).collect();
        assert_eq!(sizes, vec![3, 3, 2, 2]);
    }

    #[test]
    fn test_sharding_plan() {
        let mut plan = ShardingPlan::new();

        let spec1 =
            ShardSpec::tensor_sharded("layer1.weight".to_string(), vec![1024, 512], 0, vec![0, 1]);
        let spec2 = ShardSpec::replicated("layer1.bias".to_string(), vec![512], vec![0, 1]);

        plan.add_spec(spec1);
        plan.add_spec(spec2);

        assert_eq!(plan.num_parameters(), 2);
        assert!(plan.get_spec("layer1.weight").is_some());
        assert!(plan.get_spec("layer1.bias").is_some());
    }

    #[test]
    fn test_sharding_plan_memory_balance() {
        let mut plan = ShardingPlan::new();

        // Add balanced shards
        let spec = ShardSpec::tensor_sharded("weight".to_string(), vec![1024, 1024], 0, vec![0, 1]);
        plan.add_spec(spec);

        let balance = plan.memory_balance();
        assert!(balance < 0.1); // Should be well balanced
    }

    #[test]
    fn test_model_sharder_creation() {
        let config = ModelShardingConfig {
            num_devices: 4,
            ..Default::default()
        };
        let sharder = ModelSharder::new(config, vec![0, 1, 2, 3]);

        assert_eq!(sharder.num_devices(), 4);
        assert_eq!(sharder.devices(), vec![0, 1, 2, 3]);
    }

    #[test]
    fn test_register_parameter() {
        let config = ModelShardingConfig::default();
        let sharder = ModelSharder::new(config, vec![0, 1]);

        let param = ParameterInfo::new(
            "layer1.weight".to_string(),
            vec![1024, 512],
            "float32".to_string(),
        );
        sharder.register_parameter(param);

        assert!(sharder.get_parameter("layer1.weight").is_some());
        assert_eq!(sharder.parameter_names().len(), 1);
    }

    #[test]
    fn test_generate_plan_tensor_wise() {
        let config = ModelShardingConfig {
            num_devices: 2,
            default_strategy: ShardingStrategy::TensorWise,
            ..Default::default()
        };
        let sharder = ModelSharder::new(config, vec![0, 1]);

        sharder.register_parameter(ParameterInfo::new(
            "weight".to_string(),
            vec![1024, 512],
            "float32".to_string(),
        ));

        let plan = sharder.generate_plan().unwrap();

        assert_eq!(plan.num_parameters(), 1);
        let spec = plan.get_spec("weight").unwrap();
        assert_eq!(spec.strategy, ShardingStrategy::TensorWise);
    }

    #[test]
    fn test_generate_plan_replicated() {
        let config = ModelShardingConfig {
            num_devices: 2,
            default_strategy: ShardingStrategy::Replicated,
            ..Default::default()
        };
        let sharder = ModelSharder::new(config, vec![0, 1]);

        sharder.register_parameter(ParameterInfo::new(
            "weight".to_string(),
            vec![1024, 512],
            "float32".to_string(),
        ));

        let plan = sharder.generate_plan().unwrap();
        let spec = plan.get_spec("weight").unwrap();
        assert_eq!(spec.strategy, ShardingStrategy::Replicated);
    }

    #[test]
    fn test_small_tensor_fallback_to_replicated() {
        let config = ModelShardingConfig {
            num_devices: 4,
            default_strategy: ShardingStrategy::TensorWise,
            ..Default::default()
        };
        let sharder = ModelSharder::new(config, vec![0, 1, 2, 3]);

        // Tensor too small to shard across 4 devices
        sharder.register_parameter(ParameterInfo::new(
            "small".to_string(),
            vec![2, 512],
            "float32".to_string(),
        ));

        let plan = sharder.generate_plan().unwrap();
        let spec = plan.get_spec("small").unwrap();
        // Should fall back to replicated
        assert_eq!(spec.strategy, ShardingStrategy::Replicated);
    }

    #[test]
    fn test_shard_manager_creation() {
        let config = ModelShardingConfig::default();
        let sharder = Arc::new(ModelSharder::new(config, vec![0, 1]));
        let manager = ShardManager::new(sharder);

        assert_eq!(manager.pending_count(), 0);
    }

    #[test]
    fn test_shard_manager_store_get() {
        let config = ModelShardingConfig::default();
        let sharder = Arc::new(ModelSharder::new(config, vec![0, 1]));
        let manager = ShardManager::new(sharder);

        let data = vec![1u8, 2, 3, 4];
        manager.store_shard(0, "weight".to_string(), data.clone());

        let retrieved = manager.get_shard(0, "weight").unwrap();
        assert_eq!(retrieved, data);
    }

    #[test]
    fn test_schedule_all_gather() {
        let config = ModelShardingConfig::default();
        let sharder = Arc::new(ModelSharder::new(config, vec![0, 1]));

        sharder.register_parameter(ParameterInfo::new(
            "weight".to_string(),
            vec![1024, 512],
            "float32".to_string(),
        ));
        sharder.generate_plan().unwrap();

        let manager = ShardManager::new(sharder);
        let id = manager.schedule_all_gather("weight".to_string());

        assert!(id > 0);
        assert_eq!(manager.pending_count(), 1);
    }

    #[test]
    fn test_schedule_reduce_scatter() {
        let config = ModelShardingConfig::default();
        let sharder = Arc::new(ModelSharder::new(config, vec![0, 1]));

        sharder.register_parameter(ParameterInfo::new(
            "weight".to_string(),
            vec![1024, 512],
            "float32".to_string(),
        ));
        sharder.generate_plan().unwrap();

        let manager = ShardManager::new(sharder);
        let id = manager.schedule_reduce_scatter("weight".to_string());

        assert!(id > 0);
        assert_eq!(manager.pending_count(), 1);
    }

    #[test]
    fn test_clear_pending() {
        let config = ModelShardingConfig::default();
        let sharder = Arc::new(ModelSharder::new(config, vec![0, 1]));
        let manager = ShardManager::new(sharder.clone());

        sharder.register_parameter(ParameterInfo::new(
            "weight".to_string(),
            vec![1024, 512],
            "float32".to_string(),
        ));
        sharder.generate_plan().unwrap();

        manager.schedule_all_gather("weight".to_string());
        manager.schedule_reduce_scatter("weight".to_string());
        assert_eq!(manager.pending_count(), 2);

        manager.clear_pending();
        assert_eq!(manager.pending_count(), 0);
    }

    #[test]
    fn test_shared_model_sharder() {
        let config = ModelShardingConfig::default();
        let sharder = shared_model_sharder(config, vec![0, 1, 2, 3]);

        assert_eq!(sharder.num_devices(), 4);
    }

    #[test]
    fn test_shared_shard_manager() {
        let config = ModelShardingConfig::default();
        let sharder = shared_model_sharder(config, vec![0, 1]);
        let manager = shared_shard_manager(sharder);

        assert_eq!(manager.pending_count(), 0);
    }

    #[test]
    fn test_sharding_stats() {
        let config = ModelShardingConfig::default();
        let sharder = ModelSharder::new(config, vec![0, 1]);

        sharder.register_parameter(ParameterInfo::new(
            "weight".to_string(),
            vec![1024, 512],
            "float32".to_string(),
        ));
        sharder.generate_plan().unwrap();

        let stats = sharder.stats();
        assert_eq!(stats.plans_generated, 1);
    }

    #[test]
    fn test_layer_wise_sharding() {
        let config = ModelShardingConfig {
            num_devices: 4,
            default_strategy: ShardingStrategy::LayerWise,
            ..Default::default()
        };
        let sharder = ModelSharder::new(config, vec![0, 1, 2, 3]);

        let mut param = ParameterInfo::new(
            "layer2.weight".to_string(),
            vec![1024, 512],
            "float32".to_string(),
        );
        param.layer_index = Some(2);
        sharder.register_parameter(param);

        let plan = sharder.generate_plan().unwrap();
        let spec = plan.get_spec("layer2.weight").unwrap();

        // Layer 2 should be on device 2 (2 % 4 = 2)
        assert_eq!(spec.placements.len(), 1);
        assert_eq!(spec.placements[0].device_id, 2);
    }

    #[test]
    fn test_column_wise_sharding() {
        let config = ModelShardingConfig {
            num_devices: 2,
            default_strategy: ShardingStrategy::ColumnWise,
            ..Default::default()
        };
        let sharder = ModelSharder::new(config, vec![0, 1]);

        sharder.register_parameter(ParameterInfo::new(
            "weight".to_string(),
            vec![1024, 512],
            "float32".to_string(),
        ));

        let plan = sharder.generate_plan().unwrap();
        let spec = plan.get_spec("weight").unwrap();

        // Should shard along last dim (1)
        assert_eq!(spec.shard_dim, ShardDimension::Dim(1));
        for shape in &spec.shard_shapes {
            assert_eq!(shape[0], 1024);
            assert_eq!(shape[1], 256); // 512 / 2
        }
    }
}