//! Parallelism strategies for distributed training
//!
//! This module implements various parallelism approaches including:
//! - Data Parallel (DP)
//! - Tensor Parallel (TP)
//! - Pipeline Parallel (PP)
//! - Fully Sharded Data Parallel (FSDP/ZeRO)

use crate::error::{DistributedError, Result};
use crate::group::ProcessGroup;
use rtx_tensor::Tensor;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Data Parallel implementation
pub struct DataParallel {
    /// Process group for data parallel communication
    pub process_group: ProcessGroup,
    /// Gradient accumulation steps
    pub accumulation_steps: usize,
    /// Current accumulation step
    current_step: usize,
    /// Whether gradients are currently being accumulated
    accumulating: bool,
}

impl DataParallel {
    /// Create new data parallel instance
    pub fn new(process_group: ProcessGroup, accumulation_steps: usize) -> Self {
        Self {
            process_group,
            accumulation_steps: accumulation_steps.max(1),
            current_step: 0,
            accumulating: false,
        }
    }

    /// Start gradient accumulation
    pub fn start_accumulation(&mut self) {
        self.accumulating = true;
        self.current_step = 0;
    }

    /// Accumulate gradients (called after backward pass)
    pub async fn accumulate_gradients(&mut self, gradients: &mut [Tensor]) -> Result<bool> {
        if !self.accumulating {
            return Err(DistributedError::parallelism(
                "DataParallel",
                "not accumulating - call start_accumulation() first",
            ));
        }

        self.current_step += 1;

        // Check if we should synchronize
        let should_sync = self.current_step >= self.accumulation_steps;

        if should_sync {
            self.sync_gradients(gradients).await?;
            self.accumulating = false;
            self.current_step = 0;
        }

        Ok(should_sync)
    }

    /// Synchronize gradients across all processes
    async fn sync_gradients(&self, gradients: &mut [Tensor]) -> Result<()> {
        use crate::comm::ReduceOp;

        for gradient in gradients.iter_mut() {
            // AllReduce to sum gradients across all processes
            self.process_group
                .all_reduce(gradient, ReduceOp::Sum)
                .await?;

            // Average the gradients
            let world_size = self.process_group.world_size() as f32;
            *gradient = gradient.div_scalar(world_size)?;
        }

        Ok(())
    }
}

/// Fully Sharded Data Parallel (FSDP) implementation
pub struct Fsdp {
    /// Process group for FSDP communication
    pub process_group: ProcessGroup,
    /// FSDP configuration
    pub config: FsdpConfig,
    /// Sharded parameters
    sharded_params: HashMap<String, ShardedParameter>,
    /// Memory usage statistics
    memory_stats: FsdpMemoryStats,
}

/// FSDP configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FsdpConfig {
    /// Sharding strategy
    pub sharding_strategy: ShardingStrategy,
    /// Minimum parameter size to shard (in elements)
    pub min_param_size: usize,
    /// CPU offloading enabled
    pub cpu_offload: bool,
    /// Mixed precision configuration
    pub mixed_precision: bool,
    /// Flatten parameters for sharding
    pub flatten_parameters: bool,
}

/// Sharding strategies for FSDP
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ShardingStrategy {
    /// Full sharding (ZeRO-3)
    FullShard,
    /// Shard gradients only (ZeRO-2)
    ShardGradOp,
    /// No sharding, replicate parameters
    NoShard,
}

/// Sharded parameter representation
#[derive(Debug)]
pub struct ShardedParameter {
    /// Parameter name/identifier
    pub name: String,
    /// Local shard of the parameter
    pub local_shard: Tensor,
    /// Full parameter shape
    pub full_shape: Vec<usize>,
    /// Shard metadata
    pub shard_metadata: ShardMetadata,
}

/// Metadata for parameter sharding
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ShardMetadata {
    /// Rank that owns this shard
    pub owner_rank: i32,
    /// Start index in flattened parameter
    pub start_idx: usize,
    /// End index in flattened parameter
    pub end_idx: usize,
    /// Original parameter offset
    pub param_offset: usize,
}

/// Memory usage statistics for FSDP
#[derive(Debug, Default)]
pub struct FsdpMemoryStats {
    /// Total parameter memory (MB)
    pub total_param_memory_mb: f32,
    /// Local shard memory (MB)
    pub local_shard_memory_mb: f32,
    /// Peak memory during all-gather (MB)
    pub peak_allgather_memory_mb: f32,
    /// Memory saved compared to non-sharded (MB)
    pub memory_saved_mb: f32,
}

impl Default for FsdpConfig {
    fn default() -> Self {
        Self {
            sharding_strategy: ShardingStrategy::FullShard,
            min_param_size: 1000, // Only shard parameters with >1K elements
            cpu_offload: false,
            mixed_precision: true,
            flatten_parameters: true,
        }
    }
}

impl Fsdp {
    /// Create new FSDP instance with default config
    pub fn new(process_group: ProcessGroup) -> Result<Self> {
        let config = FsdpConfig::default();
        Ok(Self {
            process_group,
            config,
            sharded_params: HashMap::new(),
            memory_stats: FsdpMemoryStats::default(),
        })
    }

    /// Create new FSDP instance with custom config
    pub fn new_with_config(process_group: ProcessGroup, config: FsdpConfig) -> Self {
        Self {
            process_group,
            config,
            sharded_params: HashMap::new(),
            memory_stats: FsdpMemoryStats::default(),
        }
    }

    /// Shard model parameters across the process group
    pub fn shard_parameters(&mut self, parameters: &Tensor) -> Result<Tensor> {
        let world_size = self.process_group.world_size();
        let rank = self.process_group.rank();

        if world_size == 1 {
            // No sharding needed for single GPU
            return Ok(parameters.clone());
        }

        // Calculate shard size
        let total_elements = parameters.numel();
        let elements_per_shard = total_elements.div_ceil(world_size); // Round up
        let start_idx = rank * elements_per_shard;
        let end_idx = (start_idx + elements_per_shard).min(total_elements);
        let shard_size = end_idx - start_idx;

        // Create local shard (simplified - in real implementation would slice the actual tensor)
        let shard_shape = crate::TensorShape::new(vec![shard_size])?;
        let local_shard = Tensor::zeros(shard_shape, &crate::Device::default())?;

        // Update memory statistics
        let total_memory_mb = (total_elements * 4) as f32 / (1024.0 * 1024.0); // f32 = 4 bytes
        let shard_memory_mb = (shard_size * 4) as f32 / (1024.0 * 1024.0);

        self.memory_stats.total_param_memory_mb = total_memory_mb;
        self.memory_stats.local_shard_memory_mb = shard_memory_mb;
        self.memory_stats.memory_saved_mb = total_memory_mb - shard_memory_mb;

        tracing::debug!(
            "FSDP sharded {:.2}MB -> {:.2}MB ({:.1}% reduction)",
            total_memory_mb,
            shard_memory_mb,
            (self.memory_stats.memory_saved_mb / total_memory_mb) * 100.0
        );

        Ok(local_shard)
    }

    /// Synchronize gradients across all processes
    pub async fn sync_gradients(&self, gradients: &mut Tensor) -> Result<()> {
        use crate::comm::ReduceOp;

        // Perform AllReduce to sum gradients across all processes
        self.process_group
            .all_reduce(gradients, ReduceOp::Sum)
            .await?;

        // Average the gradients
        let world_size = self.process_group.world_size() as f32;
        *gradients = gradients.div_scalar(world_size)?;

        Ok(())
    }

    /// Get current memory statistics
    pub fn memory_stats(&self) -> &FsdpMemoryStats {
        &self.memory_stats
    }

    /// Calculate memory reduction percentage
    pub fn memory_reduction_percent(&self) -> f32 {
        if self.memory_stats.total_param_memory_mb > 0.0 {
            (self.memory_stats.memory_saved_mb / self.memory_stats.total_param_memory_mb) * 100.0
        } else {
            0.0
        }
    }
}

/// Tensor Parallel implementation
pub struct TensorParallel {
    /// Process group for tensor parallel communication
    pub process_group: ProcessGroup,
}

impl TensorParallel {
    /// Create new tensor parallel instance
    pub fn new(process_group: ProcessGroup) -> Result<Self> {
        Ok(Self { process_group })
    }

    /// Parallel matrix multiplication using [`ColParallelLinear`].
    ///
    /// `a` is `[batch, in_features]` and `b` is the weight matrix
    /// `[out_features, in_features]` in row-major order.  The method shards
    /// with `tp_size = 1` (this rank is the sole rank) so the result equals a
    /// full, un-sharded matmul.
    pub fn matmul(&self, a: &Tensor, b: &Tensor) -> Result<Tensor> {
        let a_dims = a.shape().dims();
        let b_dims = b.shape().dims();

        if a_dims.len() != 2 || b_dims.len() != 2 {
            return Err(DistributedError::parallelism(
                "TensorParallel::matmul",
                "expected 2-D tensors",
            ));
        }

        let batch = a_dims[0];
        let in_features = a_dims[1];
        let out_features = b_dims[0];

        if b_dims[1] != in_features {
            return Err(DistributedError::parallelism(
                "TensorParallel::matmul",
                format!(
                    "dimension mismatch: a has in_features={in_features} but b has inner dim={}",
                    b_dims[1]
                ),
            ));
        }

        let a_data = a.to_vec().map_err(|e| {
            DistributedError::parallelism("TensorParallel::matmul", format!("to_vec a: {e}"))
        })?;
        let b_data = b.to_vec().map_err(|e| {
            DistributedError::parallelism("TensorParallel::matmul", format!("to_vec b: {e}"))
        })?;

        // Use a single-rank ColParallelLinear (tp_size = 1) for the full matmul.
        let col = ColParallelLinear::new(&b_data, None, in_features, out_features, 1, 0)?;
        let output = col.forward_cpu(&a_data, batch);

        tracing::debug!(
            "Tensor parallel matmul: {batch}x{in_features} * {out_features}x{in_features}"
        );

        Tensor::from_vec(output, &[batch, out_features], &crate::Device::default()).map_err(|e| {
            DistributedError::parallelism("TensorParallel::matmul", format!("from_vec: {e}"))
        })
    }
}

// ---------------------------------------------------------------------------
// Column-parallel linear layer
// ---------------------------------------------------------------------------

/// Column-parallel linear layer for tensor parallelism.
///
/// The full weight matrix `[out_features, in_features]` is partitioned along
/// the output (row) dimension.  Each TP rank holds
/// `[out_features / tp_size, in_features]` rows.
///
/// `forward_cpu(x)` computes `x @ weight_shard.T` locally.  Because each
/// rank produces a non-overlapping output slice, the sharded outputs can be
/// concatenated (AllGather) by the caller, or fed directly into a
/// [`RowParallelLinear`] which will AllReduce the partial sums.
///
/// # Invariants
/// * `weight_full.len() == out_features * in_features`
/// * `out_features % tp_size == 0`
/// * `0 <= rank < tp_size`
pub struct ColParallelLinear {
    /// Weight shard: `[out_features / tp_size, in_features]`, row-major.
    pub weight_shard: Vec<f32>,
    /// Optional bias shard: `[out_features / tp_size]`.
    pub bias_shard: Option<Vec<f32>>,
    /// Inner dimension (columns of the weight matrix).
    pub in_features: usize,
    /// Rows owned by this rank (`out_features / tp_size`).
    pub out_features_per_rank: usize,
    /// Total tensor-parallel degree.
    pub tp_size: usize,
    /// This rank's index within the TP group (`0..tp_size`).
    pub rank: usize,
}

impl ColParallelLinear {
    /// Construct a column-parallel linear layer from the full weight matrix.
    ///
    /// # Arguments
    /// * `weight_full`  — flat, row-major `[out_features, in_features]`.
    /// * `bias_full`    — optional flat `[out_features]` bias.
    /// * `in_features`  — inner dimension K.
    /// * `out_features` — outer dimension N (must be divisible by `tp_size`).
    /// * `tp_size`      — number of tensor-parallel ranks.
    /// * `rank`         — this rank's index (`0..tp_size`).
    ///
    /// # Errors
    /// Returns [`DistributedError`] if `out_features` is not evenly divisible
    /// by `tp_size`, if the weight slice length is wrong, or if `rank >= tp_size`.
    pub fn new(
        weight_full: &[f32],
        bias_full: Option<&[f32]>,
        in_features: usize,
        out_features: usize,
        tp_size: usize,
        rank: usize,
    ) -> Result<Self> {
        if tp_size == 0 {
            return Err(DistributedError::parallelism(
                "ColParallelLinear",
                "tp_size must be >= 1",
            ));
        }
        if rank >= tp_size {
            return Err(DistributedError::parallelism(
                "ColParallelLinear",
                format!("rank {rank} is out of range for tp_size {tp_size}"),
            ));
        }
        if out_features % tp_size != 0 {
            return Err(DistributedError::parallelism(
                "ColParallelLinear",
                format!("out_features {out_features} must be divisible by tp_size {tp_size}"),
            ));
        }

        let expected = out_features * in_features;
        if weight_full.len() != expected {
            return Err(DistributedError::parallelism(
                "ColParallelLinear",
                format!(
                    "weight_full length {} does not match out_features*in_features={expected}",
                    weight_full.len()
                ),
            ));
        }

        let shard_rows = out_features / tp_size;
        let row_start = rank * shard_rows;
        let row_end = row_start + shard_rows;

        // Extract rows [row_start, row_end) from the weight matrix.
        let weight_shard = weight_full[row_start * in_features..row_end * in_features].to_vec();

        let bias_shard = bias_full.map(|bias| bias[row_start..row_end].to_vec());

        Ok(Self {
            weight_shard,
            bias_shard,
            in_features,
            out_features_per_rank: shard_rows,
            tp_size,
            rank,
        })
    }

    /// CPU reference forward pass.
    ///
    /// Computes `x [batch, in_features] @ weight_shard.T`
    /// → `[batch, out_features_per_rank]`.
    /// Adds `bias_shard` if present.
    pub fn forward_cpu(&self, x: &[f32], batch: usize) -> Vec<f32> {
        let n = self.out_features_per_rank;
        let k = self.in_features;
        let mut out = vec![0.0_f32; batch * n];

        for b in 0..batch {
            for j in 0..n {
                let mut acc = 0.0_f32;
                for i in 0..k {
                    acc += x[b * k + i] * self.weight_shard[j * k + i];
                }
                out[b * n + j] = acc + self.bias_shard.as_ref().map_or(0.0, |bias| bias[j]);
            }
        }

        out
    }
}

// ---------------------------------------------------------------------------
// Row-parallel linear layer
// ---------------------------------------------------------------------------

/// Row-parallel linear layer for tensor parallelism.
///
/// The full weight matrix `[out_features, in_features]` is partitioned along
/// the input (column) dimension.  Each TP rank holds
/// `[out_features, in_features / tp_size]` columns.
///
/// Each rank receives a matching input shard `x_shard [batch, in_features/tp_size]`
/// (typically the output of a [`ColParallelLinear`]) and computes
/// `x_shard @ weight_shard.T` locally to produce a partial sum
/// `[batch, out_features]`.  A subsequent AllReduce over all TP ranks yields
/// the full output.
///
/// # Invariants
/// * `weight_full.len() == out_features * in_features`
/// * `in_features % tp_size == 0`
/// * `0 <= rank < tp_size`
pub struct RowParallelLinear {
    /// Weight shard: `[out_features, in_features / tp_size]`, row-major.
    pub weight_shard: Vec<f32>,
    /// Bias applied after AllReduce — **only rank 0 adds it** to avoid
    /// double-counting during the reduce.  Shape: `[out_features]`.
    pub bias: Option<Vec<f32>>,
    /// Output feature count (unchanged by column-sharding).
    pub out_features: usize,
    /// Columns owned by this rank (`in_features / tp_size`).
    pub in_features_per_rank: usize,
    /// Total tensor-parallel degree.
    pub tp_size: usize,
    /// This rank's index within the TP group (`0..tp_size`).
    pub rank: usize,
    /// Process group used for the AllReduce collective.
    pub process_group: ProcessGroup,
}

impl RowParallelLinear {
    /// Construct a row-parallel linear layer from the full weight matrix.
    ///
    /// # Arguments
    /// * `weight_full`  — flat, row-major `[out_features, in_features]`.
    /// * `bias`         — optional flat `[out_features]` bias (owned; only
    ///                    rank 0 applies it after AllReduce).
    /// * `out_features` — output dimension M.
    /// * `in_features`  — inner dimension K (must be divisible by `tp_size`).
    /// * `tp_size`      — number of tensor-parallel ranks.
    /// * `rank`         — this rank's index (`0..tp_size`).
    /// * `process_group`— process group for the AllReduce collective.
    ///
    /// # Errors
    /// Returns [`DistributedError`] if `in_features` is not evenly divisible
    /// by `tp_size`, if the weight slice length is wrong, or if `rank >= tp_size`.
    pub fn new(
        weight_full: &[f32],
        bias: Option<Vec<f32>>,
        out_features: usize,
        in_features: usize,
        tp_size: usize,
        rank: usize,
        process_group: ProcessGroup,
    ) -> Result<Self> {
        if tp_size == 0 {
            return Err(DistributedError::parallelism(
                "RowParallelLinear",
                "tp_size must be >= 1",
            ));
        }
        if rank >= tp_size {
            return Err(DistributedError::parallelism(
                "RowParallelLinear",
                format!("rank {rank} is out of range for tp_size {tp_size}"),
            ));
        }
        if in_features % tp_size != 0 {
            return Err(DistributedError::parallelism(
                "RowParallelLinear",
                format!("in_features {in_features} must be divisible by tp_size {tp_size}"),
            ));
        }

        let expected = out_features * in_features;
        if weight_full.len() != expected {
            return Err(DistributedError::parallelism(
                "RowParallelLinear",
                format!(
                    "weight_full length {} does not match out_features*in_features={expected}",
                    weight_full.len()
                ),
            ));
        }

        let shard_cols = in_features / tp_size;
        let col_start = rank * shard_cols;
        let col_end = col_start + shard_cols;

        // Extract columns [col_start, col_end) from every row of the weight matrix.
        // weight_full is row-major [out_features, in_features], so row i spans
        // indices [i*in_features .. (i+1)*in_features].
        let mut weight_shard = Vec::with_capacity(out_features * shard_cols);
        for row in 0..out_features {
            let row_base = row * in_features;
            weight_shard.extend_from_slice(&weight_full[row_base + col_start..row_base + col_end]);
        }

        Ok(Self {
            weight_shard,
            bias,
            out_features,
            in_features_per_rank: shard_cols,
            tp_size,
            rank,
            process_group,
        })
    }

    /// CPU reference forward pass (no network I/O).
    ///
    /// Computes the local partial sum
    /// `x_shard [batch, in_features/tp_size] @ weight_shard.T`
    /// → `[batch, out_features]`.
    ///
    /// Bias is added **only on rank 0** so that, after an AllReduce sum, the
    /// bias is incorporated exactly once in the final result.
    ///
    /// In production this is followed by
    /// `process_group.all_reduce(&mut partial, ReduceOp::Sum)`.
    pub fn forward_cpu(&self, x_shard: &[f32], batch: usize) -> Vec<f32> {
        let n = self.out_features;
        let k = self.in_features_per_rank;
        let mut partial = vec![0.0_f32; batch * n];

        for b in 0..batch {
            for j in 0..n {
                let mut acc = 0.0_f32;
                for i in 0..k {
                    acc += x_shard[b * k + i] * self.weight_shard[j * k + i];
                }
                partial[b * n + j] = acc;
            }
        }

        // Add bias only on rank 0 to avoid double-counting across TP ranks.
        if self.rank == 0 {
            if let Some(bias) = &self.bias {
                for b in 0..batch {
                    for j in 0..n {
                        partial[b * n + j] += bias[j];
                    }
                }
            }
        }

        partial
    }

    /// Async forward pass with a real AllReduce collective.
    ///
    /// 1. Computes the local partial sum via [`forward_cpu`][Self::forward_cpu].
    /// 2. Wraps the result in a [`Tensor`] and issues an AllReduce Sum across
    ///    all TP ranks via `self.process_group`.
    /// 3. Returns the reduced flat `[batch * out_features]` buffer.
    ///
    /// In CPU simulation mode (`Backend::Cpu`, single process) the AllReduce
    /// is a no-op that scales by `world_size`; the returned data is therefore
    /// correct only when `tp_size == 1` or when multiple real processes run.
    /// Use [`forward_cpu`][Self::forward_cpu] directly for unit tests that
    /// manually sum partials.
    pub async fn forward(&self, x_shard: &[f32], batch: usize) -> Result<Vec<f32>> {
        use crate::comm::ReduceOp;

        let local = self.forward_cpu(x_shard, batch);

        // Wrap in a Tensor for the collective API.
        let mut t = Tensor::from_vec(
            local,
            &[batch, self.out_features],
            &crate::Device::default(),
        )
        .map_err(|e| {
            DistributedError::parallelism("RowParallelLinear::forward", format!("from_vec: {e}"))
        })?;

        // AllReduce: sum partial results across all TP ranks.
        self.process_group.all_reduce(&mut t, ReduceOp::Sum).await?;

        t.to_vec().map_err(|e| {
            DistributedError::parallelism(
                "RowParallelLinear::forward",
                format!("to_vec after allreduce: {e}"),
            )
        })
    }
}

/// Pipeline Parallel implementation
pub struct PipelineParallel {
    /// Process group for pipeline parallel communication
    pub process_group: ProcessGroup,
    /// Number of pipeline stages
    pub num_stages: usize,
}

impl PipelineParallel {
    /// Create new pipeline parallel instance
    pub fn new(process_group: ProcessGroup, num_stages: usize) -> Result<Self> {
        Ok(Self {
            process_group,
            num_stages,
        })
    }

    /// Forward pass through pipeline stages
    pub fn forward(&self, inputs: &Tensor) -> Result<Tensor> {
        // Simplified pipeline forward pass
        // In real implementation, this would:
        // 1. Split model into stages
        // 2. Pass activations between stages
        // 3. Handle micro-batching

        let outputs = inputs.clone();

        tracing::debug!(
            "Pipeline forward: {} stages, input shape: {:?}",
            self.num_stages,
            inputs.shape().dims()
        );

        Ok(outputs)
    }

    /// Backward pass through pipeline stages
    pub fn backward(&self, grad_outputs: &Tensor) -> Result<Tensor> {
        // Simplified pipeline backward pass
        let grad_inputs = grad_outputs.clone();

        tracing::debug!(
            "Pipeline backward: {} stages, grad_output shape: {:?}",
            self.num_stages,
            grad_outputs.shape().dims()
        );

        Ok(grad_inputs)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::group::ProcessGroup;
    use crate::{Backend, BackendConfig, Device, Tensor, TensorShape};

    #[tokio::test]
    async fn test_data_parallel() {
        let config = BackendConfig::cpu();
        let pg = ProcessGroup::new_with_config(Backend::Cpu, 2, 0, config)
            .await
            .unwrap();
        let mut dp = DataParallel::new(pg, 2);

        dp.start_accumulation();

        let mut gradients = vec![
            Tensor::ones(&TensorShape::new(vec![10, 10]).unwrap(), &Device::default()).unwrap(),
            Tensor::ones(&TensorShape::new(vec![5, 5]).unwrap(), &Device::default()).unwrap(),
        ];

        let should_sync = dp.accumulate_gradients(&mut gradients).await.unwrap();
        assert!(!should_sync); // First accumulation shouldn't sync

        let should_sync = dp.accumulate_gradients(&mut gradients).await.unwrap();
        assert!(should_sync); // Second accumulation should sync
    }

    #[test]
    fn test_fsdp_config() {
        let config = FsdpConfig::default();
        assert_eq!(config.sharding_strategy, ShardingStrategy::FullShard);
        assert_eq!(config.min_param_size, 1000);
        assert!(config.mixed_precision);
        assert!(config.flatten_parameters);
        assert!(!config.cpu_offload);
    }
}

// ---------------------------------------------------------------------------
// Tensor-parallel unit tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tp_tests {
    use super::*;
    use crate::{Backend, BackendConfig, Device, Tensor, TensorShape};

    // ------------------------------------------------------------------
    // Helpers
    // ------------------------------------------------------------------

    /// Naive CPU matmul: A [batch, k] @ W.T [k, n] → [batch, n].
    /// Weight W is stored row-major as [n, k].
    fn naive_matmul(a: &[f32], w: &[f32], batch: usize, k: usize, n: usize) -> Vec<f32> {
        let mut out = vec![0.0_f32; batch * n];
        for b in 0..batch {
            for j in 0..n {
                let mut acc = 0.0_f32;
                for i in 0..k {
                    acc += a[b * k + i] * w[j * k + i];
                }
                out[b * n + j] = acc;
            }
        }
        out
    }

    fn assert_vec_approx(got: &[f32], expected: &[f32], tol: f32, label: &str) {
        assert_eq!(
            got.len(),
            expected.len(),
            "{label}: length mismatch: got {} expected {}",
            got.len(),
            expected.len()
        );
        for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
            assert!(
                (g - e).abs() <= tol,
                "{label}: element[{i}] got={g} expected={e} diff={}",
                (g - e).abs()
            );
        }
    }

    fn cpu_pg() -> ProcessGroup {
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            let config = BackendConfig::cpu();
            ProcessGroup::new_with_config(Backend::Cpu, 1, 0, config)
                .await
                .unwrap()
        })
    }

    // ------------------------------------------------------------------
    // ColParallelLinear tests
    // ------------------------------------------------------------------

    #[test]
    fn test_col_parallel_single_rank() {
        // tp_size = 1 → weight_shard == weight_full; output matches naive matmul.
        let batch = 2;
        let in_f = 3;
        let out_f = 4;

        // Weight [out_f, in_f]
        let w: Vec<f32> = (1..=(out_f * in_f) as i32).map(|x| x as f32).collect();
        // Input [batch, in_f]
        let x: Vec<f32> = (1..=(batch * in_f) as i32).map(|x| x as f32).collect();

        let col = ColParallelLinear::new(&w, None, in_f, out_f, 1, 0).unwrap();
        assert_eq!(col.weight_shard, w, "single rank should hold full weight");
        assert_eq!(col.out_features_per_rank, out_f);

        let got = col.forward_cpu(&x, batch);
        let expected = naive_matmul(&x, &w, batch, in_f, out_f);
        assert_vec_approx(&got, &expected, 1e-5, "col_single_rank");
    }

    #[test]
    fn test_col_parallel_two_rank_sharding() {
        // out_features = 4, tp_size = 2:
        //   rank 0 → rows [0, 1]  (out cols 0,1 of full result)
        //   rank 1 → rows [2, 3]  (out cols 2,3 of full result)
        // Concatenated outputs must equal the full matmul.
        let batch = 2;
        let in_f = 3;
        let out_f = 4;

        let w: Vec<f32> = (1..=(out_f * in_f) as i32).map(|x| x as f32).collect();
        let x: Vec<f32> = (1..=(batch * in_f) as i32).map(|x| x as f32).collect();

        let col0 = ColParallelLinear::new(&w, None, in_f, out_f, 2, 0).unwrap();
        let col1 = ColParallelLinear::new(&w, None, in_f, out_f, 2, 1).unwrap();

        assert_eq!(col0.out_features_per_rank, 2);
        assert_eq!(col1.out_features_per_rank, 2);

        let out0 = col0.forward_cpu(&x, batch); // [batch, 2]
        let out1 = col1.forward_cpu(&x, batch); // [batch, 2]

        // Interleave: for each batch row, concatenate out0[b] then out1[b] → [batch, 4]
        let mut combined = Vec::with_capacity(batch * out_f);
        for b in 0..batch {
            combined.extend_from_slice(&out0[b * 2..(b + 1) * 2]);
            combined.extend_from_slice(&out1[b * 2..(b + 1) * 2]);
        }

        let expected = naive_matmul(&x, &w, batch, in_f, out_f);
        assert_vec_approx(&combined, &expected, 1e-5, "col_two_rank_sharding");
    }

    #[test]
    fn test_col_parallel_with_bias() {
        // Bias is added to the output slice owned by this rank.
        let batch = 2;
        let in_f = 3;
        let out_f = 4;

        let w: Vec<f32> = vec![1.0; out_f * in_f];
        let x: Vec<f32> = vec![1.0; batch * in_f]; // all-ones → each output is in_f
        let bias: Vec<f32> = vec![10.0, 20.0, 30.0, 40.0]; // one per output feature

        // tp_size = 1: rank 0 gets all rows and the full bias.
        let col = ColParallelLinear::new(&w, Some(&bias), in_f, out_f, 1, 0).unwrap();
        let got = col.forward_cpu(&x, batch);

        // Each output element should be in_f * 1.0 + bias[j] = 3 + bias[j]
        let expected: Vec<f32> = (0..batch)
            .flat_map(|_| bias.iter().map(|b| in_f as f32 + b))
            .collect();
        assert_vec_approx(&got, &expected, 1e-5, "col_with_bias");
    }

    #[test]
    fn test_col_parallel_bias_two_ranks() {
        // With tp_size = 2: rank 0 gets bias[0..2], rank 1 gets bias[2..4].
        let in_f = 2;
        let out_f = 4;
        let batch = 1;

        let w: Vec<f32> = vec![1.0; out_f * in_f];
        let x: Vec<f32> = vec![1.0; in_f]; // [1, in_f]
        let bias: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0];

        let col0 = ColParallelLinear::new(&w, Some(&bias), in_f, out_f, 2, 0).unwrap();
        let col1 = ColParallelLinear::new(&w, Some(&bias), in_f, out_f, 2, 1).unwrap();

        let out0 = col0.forward_cpu(&x, batch); // expects [in_f + 1, in_f + 2] = [3,4]
        let out1 = col1.forward_cpu(&x, batch); // expects [in_f + 3, in_f + 4] = [5,6]

        assert_vec_approx(&out0, &[3.0, 4.0], 1e-5, "bias_rank0");
        assert_vec_approx(&out1, &[5.0, 6.0], 1e-5, "bias_rank1");
    }

    // ------------------------------------------------------------------
    // RowParallelLinear tests
    // ------------------------------------------------------------------

    #[test]
    fn test_row_parallel_single_rank() {
        // tp_size = 1 → forward_cpu equals naive matmul.
        let batch = 2;
        let in_f = 4;
        let out_f = 3;

        let w: Vec<f32> = (1..=(out_f * in_f) as i32).map(|x| x as f32).collect();
        let x: Vec<f32> = (1..=(batch * in_f) as i32).map(|x| x as f32).collect();

        let pg = cpu_pg();
        let row = RowParallelLinear::new(&w, None, out_f, in_f, 1, 0, pg).unwrap();
        assert_eq!(row.in_features_per_rank, in_f);

        let got = row.forward_cpu(&x, batch);
        let expected = naive_matmul(&x, &w, batch, in_f, out_f);
        assert_vec_approx(&got, &expected, 1e-5, "row_single_rank");
    }

    #[test]
    fn test_row_parallel_two_rank_sharding() {
        // in_features = 4, tp_size = 2:
        //   rank 0 → cols [0,1] of W  →  partial sum from x[0..2]
        //   rank 1 → cols [2,3] of W  →  partial sum from x[2..4]
        // Adding partial sums → full matmul output.
        let batch = 2;
        let in_f = 4;
        let out_f = 3;

        let w: Vec<f32> = (1..=(out_f * in_f) as i32).map(|x| x as f32).collect();
        let x: Vec<f32> = (1..=(batch * in_f) as i32).map(|x| x as f32).collect();

        let pg0 = cpu_pg();
        let pg1 = cpu_pg();

        let row0 = RowParallelLinear::new(&w, None, out_f, in_f, 2, 0, pg0).unwrap();
        let row1 = RowParallelLinear::new(&w, None, out_f, in_f, 2, 1, pg1).unwrap();

        // x_shard for each rank: columns [0..2] and [2..4] of x per batch row.
        let shard_size = in_f / 2;
        let mut x_shard0 = Vec::with_capacity(batch * shard_size);
        let mut x_shard1 = Vec::with_capacity(batch * shard_size);
        for b in 0..batch {
            x_shard0.extend_from_slice(&x[b * in_f..b * in_f + shard_size]);
            x_shard1.extend_from_slice(&x[b * in_f + shard_size..b * in_f + in_f]);
        }

        let partial0 = row0.forward_cpu(&x_shard0, batch);
        let partial1 = row1.forward_cpu(&x_shard1, batch);

        // AllReduce (sum): add the two partials element-wise.
        let combined: Vec<f32> = partial0
            .iter()
            .zip(partial1.iter())
            .map(|(a, b)| a + b)
            .collect();

        let expected = naive_matmul(&x, &w, batch, in_f, out_f);
        assert_vec_approx(&combined, &expected, 1e-5, "row_two_rank_sharding");
    }

    #[test]
    fn test_row_parallel_bias_rank0_only() {
        // Bias is added only by rank 0; rank 1 must not add it.
        let batch = 1;
        let in_f = 2;
        let out_f = 2;

        // Identity-ish weight [out_f, in_f] = [[1,0],[0,1]]
        let w: Vec<f32> = vec![1.0, 0.0, 0.0, 1.0];
        let bias: Vec<f32> = vec![100.0, 200.0];
        let x: Vec<f32> = vec![3.0, 5.0]; // [1, 2]

        let pg0 = cpu_pg();
        let pg1 = cpu_pg();

        // rank 0 gets cols [0..1], rank 1 gets cols [1..2].
        let row0 = RowParallelLinear::new(&w, Some(bias.clone()), out_f, in_f, 2, 0, pg0).unwrap();
        let row1 = RowParallelLinear::new(&w, Some(bias.clone()), out_f, in_f, 2, 1, pg1).unwrap();

        // x_shard0 = [3.0], x_shard1 = [5.0]
        let partial0 = row0.forward_cpu(&[3.0], batch);
        let partial1 = row1.forward_cpu(&[5.0], batch);

        // partial0 = [3*1 + bias[0], 3*0 + bias[1]] = [103, 200]
        // partial1 = [5*0,            5*1           ] = [  0,   5]
        assert_vec_approx(&partial0, &[103.0, 200.0], 1e-5, "partial_rank0");
        assert_vec_approx(&partial1, &[0.0, 5.0], 1e-5, "partial_rank1");

        // AllReduce (sum) → [103, 205]  = expected [3+100, 5+200]
        let summed: Vec<f32> = partial0
            .iter()
            .zip(partial1.iter())
            .map(|(a, b)| a + b)
            .collect();
        assert_vec_approx(&summed, &[103.0, 205.0], 1e-5, "bias_rank0_only_final");
    }

    // ------------------------------------------------------------------
    // TensorParallel::matmul
    // ------------------------------------------------------------------

    #[test]
    fn test_matmul_not_zeros() {
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            let config = BackendConfig::cpu();
            let pg = ProcessGroup::new_with_config(Backend::Cpu, 1, 0, config)
                .await
                .unwrap();
            let tp = TensorParallel::new(pg).unwrap();

            let batch = 2;
            let in_f = 3;
            let out_f = 4;

            let a_data: Vec<f32> = (1..=(batch * in_f) as i32).map(|x| x as f32).collect();
            let b_data: Vec<f32> = (1..=(out_f * in_f) as i32).map(|x| x as f32).collect();

            let a = Tensor::from_vec(a_data.clone(), &[batch, in_f], &Device::default()).unwrap();
            let b = Tensor::from_vec(b_data.clone(), &[out_f, in_f], &Device::default()).unwrap();

            let result = tp.matmul(&a, &b).unwrap();
            let result_data = result.to_vec().unwrap();

            // Verify output is non-zero.
            let all_zero = result_data.iter().all(|&v| v == 0.0);
            assert!(!all_zero, "matmul output must not be all zeros");

            // Verify correctness against naive matmul.
            let expected = naive_matmul(&a_data, &b_data, batch, in_f, out_f);
            assert_eq!(result.dims(), &[batch, out_f]);
            assert_vec_approx(&result_data, &expected, 1e-4, "tp_matmul_correctness");
        });
    }

    // ------------------------------------------------------------------
    // ColParallel → RowParallel round-trip
    // ------------------------------------------------------------------

    #[test]
    fn test_col_then_row_roundtrip() {
        // Simulate the standard TP linear pipeline:
        //   x [batch, in_f]  →  ColParallel  →  x_shard [batch, in_f/tp]
        //   x_shard  →  RowParallel  →  partial [batch, out_f]
        //   AllReduce (sum of tp partials)  →  [batch, out_f]
        //
        // With tp_size = 2 this is a genuine two-stage sharded pipeline.
        // We run both ranks in the same process (CPU sim) and manually add
        // their partial outputs to mimic the AllReduce.
        let batch = 3;
        let in_f = 4;
        let hidden = 6; // ColParallel output / RowParallel input
        let out_f = 5;

        // Two weight matrices simulating a two-layer MLP.
        let w1: Vec<f32> = (1..=(hidden * in_f) as i32)
            .map(|x| x as f32 * 0.1)
            .collect();
        let w2: Vec<f32> = (1..=(out_f * hidden) as i32)
            .map(|x| x as f32 * 0.05)
            .collect();
        let x: Vec<f32> = (1..=(batch * in_f) as i32).map(|x| x as f32).collect();

        let tp_size = 2;

        // --- Stage 1: ColParallel on w1 (shard along hidden / output dim) ---
        let col0 = ColParallelLinear::new(&w1, None, in_f, hidden, tp_size, 0).unwrap();
        let col1 = ColParallelLinear::new(&w1, None, in_f, hidden, tp_size, 1).unwrap();

        let h_shard0 = col0.forward_cpu(&x, batch); // [batch, hidden/2]
        let h_shard1 = col1.forward_cpu(&x, batch); // [batch, hidden/2]

        // --- Stage 2: RowParallel on w2 (shard along hidden / input dim) ---
        // Each rank receives the corresponding column shard from Stage 1.
        let pg0 = cpu_pg();
        let pg1 = cpu_pg();
        let row0 = RowParallelLinear::new(&w2, None, out_f, hidden, tp_size, 0, pg0).unwrap();
        let row1 = RowParallelLinear::new(&w2, None, out_f, hidden, tp_size, 1, pg1).unwrap();

        let partial0 = row0.forward_cpu(&h_shard0, batch);
        let partial1 = row1.forward_cpu(&h_shard1, batch);

        // AllReduce (sum) → [batch, out_f]
        let combined: Vec<f32> = partial0
            .iter()
            .zip(partial1.iter())
            .map(|(a, b)| a + b)
            .collect();

        // Reference: single-rank two-layer matmul.
        let h_ref = naive_matmul(&x, &w1, batch, in_f, hidden);
        let out_ref = naive_matmul(&h_ref, &w2, batch, hidden, out_f);

        assert_eq!(combined.len(), batch * out_f);
        assert_vec_approx(&combined, &out_ref, 1e-3, "col_then_row_roundtrip");
    }
}