//! Automatic Model Partitioning for Distributed Training
//!
//! This module provides automatic model partitioning capabilities that
//! analyze model architecture and generate optimal sharding plans for
//! tensor parallelism (TP), pipeline parallelism (PP), and FSDP.
//!
//! # Features
//!
//! - Memory-aware partitioning to fit models on available GPU memory
//! - Compute-balanced distribution for optimal throughput
//! - Communication-minimal placement to reduce data movement
//! - Hybrid strategies combining multiple optimization goals
//!
//! # Example
//!
//! ```rust,ignore
//! use rtx_distributed::auto_partition::{AutoPartitioner, AutoPartitionConfig, ModelInfo};
//!
//! let config = AutoPartitionConfig::default();
//! let partitioner = AutoPartitioner::new(config);
//!
//! let model_info = ModelInfo::from_layer_sizes(&[
//!     ("embedding", 1_000_000_000),
//!     ("transformer.0", 500_000_000),
//!     ("transformer.1", 500_000_000),
//!     ("lm_head", 500_000_000),
//! ]);
//!
//! let plan = partitioner.partition(&model_info)?;
//! println!("{:?}", plan);
//! ```

use crate::error::{DistributedError, Result};
use crate::model_sharding::ShardingStrategy;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;

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

/// Strategy for automatic partitioning
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AutoPartitionStrategy {
    /// Minimize peak memory usage
    MemoryOptimal,
    /// Balance compute across devices
    ComputeBalanced,
    /// Minimize cross-device communication
    CommunicationMinimal,
    /// Balance all factors with weighted importance
    Hybrid,
    /// Prefer tensor parallelism for large layers
    TensorParallelFirst,
    /// Prefer pipeline parallelism for sequential models
    PipelineParallelFirst,
}

impl Default for AutoPartitionStrategy {
    fn default() -> Self {
        Self::Hybrid
    }
}

/// Configuration for the auto partitioner
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AutoPartitionConfig {
    /// Partitioning strategy
    pub strategy: AutoPartitionStrategy,
    /// Memory margin (0.1 = leave 10% headroom)
    pub memory_margin: f32,
    /// Maximum allowed compute imbalance ratio (1.2 = 20% max difference)
    pub balance_threshold: f32,
    /// Prefer tensor parallelism for large layers
    pub prefer_tp_for_large_layers: bool,
    /// Prefer pipeline parallelism for sequential structures
    pub prefer_pp_for_sequential: bool,
    /// Minimum layer size (bytes) to consider for tensor parallelism
    pub min_tp_layer_size: usize,
    /// Number of devices available
    pub num_devices: usize,
    /// Per-device memory budget (bytes)
    pub device_memory: usize,
    /// Weight for memory optimization in hybrid mode
    pub memory_weight: f32,
    /// Weight for compute balance in hybrid mode
    pub compute_weight: f32,
    /// Weight for communication minimization in hybrid mode
    pub communication_weight: f32,
}

impl Default for AutoPartitionConfig {
    fn default() -> Self {
        Self {
            strategy: AutoPartitionStrategy::Hybrid,
            memory_margin: 0.1,
            balance_threshold: 1.2,
            prefer_tp_for_large_layers: true,
            prefer_pp_for_sequential: true,
            min_tp_layer_size: 100 * 1024 * 1024, // 100 MB
            num_devices: 8,
            device_memory: 80 * 1024 * 1024 * 1024, // 80 GB (A100/H100)
            memory_weight: 0.4,
            compute_weight: 0.3,
            communication_weight: 0.3,
        }
    }
}

impl AutoPartitionConfig {
    /// Create config for memory-constrained environments
    pub fn memory_optimized(num_devices: usize, device_memory: usize) -> Self {
        Self {
            strategy: AutoPartitionStrategy::MemoryOptimal,
            num_devices,
            device_memory,
            memory_margin: 0.15,
            memory_weight: 0.6,
            compute_weight: 0.2,
            communication_weight: 0.2,
            ..Default::default()
        }
    }

    /// Create config for throughput-focused training
    pub fn throughput_optimized(num_devices: usize) -> Self {
        Self {
            strategy: AutoPartitionStrategy::ComputeBalanced,
            num_devices,
            memory_weight: 0.2,
            compute_weight: 0.6,
            communication_weight: 0.2,
            ..Default::default()
        }
    }

    /// Create config for communication-sensitive environments
    pub fn communication_optimized(num_devices: usize) -> Self {
        Self {
            strategy: AutoPartitionStrategy::CommunicationMinimal,
            num_devices,
            memory_weight: 0.2,
            compute_weight: 0.2,
            communication_weight: 0.6,
            ..Default::default()
        }
    }
}

// =============================================================================
// Model Information
// =============================================================================

/// Information about a model layer
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LayerInfo {
    /// Layer name/path
    pub name: String,
    /// Parameter count
    pub param_count: usize,
    /// Memory size in bytes (including activations estimate)
    pub memory_bytes: usize,
    /// Estimated FLOPs per forward pass
    pub flops: usize,
    /// Layer type
    pub layer_type: LayerType,
    /// Whether this layer is shardable across tensor dimension
    pub is_tp_shardable: bool,
    /// Connections to other layers (for graph analysis)
    pub connections: Vec<String>,
}

/// Types of model layers
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum LayerType {
    /// Embedding layer
    Embedding,
    /// Linear/Dense layer
    Linear,
    /// Attention layer
    Attention,
    /// LayerNorm/BatchNorm
    Normalization,
    /// Activation function
    Activation,
    /// Convolutional layer
    Convolution,
    /// Pooling layer
    Pooling,
    /// Residual connection
    Residual,
    /// Other/Unknown
    Other,
}

impl LayerInfo {
    /// Create from basic size information
    pub fn new(name: impl Into<String>, param_count: usize, layer_type: LayerType) -> Self {
        // Estimate 4 bytes per param + 3x for gradients and optimizer states
        let memory_bytes = param_count * 4 * 4;
        // Rough FLOP estimate based on layer type
        let flops = match layer_type {
            LayerType::Linear => param_count * 2,
            LayerType::Attention => param_count * 4,
            _ => param_count,
        };

        Self {
            name: name.into(),
            param_count,
            memory_bytes,
            flops,
            layer_type,
            is_tp_shardable: matches!(
                layer_type,
                LayerType::Linear | LayerType::Embedding | LayerType::Attention
            ),
            connections: Vec::new(),
        }
    }
}

/// Complete model information for partitioning
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelInfo {
    /// Model name
    pub name: String,
    /// Total parameter count
    pub total_params: usize,
    /// Total memory footprint
    pub total_memory: usize,
    /// Layers in execution order
    pub layers: Vec<LayerInfo>,
    /// Is the model purely sequential?
    pub is_sequential: bool,
    /// Detected transformer blocks
    pub transformer_blocks: Vec<TransformerBlock>,
}

/// Detected transformer block structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransformerBlock {
    /// Block index
    pub index: usize,
    /// Layer names in this block
    pub layer_names: Vec<String>,
    /// Block memory size
    pub memory_bytes: usize,
    /// Block FLOPs
    pub flops: usize,
}

impl ModelInfo {
    /// Create from layer sizes
    pub fn from_layer_sizes(layers: &[(&str, usize)]) -> Self {
        let layer_infos: Vec<LayerInfo> = layers
            .iter()
            .map(|(name, size)| {
                let layer_type = if name.contains("embed") {
                    LayerType::Embedding
                } else if name.contains("attention") || name.contains("attn") {
                    LayerType::Attention
                } else if name.contains("norm") {
                    LayerType::Normalization
                } else {
                    LayerType::Linear
                };
                LayerInfo::new(*name, *size / 4, layer_type) // Convert bytes to param count
            })
            .collect();

        let total_params = layer_infos.iter().map(|l| l.param_count).sum();
        let total_memory = layer_infos.iter().map(|l| l.memory_bytes).sum();

        Self {
            name: "model".to_string(),
            total_params,
            total_memory,
            layers: layer_infos,
            is_sequential: true,
            transformer_blocks: Vec::new(),
        }
    }

    /// Create from layer info
    pub fn from_layers(name: impl Into<String>, layers: Vec<LayerInfo>) -> Self {
        let total_params = layers.iter().map(|l| l.param_count).sum();
        let total_memory = layers.iter().map(|l| l.memory_bytes).sum();

        Self {
            name: name.into(),
            total_params,
            total_memory,
            layers,
            is_sequential: true,
            transformer_blocks: Vec::new(),
        }
    }

    /// Detect transformer blocks in the model
    pub fn detect_transformer_blocks(&mut self) {
        let mut blocks = Vec::new();
        let mut current_block_layers = Vec::new();
        let mut current_block_start = 0;

        for (i, layer) in self.layers.iter().enumerate() {
            if layer.name.contains("transformer") || layer.name.contains("block") {
                // Extract block index from name
                if let Some(idx) = extract_block_index(&layer.name) {
                    if !current_block_layers.is_empty() && idx != current_block_start {
                        // Save previous block
                        blocks.push(TransformerBlock {
                            index: current_block_start,
                            layer_names: current_block_layers.clone(),
                            memory_bytes: current_block_layers
                                .iter()
                                .filter_map(|name| self.layers.iter().find(|l| &l.name == name))
                                .map(|l| l.memory_bytes)
                                .sum(),
                            flops: current_block_layers
                                .iter()
                                .filter_map(|name| self.layers.iter().find(|l| &l.name == name))
                                .map(|l| l.flops)
                                .sum(),
                        });
                        current_block_layers.clear();
                    }
                    current_block_start = idx;
                }
                current_block_layers.push(layer.name.clone());
            }

            // Save last block
            if i == self.layers.len() - 1 && !current_block_layers.is_empty() {
                blocks.push(TransformerBlock {
                    index: current_block_start,
                    layer_names: current_block_layers.clone(),
                    memory_bytes: current_block_layers
                        .iter()
                        .filter_map(|name| self.layers.iter().find(|l| &l.name == name))
                        .map(|l| l.memory_bytes)
                        .sum(),
                    flops: current_block_layers
                        .iter()
                        .filter_map(|name| self.layers.iter().find(|l| &l.name == name))
                        .map(|l| l.flops)
                        .sum(),
                });
            }
        }

        self.transformer_blocks = blocks;
    }
}

/// Extract block index from layer name (e.g., "transformer.12.attention" -> 12)
fn extract_block_index(name: &str) -> Option<usize> {
    name.split('.').find_map(|s| s.parse::<usize>().ok())
}

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

/// How a layer should be partitioned
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LayerPartition {
    /// Layer name
    pub layer_name: String,
    /// Sharding strategy for this layer
    pub strategy: ShardingStrategy,
    /// Device assignments (for PP)
    pub device_ids: Vec<usize>,
    /// TP dimension (for TP sharding)
    pub tp_size: usize,
    /// PP stage (for PP sharding)
    pub pp_stage: usize,
}

/// Complete sharding plan for a model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ShardingPlan {
    /// Model name
    pub model_name: String,
    /// Per-layer partitioning
    pub layer_partitions: Vec<LayerPartition>,
    /// Total number of pipeline stages
    pub num_pp_stages: usize,
    /// Tensor parallel size
    pub tp_size: usize,
    /// Data parallel size (derived from devices / (tp * pp))
    pub dp_size: usize,
    /// Estimated peak memory per device
    pub estimated_memory_per_device: usize,
    /// Estimated compute imbalance factor
    pub compute_imbalance: f32,
    /// Estimated communication volume
    pub communication_volume: usize,
    /// Score from optimization
    pub optimization_score: f64,
}

impl ShardingPlan {
    /// Create a new sharding plan
    pub fn new(model_name: impl Into<String>) -> Self {
        Self {
            model_name: model_name.into(),
            layer_partitions: Vec::new(),
            num_pp_stages: 1,
            tp_size: 1,
            dp_size: 1,
            estimated_memory_per_device: 0,
            compute_imbalance: 1.0,
            communication_volume: 0,
            optimization_score: 0.0,
        }
    }

    /// Add a layer partition
    pub fn add_partition(&mut self, partition: LayerPartition) {
        self.layer_partitions.push(partition);
    }
}

// =============================================================================
// Layer Profile
// =============================================================================

/// Profile of a layer's resource requirements
#[derive(Debug, Clone)]
struct LayerProfile {
    /// Layer info
    info: LayerInfo,
    /// Memory required (bytes)
    memory: usize,
    /// Compute cost (FLOPs)
    compute: usize,
    /// Communication cost if sharded
    comm_cost: usize,
    /// Can this layer be sharded with TP?
    can_tp_shard: bool,
}

// =============================================================================
// Auto Partitioner
// =============================================================================

/// Automatic model partitioner
pub struct AutoPartitioner {
    /// Configuration
    config: AutoPartitionConfig,
    /// Per-device memory budgets
    memory_budget: RwLock<HashMap<usize, usize>>,
    /// Cached layer profiles
    layer_profiles: RwLock<Vec<LayerProfile>>,
}

impl AutoPartitioner {
    /// Create a new auto partitioner
    pub fn new(config: AutoPartitionConfig) -> Self {
        let mut memory_budget = HashMap::new();
        let effective_memory =
            (config.device_memory as f32 * (1.0 - config.memory_margin)) as usize;
        for i in 0..config.num_devices {
            memory_budget.insert(i, effective_memory);
        }

        Self {
            config,
            memory_budget: RwLock::new(memory_budget),
            layer_profiles: RwLock::new(Vec::new()),
        }
    }

    /// Analyze model and generate optimal sharding plan
    pub fn partition(&self, model_info: &ModelInfo) -> Result<ShardingPlan> {
        // Profile all layers
        let profiles = self.profile_layers(model_info);
        *self.layer_profiles.write() = profiles.clone();

        // Check if model fits on single device
        if model_info.total_memory <= self.effective_memory_per_device() {
            return self.create_single_device_plan(model_info);
        }

        // Choose partitioning approach based on strategy
        let plan = match self.config.strategy {
            AutoPartitionStrategy::MemoryOptimal => {
                self.partition_memory_optimal(model_info, &profiles)?
            }
            AutoPartitionStrategy::ComputeBalanced => {
                self.partition_compute_balanced(model_info, &profiles)?
            }
            AutoPartitionStrategy::CommunicationMinimal => {
                self.partition_comm_minimal(model_info, &profiles)?
            }
            AutoPartitionStrategy::Hybrid => self.partition_hybrid(model_info, &profiles)?,
            AutoPartitionStrategy::TensorParallelFirst => {
                self.partition_tp_first(model_info, &profiles)?
            }
            AutoPartitionStrategy::PipelineParallelFirst => {
                self.partition_pp_first(model_info, &profiles)?
            }
        };

        Ok(plan)
    }

    /// Get effective memory per device after margin
    fn effective_memory_per_device(&self) -> usize {
        (self.config.device_memory as f32 * (1.0 - self.config.memory_margin)) as usize
    }

    /// Profile layers for resource requirements
    fn profile_layers(&self, model: &ModelInfo) -> Vec<LayerProfile> {
        model
            .layers
            .iter()
            .map(|layer| {
                let comm_cost = if layer.is_tp_shardable {
                    // Communication cost proportional to layer size
                    layer.memory_bytes / self.config.num_devices
                } else {
                    0
                };

                LayerProfile {
                    info: layer.clone(),
                    memory: layer.memory_bytes,
                    compute: layer.flops,
                    comm_cost,
                    can_tp_shard: layer.is_tp_shardable,
                }
            })
            .collect()
    }

    /// Create plan for single-device execution
    fn create_single_device_plan(&self, model: &ModelInfo) -> Result<ShardingPlan> {
        let mut plan = ShardingPlan::new(&model.name);
        plan.num_pp_stages = 1;
        plan.tp_size = 1;
        plan.dp_size = self.config.num_devices;
        plan.estimated_memory_per_device = model.total_memory;
        plan.compute_imbalance = 1.0;

        for layer in &model.layers {
            plan.add_partition(LayerPartition {
                layer_name: layer.name.clone(),
                strategy: ShardingStrategy::Replicated,
                device_ids: vec![0],
                tp_size: 1,
                pp_stage: 0,
            });
        }

        Ok(plan)
    }

    /// Memory-optimal partitioning (minimize peak memory)
    fn partition_memory_optimal(
        &self,
        model: &ModelInfo,
        profiles: &[LayerProfile],
    ) -> Result<ShardingPlan> {
        let effective_memory = self.effective_memory_per_device();

        // Calculate required TP degree to fit model
        let mut tp_size = 1;
        while model.total_memory / tp_size > effective_memory && tp_size < self.config.num_devices {
            tp_size *= 2;
        }

        // If still doesn't fit, add PP
        let mut pp_stages = 1;
        let memory_per_tp = model.total_memory / tp_size;
        if memory_per_tp > effective_memory {
            pp_stages = (memory_per_tp + effective_memory - 1) / effective_memory;
            pp_stages = pp_stages.min(self.config.num_devices / tp_size);
        }

        self.create_tp_pp_plan(model, profiles, tp_size, pp_stages)
    }

    /// Compute-balanced partitioning
    fn partition_compute_balanced(
        &self,
        model: &ModelInfo,
        profiles: &[LayerProfile],
    ) -> Result<ShardingPlan> {
        let total_compute: usize = profiles.iter().map(|p| p.compute).sum();
        let _target_per_device = total_compute / self.config.num_devices;

        // Greedy assignment to balance compute
        let mut device_compute = vec![0usize; self.config.num_devices];
        let mut assignments = Vec::new();

        for profile in profiles {
            // Find device with least compute
            let min_device = device_compute
                .iter()
                .enumerate()
                .min_by_key(|(_, c)| *c)
                .map_or(0, |(i, _)| i);

            assignments.push((profile.info.name.clone(), min_device));
            device_compute[min_device] += profile.compute;
        }

        // Convert to sharding plan
        let mut plan = ShardingPlan::new(&model.name);
        plan.num_pp_stages = 1;
        plan.tp_size = 1;
        plan.dp_size = self.config.num_devices;

        for (layer_name, device) in assignments {
            plan.add_partition(LayerPartition {
                layer_name,
                strategy: ShardingStrategy::Replicated,
                device_ids: vec![device],
                tp_size: 1,
                pp_stage: device,
            });
        }

        // Calculate imbalance
        let max_compute = device_compute.iter().max().copied().unwrap_or(0);
        let min_compute = device_compute.iter().min().copied().unwrap_or(1);
        plan.compute_imbalance = if min_compute > 0 {
            max_compute as f32 / min_compute as f32
        } else {
            f32::MAX
        };

        Ok(plan)
    }

    /// Communication-minimal partitioning
    fn partition_comm_minimal(
        &self,
        model: &ModelInfo,
        profiles: &[LayerProfile],
    ) -> Result<ShardingPlan> {
        // Prefer PP over TP to minimize communication
        let effective_memory = self.effective_memory_per_device();
        let memory_per_device = model.total_memory / self.config.num_devices;

        if memory_per_device <= effective_memory {
            // Pure PP is possible
            self.create_pp_only_plan(model, profiles)
        } else {
            // Need some TP
            let tp_size = (model.total_memory + effective_memory - 1) / effective_memory;
            let tp_size = tp_size.min(self.config.num_devices).next_power_of_two();
            self.create_tp_pp_plan(model, profiles, tp_size, 1)
        }
    }

    /// Hybrid partitioning (balance all factors)
    fn partition_hybrid(
        &self,
        model: &ModelInfo,
        profiles: &[LayerProfile],
    ) -> Result<ShardingPlan> {
        let effective_memory = self.effective_memory_per_device();

        // Try different TP/PP combinations and score them
        let mut best_plan = None;
        let mut best_score = f64::MIN;

        for tp in [1, 2, 4, 8]
            .iter()
            .filter(|&&t| t <= self.config.num_devices)
        {
            let max_pp = self.config.num_devices / tp;
            for pp in 1..=max_pp {
                let memory_per_device = model.total_memory / (tp * pp);
                if memory_per_device > effective_memory {
                    continue;
                }

                if let Ok(plan) = self.create_tp_pp_plan(model, profiles, *tp, pp) {
                    let score = self.score_plan(&plan);
                    if score > best_score {
                        best_score = score;
                        best_plan = Some(plan);
                    }
                }
            }
        }

        best_plan.ok_or_else(|| {
            DistributedError::configuration("Could not find valid partitioning scheme")
        })
    }

    /// Tensor-parallel first partitioning
    fn partition_tp_first(
        &self,
        model: &ModelInfo,
        profiles: &[LayerProfile],
    ) -> Result<ShardingPlan> {
        let effective_memory = self.effective_memory_per_device();

        // Maximize TP before using PP
        let mut tp_size = self.config.num_devices;
        while tp_size > 1 {
            let memory_per_device = model.total_memory / tp_size;
            if memory_per_device <= effective_memory {
                break;
            }
            tp_size /= 2;
        }

        if tp_size == 0 {
            tp_size = 1;
        }

        let pp_stages = self.config.num_devices / tp_size;
        self.create_tp_pp_plan(model, profiles, tp_size, pp_stages)
    }

    /// Pipeline-parallel first partitioning
    fn partition_pp_first(
        &self,
        model: &ModelInfo,
        profiles: &[LayerProfile],
    ) -> Result<ShardingPlan> {
        let effective_memory = self.effective_memory_per_device();

        // Maximize PP before using TP
        let mut pp_stages = self.config.num_devices;
        while pp_stages > 1 {
            let memory_per_device = model.total_memory / pp_stages;
            if memory_per_device <= effective_memory {
                break;
            }
            pp_stages /= 2;
        }

        if pp_stages == 0 {
            pp_stages = 1;
        }

        let tp_size = self.config.num_devices / pp_stages;
        self.create_tp_pp_plan(model, profiles, tp_size, pp_stages)
    }

    /// Create pure pipeline parallel plan
    fn create_pp_only_plan(
        &self,
        model: &ModelInfo,
        profiles: &[LayerProfile],
    ) -> Result<ShardingPlan> {
        let total_memory: usize = profiles.iter().map(|p| p.memory).sum();
        let target_per_stage = total_memory / self.config.num_devices;

        let mut plan = ShardingPlan::new(&model.name);
        plan.tp_size = 1;
        plan.num_pp_stages = self.config.num_devices;
        plan.dp_size = 1;

        let mut current_stage = 0;
        let mut stage_memory = 0;

        for profile in profiles {
            if stage_memory + profile.memory > target_per_stage * 2
                && current_stage < self.config.num_devices - 1
            {
                current_stage += 1;
                stage_memory = 0;
            }

            plan.add_partition(LayerPartition {
                layer_name: profile.info.name.clone(),
                strategy: ShardingStrategy::LayerWise,
                device_ids: vec![current_stage],
                tp_size: 1,
                pp_stage: current_stage,
            });

            stage_memory += profile.memory;
        }

        plan.estimated_memory_per_device = target_per_stage;
        Ok(plan)
    }

    /// Create combined TP + PP plan
    fn create_tp_pp_plan(
        &self,
        model: &ModelInfo,
        profiles: &[LayerProfile],
        tp_size: usize,
        pp_stages: usize,
    ) -> Result<ShardingPlan> {
        let mut plan = ShardingPlan::new(&model.name);
        plan.tp_size = tp_size;
        plan.num_pp_stages = pp_stages;
        plan.dp_size = self.config.num_devices / (tp_size * pp_stages);

        let total_memory: usize = profiles.iter().map(|p| p.memory).sum();
        let memory_per_stage = total_memory / pp_stages;
        let target_per_stage = memory_per_stage / tp_size;

        let mut current_stage = 0;
        let mut stage_memory = 0;

        for profile in profiles {
            if stage_memory + profile.memory > target_per_stage * 2 && current_stage < pp_stages - 1
            {
                current_stage += 1;
                stage_memory = 0;
            }

            let strategy = if profile.can_tp_shard && tp_size > 1 {
                ShardingStrategy::TensorWise
            } else if pp_stages > 1 {
                ShardingStrategy::LayerWise
            } else {
                ShardingStrategy::Replicated
            };

            // Device IDs for this layer
            let device_ids: Vec<usize> =
                (0..tp_size).map(|t| current_stage * tp_size + t).collect();

            plan.add_partition(LayerPartition {
                layer_name: profile.info.name.clone(),
                strategy,
                device_ids,
                tp_size,
                pp_stage: current_stage,
            });

            stage_memory += profile.memory;
        }

        plan.estimated_memory_per_device = target_per_stage;
        plan.communication_volume = profiles
            .iter()
            .filter(|p| p.can_tp_shard)
            .map(|p| p.comm_cost)
            .sum();

        Ok(plan)
    }

    /// Score a sharding plan (higher is better)
    fn score_plan(&self, plan: &ShardingPlan) -> f64 {
        let effective_memory = self.effective_memory_per_device() as f64;

        // Memory score (1.0 if fits well, decreasing as it gets tighter)
        let memory_score = if plan.estimated_memory_per_device as f64 <= effective_memory {
            1.0 - (plan.estimated_memory_per_device as f64 / effective_memory) * 0.5
        } else {
            0.0
        };

        // Compute balance score
        let balance_score = if plan.compute_imbalance <= self.config.balance_threshold {
            1.0 / plan.compute_imbalance as f64
        } else {
            0.5 / plan.compute_imbalance as f64
        };

        // Communication score (lower volume is better)
        let max_comm = plan.estimated_memory_per_device as f64;
        let comm_score = 1.0 - (plan.communication_volume as f64 / max_comm).min(1.0);

        // Weighted combination
        self.config.memory_weight as f64 * memory_score
            + self.config.compute_weight as f64 * balance_score
            + self.config.communication_weight as f64 * comm_score
    }
}

/// Thread-safe shared auto partitioner
pub type SharedAutoPartitioner = Arc<AutoPartitioner>;

/// Create a shared auto partitioner
pub fn shared_auto_partitioner(config: AutoPartitionConfig) -> SharedAutoPartitioner {
    Arc::new(AutoPartitioner::new(config))
}

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

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

    #[test]
    fn test_auto_partition_config_default() {
        let config = AutoPartitionConfig::default();
        assert_eq!(config.strategy, AutoPartitionStrategy::Hybrid);
        assert!((config.memory_margin - 0.1).abs() < 0.01);
    }

    #[test]
    fn test_auto_partition_config_memory_optimized() {
        let config = AutoPartitionConfig::memory_optimized(8, 80 * 1024 * 1024 * 1024);
        assert_eq!(config.strategy, AutoPartitionStrategy::MemoryOptimal);
        assert_eq!(config.num_devices, 8);
    }

    #[test]
    fn test_layer_info_creation() {
        let layer = LayerInfo::new("fc1", 1000000, LayerType::Linear);
        assert_eq!(layer.param_count, 1000000);
        assert!(layer.is_tp_shardable);
    }

    #[test]
    fn test_model_info_from_layers() {
        let model = ModelInfo::from_layer_sizes(&[
            ("embed", 1_000_000),
            ("fc1", 500_000),
            ("fc2", 500_000),
        ]);

        assert_eq!(model.layers.len(), 3);
        assert_eq!(model.total_params, 500_000); // Divided by 4 for bytes->params
    }

    #[test]
    fn test_sharding_plan_creation() {
        let mut plan = ShardingPlan::new("test_model");
        plan.add_partition(LayerPartition {
            layer_name: "layer1".to_string(),
            strategy: ShardingStrategy::Replicated,
            device_ids: vec![0],
            tp_size: 1,
            pp_stage: 0,
        });

        assert_eq!(plan.layer_partitions.len(), 1);
    }

    #[test]
    fn test_auto_partitioner_single_device() {
        let config = AutoPartitionConfig {
            num_devices: 1,
            device_memory: 80 * 1024 * 1024 * 1024,
            ..Default::default()
        };
        let partitioner = AutoPartitioner::new(config);

        let model = ModelInfo::from_layer_sizes(&[("embed", 1_000_000), ("fc1", 500_000)]);

        let plan = partitioner.partition(&model).unwrap();
        assert_eq!(plan.tp_size, 1);
        assert_eq!(plan.num_pp_stages, 1);
    }

    #[test]
    fn test_auto_partitioner_large_model() {
        let config = AutoPartitionConfig {
            num_devices: 8,
            device_memory: 10 * 1024 * 1024 * 1024, // 10 GB each
            ..Default::default()
        };
        let partitioner = AutoPartitioner::new(config);

        // Create a model that needs distribution (but is tractable)
        // With 10% margin, each device can hold ~9 GB
        // Each layer should be ~1 GB of params = 4 GB memory (4 bytes * 4x for grad+optim)
        // So 10 layers = 40 GB total, needs 8 devices
        let layers: Vec<_> = (0..10)
            .map(|i| (format!("layer{}", i), 1024 * 1024 * 1024usize)) // 1 GB params each
            .collect();
        let layer_refs: Vec<_> = layers.iter().map(|(n, s)| (n.as_str(), *s)).collect();
        let model = ModelInfo::from_layer_sizes(&layer_refs);

        let plan = partitioner.partition(&model).unwrap();
        // Should use multiple devices
        assert!(plan.tp_size > 1 || plan.num_pp_stages > 1);
    }

    #[test]
    fn test_partition_strategy_variants() {
        assert_eq!(
            AutoPartitionStrategy::default(),
            AutoPartitionStrategy::Hybrid
        );
        assert_ne!(
            AutoPartitionStrategy::MemoryOptimal,
            AutoPartitionStrategy::ComputeBalanced
        );
    }

    #[test]
    fn test_extract_block_index() {
        assert_eq!(extract_block_index("transformer.12.attention"), Some(12));
        assert_eq!(extract_block_index("block.5.mlp"), Some(5));
        assert_eq!(extract_block_index("embedding"), None);
    }

    #[test]
    fn test_layer_type() {
        assert!(matches!(LayerType::Linear, LayerType::Linear));
        assert!(!matches!(LayerType::Linear, LayerType::Attention));
    }

    #[test]
    fn test_effective_memory() {
        let config = AutoPartitionConfig {
            device_memory: 100,
            memory_margin: 0.1,
            ..Default::default()
        };
        let partitioner = AutoPartitioner::new(config);
        assert_eq!(partitioner.effective_memory_per_device(), 90);
    }
}