//! Parallelism strategies for distributed LLM inference.
//!
//! This module implements tensor parallelism and pipeline parallelism
//! for distributing large language models across multiple nodes.

use distllm_shared::LayerAssignment;

// ============================================================================
// Tensor Parallelism
// ============================================================================

/// Tensor parallelism for splitting linear layers across nodes.
///
/// In tensor parallelism, each attention head and FFN dimension is split
/// across multiple devices. This requires all-reduce after each layer.
#[derive(Debug, Clone)]
pub struct TensorParallel {
    /// Tensor parallel world size.
    pub world_size: usize,
    /// This node's rank in tensor parallel group.
    pub rank: usize,
    /// Shard size for attention heads.
    pub head_shard_size: usize,
    /// Shard size for FFN hidden dimension.
    pub ffn_shard_size: usize,
}

impl TensorParallel {
    /// Create a new tensor parallel configuration.
    #[must_use]
    pub fn new(world_size: usize, rank: usize) -> Self {
        Self {
            world_size,
            rank,
            head_shard_size: 0,
            ffn_shard_size: 0,
        }
    }

    /// Configure for a specific model.
    pub fn configure(&mut self, num_heads: usize, ffn_dim: usize) {
        self.head_shard_size = num_heads / self.world_size;
        self.ffn_shard_size = ffn_dim / self.world_size;
    }

    /// Get the attention head range for this rank.
    #[must_use]
    pub fn head_range(&self, total_heads: usize) -> (usize, usize) {
        let heads_per_rank = total_heads / self.world_size;
        let start = self.rank * heads_per_rank;
        let end = if self.rank == self.world_size - 1 {
            total_heads
        } else {
            (self.rank + 1) * heads_per_rank
        };
        (start, end)
    }

    /// Get the FFN dimension range for this rank.
    #[must_use]
    pub fn ffn_range(&self, total_dim: usize) -> (usize, usize) {
        let dim_per_rank = total_dim / self.world_size;
        let start = self.rank * dim_per_rank;
        let end = if self.rank == self.world_size - 1 {
            total_dim
        } else {
            (self.rank + 1) * dim_per_rank
        };
        (start, end)
    }

    /// Split a linear layer weight matrix column-wise.
    #[must_use]
    pub fn column_split(&self, weight_shape: (usize, usize)) -> (usize, usize) {
        let (rows, cols) = weight_shape;
        let cols_per_rank = cols / self.world_size;
        (rows, cols_per_rank)
    }

    /// Split a linear layer weight matrix row-wise.
    #[must_use]
    pub fn row_split(&self, weight_shape: (usize, usize)) -> (usize, usize) {
        let (rows, cols) = weight_shape;
        let rows_per_rank = rows / self.world_size;
        (rows_per_rank, cols)
    }

    /// Calculate communication volume for all-reduce in MB.
    #[must_use]
    pub fn all_reduce_volume_mb(
        &self,
        hidden_dim: usize,
        seq_len: usize,
        dtype_bytes: usize,
    ) -> f64 {
        let elements = hidden_dim * seq_len;
        let bytes = elements * dtype_bytes;
        bytes as f64 / (1024.0 * 1024.0)
    }
}

// ============================================================================
// Pipeline Parallelism
// ============================================================================

/// Pipeline parallelism for layer-wise distribution.
///
/// In pipeline parallelism, consecutive layers are assigned to different
/// stages. This requires point-to-point communication between stages.
#[derive(Debug, Clone)]
pub struct PipelineParallel {
    /// Number of pipeline stages.
    pub num_stages: usize,
    /// This node's stage index.
    pub stage: usize,
    /// Layers assigned to each stage.
    pub layers_per_stage: Vec<usize>,
}

impl PipelineParallel {
    /// Create a new pipeline parallel configuration.
    #[must_use]
    pub fn new(num_stages: usize, stage: usize) -> Self {
        Self {
            num_stages,
            stage,
            layers_per_stage: vec![],
        }
    }

    /// Configure for a specific number of layers.
    pub fn configure(&mut self, num_layers: usize) {
        let base_layers = num_layers / self.num_stages;
        let extra = num_layers % self.num_stages;

        self.layers_per_stage = (0..self.num_stages)
            .map(|s| base_layers + usize::from(s < extra))
            .collect();
    }

    /// Get the layer range for this stage.
    #[must_use]
    pub fn layer_range(&self) -> (usize, usize) {
        if self.layers_per_stage.is_empty() {
            return (0, 0);
        }

        let start: usize = self.layers_per_stage[..self.stage].iter().sum();
        let end = start + self.layers_per_stage[self.stage];
        (start, end)
    }

    /// Get the previous stage (for receiving activations).
    #[must_use]
    pub fn prev_stage(&self) -> Option<usize> {
        if self.stage > 0 {
            Some(self.stage - 1)
        } else {
            None
        }
    }

    /// Get the next stage (for sending activations).
    #[must_use]
    pub fn next_stage(&self) -> Option<usize> {
        if self.stage < self.num_stages - 1 {
            Some(self.stage + 1)
        } else {
            None
        }
    }

    /// Check if this is the first stage.
    #[must_use]
    pub fn is_first_stage(&self) -> bool {
        self.stage == 0
    }

    /// Check if this is the last stage.
    #[must_use]
    pub fn is_last_stage(&self) -> bool {
        self.stage == self.num_stages - 1
    }

    /// Calculate activation transfer size in MB.
    #[must_use]
    pub fn activation_size_mb(
        &self,
        hidden_dim: usize,
        seq_len: usize,
        batch_size: usize,
        dtype_bytes: usize,
    ) -> f64 {
        let elements = hidden_dim * seq_len * batch_size;
        let bytes = elements * dtype_bytes;
        bytes as f64 / (1024.0 * 1024.0)
    }

    /// Calculate pipeline bubble ratio.
    #[must_use]
    pub fn bubble_ratio(&self, num_micro_batches: usize) -> f64 {
        let p = self.num_stages as f64;
        let m = num_micro_batches as f64;
        (p - 1.0) / (p + m - 1.0)
    }
}

// ============================================================================
// Layer Partitioner
// ============================================================================

/// Partitioner for optimal layer assignment across nodes.
#[derive(Debug)]
pub struct LayerPartitioner {
    /// Total number of layers.
    pub num_layers: usize,
    /// Total number of nodes.
    pub num_nodes: usize,
    /// Tensor parallelism degree.
    pub tensor_parallel: usize,
    /// Pipeline parallelism degree.
    pub pipeline_parallel: usize,
    /// Layer memory costs (optional).
    pub layer_costs: Option<Vec<f64>>,
}

impl LayerPartitioner {
    /// Create a new layer partitioner.
    #[must_use]
    pub fn new(
        num_layers: usize,
        num_nodes: usize,
        tensor_parallel: usize,
        pipeline_parallel: usize,
    ) -> Self {
        Self {
            num_layers,
            num_nodes,
            tensor_parallel,
            pipeline_parallel,
            layer_costs: None,
        }
    }

    /// Set custom layer memory costs.
    pub fn set_layer_costs(&mut self, costs: Vec<f64>) {
        self.layer_costs = Some(costs);
    }

    /// Partition layers across nodes.
    #[must_use]
    pub fn partition(&self) -> Vec<LayerAssignment> {
        let mut assignments = Vec::with_capacity(self.num_layers);

        // Calculate layers per pipeline stage
        let layers_per_stage = self.num_layers / self.pipeline_parallel;
        let extra_layers = self.num_layers % self.pipeline_parallel;

        let mut layer_idx = 0;
        for stage in 0..self.pipeline_parallel {
            let stage_layers = layers_per_stage + usize::from(stage < extra_layers);

            for local_layer in 0..stage_layers {
                // For tensor parallel, each layer is replicated across TP ranks
                // but we assign to the primary node (TP rank 0)
                let node_id = stage * self.tensor_parallel;

                let memory_mb = self
                    .layer_costs
                    .as_ref()
                    .map_or(1024.0, |c| c.get(layer_idx).copied().unwrap_or(1024.0));

                assignments.push(LayerAssignment {
                    layer_id: layer_idx,
                    node_id,
                    memory_mb,
                    pipeline_stage: stage,
                    tensor_rank: 0,
                });

                layer_idx += 1;
                let _ = local_layer; // Suppress unused warning
            }
        }

        assignments
    }

    /// Partition with load balancing based on layer costs.
    #[must_use]
    pub fn partition_balanced(&self) -> Vec<LayerAssignment> {
        if self.layer_costs.is_none() {
            return self.partition();
        }

        let costs = self.layer_costs.as_ref().unwrap();
        let total_cost: f64 = costs.iter().sum();
        let target_cost_per_stage = total_cost / self.pipeline_parallel as f64;

        let mut assignments = Vec::with_capacity(self.num_layers);
        let mut current_stage = 0;
        let mut current_stage_cost = 0.0;

        for (layer_idx, &cost) in costs.iter().enumerate() {
            let node_id = current_stage * self.tensor_parallel;

            assignments.push(LayerAssignment {
                layer_id: layer_idx,
                node_id,
                memory_mb: cost,
                pipeline_stage: current_stage,
                tensor_rank: 0,
            });

            current_stage_cost += cost;

            // Move to next stage if we've exceeded target and not on last stage
            if current_stage_cost >= target_cost_per_stage
                && current_stage < self.pipeline_parallel - 1
            {
                current_stage += 1;
                current_stage_cost = 0.0;
            }
        }

        assignments
    }

    /// Get memory usage per node.
    #[must_use]
    pub fn memory_per_node(&self, assignments: &[LayerAssignment]) -> Vec<f64> {
        let mut memory = vec![0.0; self.num_nodes];
        for a in assignments {
            if a.node_id < memory.len() {
                memory[a.node_id] += a.memory_mb;
            }
        }
        memory
    }

    /// Calculate load imbalance ratio.
    #[must_use]
    pub fn imbalance_ratio(&self, assignments: &[LayerAssignment]) -> f64 {
        let memory = self.memory_per_node(assignments);
        let max_mem = memory.iter().copied().fold(0.0, f64::max);
        let min_mem = memory.iter().copied().fold(f64::MAX, f64::min);
        let avg_mem = memory.iter().sum::<f64>() / memory.len() as f64;

        if avg_mem > 0.0 {
            (max_mem - min_mem) / avg_mem
        } else {
            0.0
        }
    }
}

// ============================================================================
// Micro-batch Schedule
// ============================================================================

/// Micro-batch schedule for pipeline execution.
#[derive(Debug, Clone)]
pub struct MicroBatchSchedule {
    /// Number of micro-batches.
    pub num_micro_batches: usize,
    /// Number of pipeline stages.
    pub num_stages: usize,
    /// Schedule entries (stage, micro_batch, is_forward).
    pub entries: Vec<(usize, usize, bool)>,
}

impl MicroBatchSchedule {
    /// Create a GPipe schedule (all forwards then all backwards).
    #[must_use]
    pub fn gpipe(num_stages: usize, num_micro_batches: usize) -> Self {
        let mut entries = Vec::new();

        // Forward passes
        for mb in 0..num_micro_batches {
            for stage in 0..num_stages {
                entries.push((stage, mb, true));
            }
        }

        // Backward passes
        for mb in (0..num_micro_batches).rev() {
            for stage in (0..num_stages).rev() {
                entries.push((stage, mb, false));
            }
        }

        Self {
            num_micro_batches,
            num_stages,
            entries,
        }
    }

    /// Create a 1F1B schedule (interleaved forward/backward).
    #[must_use]
    pub fn one_f_one_b(num_stages: usize, num_micro_batches: usize) -> Self {
        let mut entries = Vec::new();

        // Warmup: forward passes for first (num_stages) micro-batches
        for mb in 0..num_stages.min(num_micro_batches) {
            for stage in 0..num_stages {
                entries.push((stage, mb, true));
            }
        }

        // Steady state: alternate 1F1B
        for mb in num_stages..num_micro_batches {
            for stage in 0..num_stages {
                entries.push((stage, mb, true));
            }
            let backward_mb = mb - num_stages;
            for stage in (0..num_stages).rev() {
                entries.push((stage, backward_mb, false));
            }
        }

        // Cooldown: remaining backward passes
        for mb in (num_micro_batches.saturating_sub(num_stages)..num_micro_batches).rev() {
            for stage in (0..num_stages).rev() {
                entries.push((stage, mb, false));
            }
        }

        Self {
            num_micro_batches,
            num_stages,
            entries,
        }
    }

    /// Get schedule for a specific stage.
    #[must_use]
    pub fn stage_schedule(&self, stage: usize) -> Vec<(usize, bool)> {
        self.entries
            .iter()
            .filter(|(s, _, _)| *s == stage)
            .map(|(_, mb, fwd)| (*mb, *fwd))
            .collect()
    }

    /// Calculate bubble time slots.
    #[must_use]
    pub fn bubble_slots(&self) -> usize {
        // In ideal case with no bubbles, each stage would execute
        // 2 * num_micro_batches operations (fwd + bwd)
        let ideal = self.num_stages * self.num_micro_batches * 2;
        let actual = self.entries.len();
        actual.saturating_sub(ideal) / self.num_stages
    }
}

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

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

    #[test]
    fn test_tensor_parallel_creation() {
        let tp = TensorParallel::new(4, 1);
        assert_eq!(tp.world_size, 4);
        assert_eq!(tp.rank, 1);
    }

    #[test]
    fn test_tensor_parallel_head_range() {
        let tp = TensorParallel::new(4, 1);
        let (start, end) = tp.head_range(64);
        assert_eq!(start, 16);
        assert_eq!(end, 32);
    }

    #[test]
    fn test_tensor_parallel_ffn_range() {
        let tp = TensorParallel::new(4, 2);
        let (start, end) = tp.ffn_range(28672);
        assert_eq!(start, 14336);
        assert_eq!(end, 21504);
    }

    #[test]
    fn test_tensor_parallel_column_split() {
        let tp = TensorParallel::new(4, 0);
        let (rows, cols) = tp.column_split((4096, 16384));
        assert_eq!(rows, 4096);
        assert_eq!(cols, 4096);
    }

    #[test]
    fn test_pipeline_parallel_creation() {
        let pp = PipelineParallel::new(4, 1);
        assert_eq!(pp.num_stages, 4);
        assert_eq!(pp.stage, 1);
    }

    #[test]
    fn test_pipeline_parallel_configure() {
        let mut pp = PipelineParallel::new(4, 0);
        pp.configure(80);
        assert_eq!(pp.layers_per_stage.iter().sum::<usize>(), 80);
    }

    #[test]
    fn test_pipeline_parallel_layer_range() {
        let mut pp = PipelineParallel::new(4, 1);
        pp.configure(80);
        let (start, end) = pp.layer_range();
        assert_eq!(start, 20);
        assert_eq!(end, 40);
    }

    #[test]
    fn test_pipeline_parallel_neighbors() {
        let pp = PipelineParallel::new(4, 1);
        assert_eq!(pp.prev_stage(), Some(0));
        assert_eq!(pp.next_stage(), Some(2));

        let first = PipelineParallel::new(4, 0);
        assert!(first.is_first_stage());
        assert_eq!(first.prev_stage(), None);

        let last = PipelineParallel::new(4, 3);
        assert!(last.is_last_stage());
        assert_eq!(last.next_stage(), None);
    }

    #[test]
    fn test_pipeline_bubble_ratio() {
        let pp = PipelineParallel::new(4, 0);
        let ratio = pp.bubble_ratio(16);
        // (4-1)/(4+16-1) = 3/19 = 0.157...
        assert!((ratio - 0.158).abs() < 0.01);
    }

    #[test]
    fn test_layer_partitioner() {
        let partitioner = LayerPartitioner::new(80, 4, 1, 4);
        let assignments = partitioner.partition();
        assert_eq!(assignments.len(), 80);
    }

    #[test]
    fn test_layer_partitioner_stages() {
        let partitioner = LayerPartitioner::new(80, 4, 1, 4);
        let assignments = partitioner.partition();

        // First 20 layers should be stage 0
        assert!(assignments[..20].iter().all(|a| a.pipeline_stage == 0));
        // Next 20 should be stage 1
        assert!(assignments[20..40].iter().all(|a| a.pipeline_stage == 1));
    }

    #[test]
    fn test_layer_partitioner_balanced() {
        let mut partitioner = LayerPartitioner::new(8, 4, 1, 4);
        partitioner.set_layer_costs(vec![1.0, 2.0, 3.0, 4.0, 4.0, 3.0, 2.0, 1.0]);
        let assignments = partitioner.partition_balanced();
        assert_eq!(assignments.len(), 8);
    }

    #[test]
    fn test_memory_per_node() {
        let partitioner = LayerPartitioner::new(8, 2, 1, 2);
        let assignments = partitioner.partition();
        let memory = partitioner.memory_per_node(&assignments);
        assert_eq!(memory.len(), 2);
    }

    #[test]
    fn test_gpipe_schedule() {
        let schedule = MicroBatchSchedule::gpipe(4, 8);
        assert!(!schedule.entries.is_empty());

        // All forwards should come before backwards
        let first_backward = schedule
            .entries
            .iter()
            .position(|(_, _, fwd)| !fwd)
            .unwrap();
        let last_forward = schedule
            .entries
            .iter()
            .rposition(|(_, _, fwd)| *fwd)
            .unwrap();
        assert!(last_forward < first_backward);
    }

    #[test]
    fn test_1f1b_schedule() {
        let schedule = MicroBatchSchedule::one_f_one_b(4, 8);
        assert!(!schedule.entries.is_empty());
    }

    #[test]
    fn test_stage_schedule() {
        let schedule = MicroBatchSchedule::gpipe(4, 4);
        let stage_0 = schedule.stage_schedule(0);
        assert!(!stage_0.is_empty());
        // Stage 0 should have 4 forwards and 4 backwards
        let forwards = stage_0.iter().filter(|(_, fwd)| *fwd).count();
        let backwards = stage_0.iter().filter(|(_, fwd)| !*fwd).count();
        assert_eq!(forwards, 4);
        assert_eq!(backwards, 4);
    }
}