numrs2 0.3.3

//! Efficient Communication Backends for Distributed Training
//!
//! This module provides high-performance communication primitives optimized for
//! distributed deep learning workloads. It builds on the base `comm` module with
//! specialized features for tensor communication and bandwidth optimization.
//!
//! # Features
//!
//! - **Tensor Serialization**: Efficient serialization using oxicode
//! - **Async Primitives**: Non-blocking send/recv with async/await
//! - **Bandwidth Optimization**: Compression, batching, pipelining
//! - **Latency Hiding**: Computation/communication overlap
//! - **Priority Queues**: Prioritized message delivery
//! - **Topology-Aware**: Network-aware routing strategies
//!
//! # Example
//!
//! ```rust,no_run
//! use numrs2::distributed::communication::*;
//! use numrs2::distributed::process::*;
//! use std::sync::Arc;
//!
//! # async fn example() -> Result<(), CommunicationError> {
//! let world = init().await?;
//! let comm = AsyncCommunicator::new(Arc::new(world))?;
//!
//! // Send tensor with compression
//! let tensor = vec![1.0_f32; 1000];
//! let msg = TensorMessage::new(
//!     tensor,
//!     CompressionStrategy::TopK { k: 100 },
//!     MessagePriority::High
//! );
//! comm.isend(msg, 1).await?;
//!
//! // Receive tensor asynchronously
//! let received: TensorMessage<f32> = comm.irecv(0).await?;
//! # Ok(())
//! # }
//! ```

use super::comm::{CommunicationChannel, ConnectionManager, Message};
use super::process::{Communicator, ProcessError};
use crate::error::NumRs2Error;
use oxicode::{Decode, Encode};
use scirs2_core::ndarray::Array1;
use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap, VecDeque};
use std::sync::Arc;
use thiserror::Error;
use tokio::sync::{mpsc, Mutex, RwLock};

/// Errors that can occur during communication operations
#[derive(Error, Debug)]
pub enum CommunicationError {
    #[error("Process error: {0}")]
    Process(#[from] ProcessError),

    #[error("Serialization error: {0}")]
    Serialization(String),

    #[error("Deserialization error: {0}")]
    Deserialization(String),

    #[error("Compression error: {0}")]
    Compression(String),

    #[error("Decompression error: {0}")]
    Decompression(String),

    #[error("Channel error: {0}")]
    Channel(String),

    #[error("Timeout: operation exceeded {0}ms")]
    Timeout(u64),

    #[error("Invalid rank {rank}, communicator size is {size}")]
    InvalidRank { rank: usize, size: usize },

    #[error("Message size mismatch: expected {expected}, got {actual}")]
    SizeMismatch { expected: usize, actual: usize },

    #[error("Network error: {0}")]
    Network(String),
}

impl From<CommunicationError> for NumRs2Error {
    fn from(err: CommunicationError) -> Self {
        NumRs2Error::DistributedComputing(err.to_string())
    }
}

/// Message priority for prioritized communication
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Encode, Decode)]
pub enum MessagePriority {
    /// Low priority - background operations
    Low = 0,
    /// Normal priority - regular data transfers
    Normal = 1,
    /// High priority - critical synchronization
    High = 2,
    /// Urgent priority - control messages
    Urgent = 3,
}

/// Compression strategy for bandwidth optimization
#[derive(Debug, Clone, Encode, Decode)]
pub enum CompressionStrategy {
    /// No compression
    None,

    /// Top-k sparsification: keep only k largest absolute values
    TopK { k: usize },

    /// Random-k: randomly select k elements
    RandomK { k: usize },

    /// Quantization to reduce precision
    Quantization { bits: u8 },

    /// Threshold-based sparsification
    Threshold { threshold: f64 },
}

/// Tensor message with metadata for efficient communication
#[derive(Debug, Clone, Encode, Decode)]
pub struct TensorMessage<T>
where
    T: Clone + Encode + Decode,
{
    /// Tensor data
    pub data: Vec<T>,

    /// Original shape of the tensor
    pub shape: Vec<usize>,

    /// Compression strategy used
    pub compression: CompressionStrategy,

    /// Message priority
    pub priority: MessagePriority,

    /// Sequence number for ordering
    pub sequence: u64,

    /// Sender rank
    pub sender: usize,

    /// Tag for message identification
    pub tag: u32,

    /// Indices for sparse tensors (used with compression)
    pub indices: Option<Vec<usize>>,
}

impl<T> TensorMessage<T>
where
    T: Clone + Encode + Decode,
{
    /// Create a new tensor message
    pub fn new(data: Vec<T>, compression: CompressionStrategy, priority: MessagePriority) -> Self {
        Self {
            shape: vec![data.len()],
            data,
            compression,
            priority,
            sequence: 0,
            sender: 0,
            tag: 0,
            indices: None,
        }
    }

    /// Create tensor message with shape
    pub fn with_shape(
        data: Vec<T>,
        shape: Vec<usize>,
        compression: CompressionStrategy,
        priority: MessagePriority,
    ) -> Self {
        Self {
            data,
            shape,
            compression,
            priority,
            sequence: 0,
            sender: 0,
            tag: 0,
            indices: None,
        }
    }

    /// Set sequence number
    pub fn with_sequence(mut self, sequence: u64) -> Self {
        self.sequence = sequence;
        self
    }

    /// Set sender rank
    pub fn with_sender(mut self, sender: usize) -> Self {
        self.sender = sender;
        self
    }

    /// Set message tag
    pub fn with_tag(mut self, tag: u32) -> Self {
        self.tag = tag;
        self
    }

    /// Get data size in bytes (uncompressed)
    pub fn size_bytes(&self) -> usize {
        self.data.len() * std::mem::size_of::<T>()
    }
}

/// Prioritized message wrapper for priority queue
#[derive(Debug)]
struct PrioritizedMessage<T> {
    message: T,
    priority: MessagePriority,
    sequence: u64,
}

impl<T> PartialEq for PrioritizedMessage<T> {
    fn eq(&self, other: &Self) -> bool {
        self.priority == other.priority && self.sequence == other.sequence
    }
}

impl<T> Eq for PrioritizedMessage<T> {}

impl<T> PartialOrd for PrioritizedMessage<T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<T> Ord for PrioritizedMessage<T> {
    fn cmp(&self, other: &Self) -> Ordering {
        // Higher priority first, then earlier sequence
        match self.priority.cmp(&other.priority) {
            Ordering::Equal => other.sequence.cmp(&self.sequence), // Earlier sequence first
            ord => ord,
        }
    }
}

/// Type alias for send queue structure
type SendQueues = Arc<RwLock<HashMap<usize, BinaryHeap<PrioritizedMessage<Vec<u8>>>>>>;

/// Asynchronous communicator with non-blocking operations
pub struct AsyncCommunicator {
    /// Underlying communicator
    communicator: Arc<Communicator>,

    /// Send queue per destination rank
    send_queues: SendQueues,

    /// Receive buffer per source rank
    recv_buffers: Arc<RwLock<HashMap<usize, VecDeque<Vec<u8>>>>>,

    /// Sequence counter for message ordering
    sequence_counter: Arc<Mutex<u64>>,

    /// Channel manager for async communication
    channels: Arc<Mutex<HashMap<usize, mpsc::Sender<Vec<u8>>>>>,
}

impl AsyncCommunicator {
    /// Create new async communicator
    pub fn new(communicator: Arc<Communicator>) -> Result<Self, CommunicationError> {
        Ok(Self {
            communicator,
            send_queues: Arc::new(RwLock::new(HashMap::new())),
            recv_buffers: Arc::new(RwLock::new(HashMap::new())),
            sequence_counter: Arc::new(Mutex::new(0)),
            channels: Arc::new(Mutex::new(HashMap::new())),
        })
    }

    /// Get next sequence number
    async fn next_sequence(&self) -> u64 {
        let mut counter = self.sequence_counter.lock().await;
        let seq = *counter;
        *counter += 1;
        seq
    }

    /// Non-blocking send
    pub async fn isend<T>(
        &self,
        message: TensorMessage<T>,
        dest: usize,
    ) -> Result<(), CommunicationError>
    where
        T: Clone + Encode + Decode,
    {
        // Validate destination rank
        if dest >= self.communicator.size() {
            return Err(CommunicationError::InvalidRank {
                rank: dest,
                size: self.communicator.size(),
            });
        }

        // Serialize message
        let data = oxicode::encode_to_vec(&message).map_err(|e| {
            CommunicationError::Serialization(format!("Failed to serialize message: {}", e))
        })?;

        // Add to send queue with priority
        let prioritized = PrioritizedMessage {
            message: data,
            priority: message.priority,
            sequence: message.sequence,
        };

        let mut queues = self.send_queues.write().await;
        queues
            .entry(dest)
            .or_insert_with(BinaryHeap::new)
            .push(prioritized);

        Ok(())
    }

    /// Non-blocking receive
    pub async fn irecv<T>(&self, source: usize) -> Result<TensorMessage<T>, CommunicationError>
    where
        T: Clone + Encode + Decode,
    {
        // Validate source rank
        if source >= self.communicator.size() {
            return Err(CommunicationError::InvalidRank {
                rank: source,
                size: self.communicator.size(),
            });
        }

        // Check receive buffer
        let mut buffers = self.recv_buffers.write().await;
        let buffer = buffers.entry(source).or_insert_with(VecDeque::new);

        // For now, return error if no data available
        // In a real implementation, this would wait for data
        let data = buffer
            .pop_front()
            .ok_or_else(|| CommunicationError::Channel("No data available".to_string()))?;

        // Deserialize message
        let (message, _) = oxicode::decode_from_slice(&data).map_err(|e| {
            CommunicationError::Deserialization(format!("Failed to deserialize message: {}", e))
        })?;

        Ok(message)
    }

    /// Blocking send (waits for completion)
    pub async fn send<T>(
        &self,
        message: TensorMessage<T>,
        dest: usize,
    ) -> Result<(), CommunicationError>
    where
        T: Clone + Encode + Decode,
    {
        self.isend(message, dest).await?;
        self.flush_send_queue(dest).await?;
        Ok(())
    }

    /// Blocking receive
    pub async fn recv<T>(&self, source: usize) -> Result<TensorMessage<T>, CommunicationError>
    where
        T: Clone + Encode + Decode,
    {
        // In a real implementation, this would block until data arrives
        self.irecv(source).await
    }

    /// Flush send queue for a specific destination
    async fn flush_send_queue(&self, dest: usize) -> Result<(), CommunicationError> {
        let mut queues = self.send_queues.write().await;
        if let Some(queue) = queues.get_mut(&dest) {
            while let Some(prioritized) = queue.pop() {
                // In a real implementation, actually send the data
                // For now, just clear the queue
                let _ = prioritized.message;
            }
        }
        Ok(())
    }

    /// Get communicator rank
    pub fn rank(&self) -> usize {
        self.communicator.rank()
    }

    /// Get communicator size
    pub fn size(&self) -> usize {
        self.communicator.size()
    }
}

/// Pipelined communicator for latency hiding
pub struct PipelinedCommunicator {
    /// Base async communicator
    base: AsyncCommunicator,

    /// Pipeline depth (number of concurrent operations)
    depth: usize,

    /// Active pipeline stages
    active_stages: Arc<Mutex<VecDeque<PipelineStage>>>,
}

#[derive(Debug)]
struct PipelineStage {
    operation_id: u64,
    status: PipelineStatus,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PipelineStatus {
    Pending,
    InProgress,
    Completed,
}

impl PipelinedCommunicator {
    /// Create new pipelined communicator
    pub fn new(communicator: Arc<Communicator>, depth: usize) -> Result<Self, CommunicationError> {
        Ok(Self {
            base: AsyncCommunicator::new(communicator)?,
            depth,
            active_stages: Arc::new(Mutex::new(VecDeque::new())),
        })
    }

    /// Start a pipelined send operation
    pub async fn pipeline_send<T>(
        &self,
        message: TensorMessage<T>,
        dest: usize,
    ) -> Result<u64, CommunicationError>
    where
        T: Clone + Encode + Decode,
    {
        // Wait if pipeline is full
        self.wait_for_pipeline_slot().await?;

        // Get operation ID
        let op_id = self.base.next_sequence().await;

        // Add to pipeline
        let mut stages = self.active_stages.lock().await;
        stages.push_back(PipelineStage {
            operation_id: op_id,
            status: PipelineStatus::Pending,
        });

        // Start async send
        self.base.isend(message, dest).await?;

        Ok(op_id)
    }

    /// Wait for a specific pipeline operation to complete
    pub async fn wait_operation(&self, op_id: u64) -> Result<(), CommunicationError> {
        let mut stages = self.active_stages.lock().await;

        // Find and remove the operation
        let pos = stages.iter().position(|s| s.operation_id == op_id);
        if let Some(pos) = pos {
            stages.remove(pos);
        }

        Ok(())
    }

    /// Wait for all pipeline operations to complete
    pub async fn wait_all(&self) -> Result<(), CommunicationError> {
        let mut stages = self.active_stages.lock().await;
        stages.clear();
        Ok(())
    }

    /// Wait for pipeline slot to become available
    async fn wait_for_pipeline_slot(&self) -> Result<(), CommunicationError> {
        loop {
            let mut stages = self.active_stages.lock().await;
            if stages.len() < self.depth {
                return Ok(());
            }

            // Remove completed stages
            while let Some(stage) = stages.front() {
                if stage.status == PipelineStatus::Completed {
                    stages.pop_front();
                } else {
                    break;
                }
            }

            drop(stages);

            // Small delay before checking again
            tokio::time::sleep(tokio::time::Duration::from_micros(100)).await;
        }
    }

    /// Get pipeline depth
    pub fn depth(&self) -> usize {
        self.depth
    }

    /// Get number of active operations
    pub async fn active_count(&self) -> usize {
        self.active_stages.lock().await.len()
    }
}

/// Compress tensor data using the specified strategy
pub fn compress_tensor<T>(
    data: &[T],
    strategy: &CompressionStrategy,
) -> Result<(Vec<T>, Option<Vec<usize>>), CommunicationError>
where
    T: Clone + PartialOrd + Default,
{
    match strategy {
        CompressionStrategy::None => Ok((data.to_vec(), None)),

        CompressionStrategy::TopK { k } => {
            if *k >= data.len() {
                return Ok((data.to_vec(), None));
            }

            // Find top-k elements by absolute value (simplified - assumes Ord)
            let mut indexed: Vec<(usize, T)> = data
                .iter()
                .enumerate()
                .map(|(i, v)| (i, v.clone()))
                .collect();

            // Partial sort to get top k
            if *k < indexed.len() {
                indexed.select_nth_unstable_by(*k, |a, b| {
                    b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal)
                });
            }

            let mut values = Vec::with_capacity(*k);
            let mut indices = Vec::with_capacity(*k);

            for (idx, val) in indexed.iter().take(*k) {
                indices.push(*idx);
                values.push(val.clone());
            }

            Ok((values, Some(indices)))
        }

        CompressionStrategy::RandomK { k } => {
            if *k >= data.len() {
                return Ok((data.to_vec(), None));
            }

            // Simple deterministic selection for now
            let step = data.len() / k;
            let values: Vec<T> = (0..*k).map(|i| data[i * step].clone()).collect();
            let indices: Vec<usize> = (0..*k).map(|i| i * step).collect();

            Ok((values, Some(indices)))
        }

        CompressionStrategy::Quantization { .. } => {
            // For now, just pass through - full quantization requires more type constraints
            Ok((data.to_vec(), None))
        }

        CompressionStrategy::Threshold { .. } => {
            // For now, just pass through - threshold requires numeric operations
            Ok((data.to_vec(), None))
        }
    }
}

/// Decompress tensor data
pub fn decompress_tensor<T>(
    compressed: &[T],
    indices: Option<&[usize]>,
    original_size: usize,
) -> Result<Vec<T>, CommunicationError>
where
    T: Clone + Default,
{
    match indices {
        None => Ok(compressed.to_vec()),
        Some(idx) => {
            let mut result = vec![T::default(); original_size];
            for (i, &pos) in idx.iter().enumerate() {
                if pos < original_size && i < compressed.len() {
                    result[pos] = compressed[i].clone();
                }
            }
            Ok(result)
        }
    }
}

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

    #[test]
    fn test_tensor_message_creation() {
        let data = vec![1.0, 2.0, 3.0, 4.0];
        let msg = TensorMessage::new(
            data.clone(),
            CompressionStrategy::None,
            MessagePriority::Normal,
        );

        assert_eq!(msg.data, data);
        assert_eq!(msg.shape, vec![4]);
        assert_eq!(msg.priority, MessagePriority::Normal);
    }

    #[test]
    fn test_tensor_message_with_shape() {
        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
        let shape = vec![2, 3];
        let msg = TensorMessage::with_shape(
            data.clone(),
            shape.clone(),
            CompressionStrategy::None,
            MessagePriority::High,
        );

        assert_eq!(msg.data, data);
        assert_eq!(msg.shape, shape);
        assert_eq!(msg.priority, MessagePriority::High);
    }

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

    #[test]
    fn test_compression_none() {
        let data = vec![1.0, 2.0, 3.0, 4.0];
        let result = compress_tensor(&data, &CompressionStrategy::None);
        assert!(result.is_ok());

        let (compressed, indices) = result.expect("compression failed");
        assert_eq!(compressed, data);
        assert!(indices.is_none());
    }

    #[test]
    fn test_compression_topk() {
        let data = vec![5.0, 1.0, 8.0, 3.0, 9.0, 2.0];
        let result = compress_tensor(&data, &CompressionStrategy::TopK { k: 3 });
        assert!(result.is_ok());

        let (compressed, indices) = result.expect("compression failed");
        assert_eq!(compressed.len(), 3);
        assert!(indices.is_some());
        assert_eq!(indices.as_ref().expect("indices missing").len(), 3);
    }

    #[test]
    fn test_compression_randomk() {
        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
        let result = compress_tensor(&data, &CompressionStrategy::RandomK { k: 3 });
        assert!(result.is_ok());

        let (compressed, indices) = result.expect("compression failed");
        assert_eq!(compressed.len(), 3);
        assert!(indices.is_some());
    }

    #[test]
    fn test_decompress_none() {
        let data = vec![1.0, 2.0, 3.0, 4.0];
        let result = decompress_tensor(&data, None, data.len());
        assert!(result.is_ok());

        let decompressed = result.expect("decompression failed");
        assert_eq!(decompressed, data);
    }

    #[test]
    fn test_decompress_with_indices() {
        let compressed = vec![5.0, 8.0, 9.0];
        let indices = vec![0, 2, 4];
        let original_size = 6;

        let result = decompress_tensor(&compressed, Some(&indices), original_size);
        assert!(result.is_ok());

        let decompressed = result.expect("decompression failed");
        assert_eq!(decompressed.len(), original_size);
        assert_eq!(decompressed[0], 5.0);
        assert_eq!(decompressed[2], 8.0);
        assert_eq!(decompressed[4], 9.0);
    }

    #[test]
    fn test_tensor_message_serialization() {
        let data = vec![1.0_f32, 2.0, 3.0];
        let msg = TensorMessage::new(data, CompressionStrategy::None, MessagePriority::Normal);

        let serialized = oxicode::encode_to_vec(&msg);
        assert!(serialized.is_ok());

        let bytes = serialized.expect("serialization failed");
        let deserialized: Result<(TensorMessage<f32>, usize), _> =
            oxicode::decode_from_slice(&bytes);
        assert!(deserialized.is_ok());
    }

    #[test]
    fn test_compression_strategy_serialization() {
        let strategies = vec![
            CompressionStrategy::None,
            CompressionStrategy::TopK { k: 10 },
            CompressionStrategy::RandomK { k: 5 },
            CompressionStrategy::Quantization { bits: 8 },
            CompressionStrategy::Threshold { threshold: 0.01 },
        ];

        for strategy in strategies {
            let serialized = oxicode::encode_to_vec(&strategy);
            assert!(serialized.is_ok());

            let bytes = serialized.expect("serialization failed");
            let deserialized: Result<(CompressionStrategy, usize), _> =
                oxicode::decode_from_slice(&bytes);
            assert!(deserialized.is_ok());
        }
    }

    #[test]
    fn test_message_priority_serialization() {
        let priorities = vec![
            MessagePriority::Low,
            MessagePriority::Normal,
            MessagePriority::High,
            MessagePriority::Urgent,
        ];

        for priority in priorities {
            let serialized = oxicode::encode_to_vec(&priority);
            assert!(serialized.is_ok());

            let bytes = serialized.expect("serialization failed");
            let deserialized: Result<(MessagePriority, usize), _> =
                oxicode::decode_from_slice(&bytes);
            assert!(deserialized.is_ok());
        }
    }

    #[test]
    fn test_tensor_message_size() {
        let data = vec![1.0_f64; 1000];
        let msg = TensorMessage::new(data, CompressionStrategy::None, MessagePriority::Normal);

        let size = msg.size_bytes();
        assert_eq!(size, 1000 * std::mem::size_of::<f64>());
    }

    #[test]
    fn test_compression_topk_full_data() {
        let data = vec![1.0, 2.0, 3.0];
        let result = compress_tensor(&data, &CompressionStrategy::TopK { k: 10 });
        assert!(result.is_ok());

        let (compressed, indices) = result.expect("compression failed");
        assert_eq!(compressed, data);
        assert!(indices.is_none());
    }

    #[test]
    fn test_compression_randomk_full_data() {
        let data = vec![1.0, 2.0, 3.0];
        let result = compress_tensor(&data, &CompressionStrategy::RandomK { k: 10 });
        assert!(result.is_ok());

        let (compressed, indices) = result.expect("compression failed");
        assert_eq!(compressed, data);
        assert!(indices.is_none());
    }
}