numrs2 0.3.3

//! Network Optimization for Distributed Computing
//!
//! This module provides network-aware optimizations to improve communication
//! efficiency in distributed computations.
//!
//! # Features
//!
//! - Network topology detection and modeling
//! - Bandwidth and latency measurement
//! - Communication pattern optimization
//! - Computation-communication overlap
//! - Data compression for network transfer
//!
//! # Example
//!
//! ```rust,no_run
//! use numrs2::distributed::optimization::*;
//! use numrs2::distributed::process::*;
//!
//! # async fn example() -> Result<(), OptimizationError> {
//! let world = init().await?;
//!
//! // Detect network topology
//! let topology = detect_topology(&world).await?;
//! println!("Network topology: {:?}", topology);
//!
//! // Measure network characteristics
//! if world.rank() == 0 && world.size() > 1 {
//!     let bandwidth = measure_bandwidth(0, 1, &world).await?;
//!     let latency = measure_latency(0, 1, &world).await?;
//!     println!("Bandwidth: {} MB/s, Latency: {} μs", bandwidth, latency);
//! }
//!
//! finalize(world).await?;
//! # Ok(())
//! # }
//! ```

use super::collective::CollectiveError;
use super::process::{Communicator, ProcessError};
use oxiarc_lz4::{compress as lz4_compress, decompress as lz4_decompress};
use serde::{Deserialize, Serialize};
use thiserror::Error;

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

    #[error("Collective operation error: {0}")]
    Collective(#[from] CollectiveError),

    #[error("Topology detection failed: {0}")]
    TopologyError(String),

    #[error("Measurement failed: {0}")]
    MeasurementError(String),

    #[error("Optimization failed: {0}")]
    OptimizationFailed(String),

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

/// Network topology types
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum NetworkTopology {
    /// Fully connected network (all-to-all)
    FullyConnected,

    /// Tree topology
    Tree { arity: usize },

    /// Ring topology
    Ring,

    /// 2D/3D mesh topology
    Mesh { dims: [usize; 3] },

    /// Hypercube topology
    Hypercube { dimension: usize },

    /// Fat-tree topology (common in data centers)
    FatTree { levels: usize },

    /// Custom topology
    Custom,
}

impl NetworkTopology {
    /// Get optimal algorithm for collective operations on this topology
    pub fn optimal_algorithm(&self, op: &str) -> Algorithm {
        match (self, op) {
            (NetworkTopology::Tree { .. }, "broadcast") => Algorithm::TreeBroadcast,
            (NetworkTopology::Ring, "reduce") => Algorithm::RingReduce,
            (NetworkTopology::Hypercube { .. }, "allreduce") => Algorithm::HypercubeAllReduce,
            _ => Algorithm::Default,
        }
    }

    /// Check if topology supports efficient point-to-point for given rank pair
    pub fn has_direct_connection(&self, src: usize, dst: usize, _size: usize) -> bool {
        match self {
            NetworkTopology::FullyConnected => true,
            NetworkTopology::Ring => (src as i64 - dst as i64).abs() == 1,
            _ => false, // Conservative for other topologies
        }
    }
}

/// Communication algorithm variants
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Algorithm {
    /// Default algorithm
    Default,

    /// Tree-based broadcast
    TreeBroadcast,

    /// Ring-based reduction
    RingReduce,

    /// Hypercube all-reduce
    HypercubeAllReduce,

    /// Recursive doubling
    RecursiveDoubling,

    /// Pairwise exchange
    PairwiseExchange,
}

/// Network bandwidth model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BandwidthModel {
    /// Measured bandwidths between process pairs (in MB/s)
    measurements: Vec<(usize, usize, f64)>,

    /// Average bandwidth
    average: f64,

    /// Minimum bandwidth (bottleneck)
    min: f64,

    /// Maximum bandwidth
    max: f64,
}

impl BandwidthModel {
    /// Create a new bandwidth model
    pub fn new() -> Self {
        Self {
            measurements: Vec::new(),
            average: 0.0,
            min: 0.0,
            max: 0.0,
        }
    }

    /// Add a measurement
    pub fn add_measurement(&mut self, src: usize, dst: usize, bandwidth: f64) {
        self.measurements.push((src, dst, bandwidth));
        self.update_statistics();
    }

    /// Update statistics after adding measurements
    fn update_statistics(&mut self) {
        if self.measurements.is_empty() {
            return;
        }

        let values: Vec<f64> = self.measurements.iter().map(|(_, _, bw)| *bw).collect();

        self.average = values.iter().sum::<f64>() / values.len() as f64;
        self.min = values.iter().copied().fold(f64::INFINITY, f64::min);
        self.max = values.iter().copied().fold(f64::NEG_INFINITY, f64::max);
    }

    /// Get estimated bandwidth between two processes
    pub fn estimate(&self, src: usize, dst: usize) -> f64 {
        // Look for exact measurement
        for &(s, d, bw) in &self.measurements {
            if (s == src && d == dst) || (s == dst && d == src) {
                return bw;
            }
        }

        // Return average as fallback
        self.average
    }
}

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

/// Network latency model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LatencyModel {
    /// Measured latencies between process pairs (in microseconds)
    measurements: Vec<(usize, usize, f64)>,

    /// Average latency
    average: f64,

    /// Minimum latency
    min: f64,

    /// Maximum latency
    max: f64,
}

impl LatencyModel {
    /// Create a new latency model
    pub fn new() -> Self {
        Self {
            measurements: Vec::new(),
            average: 0.0,
            min: 0.0,
            max: 0.0,
        }
    }

    /// Add a measurement
    pub fn add_measurement(&mut self, src: usize, dst: usize, latency: f64) {
        self.measurements.push((src, dst, latency));
        self.update_statistics();
    }

    /// Update statistics
    fn update_statistics(&mut self) {
        if self.measurements.is_empty() {
            return;
        }

        let values: Vec<f64> = self.measurements.iter().map(|(_, _, lat)| *lat).collect();

        self.average = values.iter().sum::<f64>() / values.len() as f64;
        self.min = values.iter().copied().fold(f64::INFINITY, f64::min);
        self.max = values.iter().copied().fold(f64::NEG_INFINITY, f64::max);
    }

    /// Get estimated latency between two processes
    pub fn estimate(&self, src: usize, dst: usize) -> f64 {
        // Look for exact measurement
        for &(s, d, lat) in &self.measurements {
            if (s == src && d == dst) || (s == dst && d == src) {
                return lat;
            }
        }

        // Return average as fallback
        self.average
    }
}

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

/// Detect network topology
///
/// Attempts to detect the network topology by analyzing process connectivity.
pub async fn detect_topology(comm: &Communicator) -> Result<NetworkTopology, OptimizationError> {
    let size = comm.size();

    // Simple heuristic based on size
    if size.is_power_of_two() && size >= 8 {
        Ok(NetworkTopology::Hypercube {
            dimension: (size as f64).log2() as usize,
        })
    } else {
        // Default to fully connected for small clusters
        Ok(NetworkTopology::FullyConnected)
    }
}

/// Measure bandwidth between two processes
///
/// Sends test messages to measure network bandwidth.
pub async fn measure_bandwidth(
    _src: usize,
    _dst: usize,
    _comm: &Communicator,
) -> Result<f64, OptimizationError> {
    // Placeholder implementation
    // Real implementation would:
    // 1. Send multiple test messages of varying sizes
    // 2. Measure transfer time
    // 3. Calculate bandwidth = size / time

    // Return estimated bandwidth (in MB/s)
    Ok(1000.0) // Placeholder: 1 GB/s
}

/// Measure latency between two processes
///
/// Sends small test messages to measure round-trip latency.
pub async fn measure_latency(
    _src: usize,
    _dst: usize,
    _comm: &Communicator,
) -> Result<f64, OptimizationError> {
    // Placeholder implementation
    // Real implementation would:
    // 1. Send small ping messages
    // 2. Measure round-trip time
    // 3. Calculate one-way latency = RTT / 2

    // Return estimated latency (in microseconds)
    Ok(10.0) // Placeholder: 10 μs
}

/// Optimize collective operation for given topology
///
/// Returns the optimal algorithm for a collective operation on the given topology.
pub fn optimize_collective(
    _op: &str,
    topology: &NetworkTopology,
) -> Result<Algorithm, OptimizationError> {
    Ok(topology.optimal_algorithm(_op))
}

/// Overlap computation and communication using async operations
///
/// This is a placeholder for future async computation-communication overlap implementation.
pub async fn overlap_compute_communicate() -> Result<(), OptimizationError> {
    // Placeholder implementation
    // Real implementation would:
    // 1. Launch communication operations asynchronously
    // 2. Perform computation while communication is in progress
    // 3. Synchronize when both complete

    Ok(())
}

/// Compress data for network transfer
///
/// Serializes the data to JSON and compresses with LZ4 (fast, low-latency).
/// Wire format: [8-byte u64 LE uncompressed size][LZ4 frame compressed bytes]
pub fn compress_data<T: Serialize>(data: &[T]) -> Result<Vec<u8>, OptimizationError> {
    let json_bytes = serde_json::to_vec(data)
        .map_err(|e| OptimizationError::CompressionError(format!("Serialization error: {}", e)))?;

    let uncompressed_size = json_bytes.len() as u64;

    let compressed = lz4_compress(&json_bytes).map_err(|e| {
        OptimizationError::CompressionError(format!("LZ4 compression error: {}", e))
    })?;

    let mut result = Vec::with_capacity(8 + compressed.len());
    result.extend_from_slice(&uncompressed_size.to_le_bytes());
    result.extend_from_slice(&compressed);

    Ok(result)
}

/// Decompress data after network transfer
///
/// Decompresses data that was compressed with `compress_data`.
/// Wire format: [8-byte u64 LE uncompressed size][LZ4 frame compressed bytes]
pub fn decompress_data<T: for<'de> Deserialize<'de>>(
    data: &[u8],
) -> Result<Vec<T>, OptimizationError> {
    if data.len() < 8 {
        return Err(OptimizationError::CompressionError(format!(
            "Data too short: expected at least 8 bytes, got {}",
            data.len()
        )));
    }

    let size_bytes: [u8; 8] = data[..8].try_into().map_err(|_| {
        OptimizationError::CompressionError("Failed to read uncompressed size header".to_string())
    })?;
    let uncompressed_size = u64::from_le_bytes(size_bytes) as usize;

    let json_bytes = lz4_decompress(&data[8..], uncompressed_size).map_err(|e| {
        OptimizationError::CompressionError(format!("LZ4 decompression error: {}", e))
    })?;

    serde_json::from_slice(&json_bytes)
        .map_err(|e| OptimizationError::CompressionError(format!("Deserialization error: {}", e)))
}

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

    #[test]
    fn test_network_topology() {
        let topology = NetworkTopology::Tree { arity: 2 };
        assert_eq!(
            topology.optimal_algorithm("broadcast"),
            Algorithm::TreeBroadcast
        );
    }

    #[test]
    fn test_bandwidth_model() {
        let mut model = BandwidthModel::new();
        model.add_measurement(0, 1, 1000.0);
        model.add_measurement(1, 2, 950.0);
        model.add_measurement(2, 3, 1050.0);

        assert_eq!(model.estimate(0, 1), 1000.0);
        assert!((model.average - 1000.0).abs() < 50.0);
    }

    #[test]
    fn test_latency_model() {
        let mut model = LatencyModel::new();
        model.add_measurement(0, 1, 10.0);
        model.add_measurement(1, 2, 12.0);
        model.add_measurement(2, 3, 11.0);

        assert_eq!(model.estimate(0, 1), 10.0);
        assert!((model.average - 11.0).abs() < 1.0);
    }

    #[test]
    fn test_topology_direct_connection() {
        let topology = NetworkTopology::FullyConnected;
        assert!(topology.has_direct_connection(0, 1, 4));
        assert!(topology.has_direct_connection(0, 3, 4));

        let ring = NetworkTopology::Ring;
        assert!(ring.has_direct_connection(0, 1, 4));
        assert!(!ring.has_direct_connection(0, 2, 4));
    }

    #[test]
    fn test_compress_decompress_roundtrip_floats() {
        let data: Vec<f64> = vec![1.0, 2.5, -3.14, 0.0, f64::MAX];
        let compressed = compress_data(&data).expect("compression should succeed");
        let recovered: Vec<f64> =
            decompress_data(&compressed).expect("decompression should succeed");
        assert_eq!(data.len(), recovered.len());
        for (a, b) in data.iter().zip(recovered.iter()) {
            assert!(
                (a - b).abs() < f64::EPSILON * 100.0,
                "mismatch: {} vs {}",
                a,
                b
            );
        }
    }

    #[test]
    fn test_compress_decompress_roundtrip_strings() {
        let data: Vec<String> = vec![
            "hello".to_string(),
            "world".to_string(),
            "oxiarc".to_string(),
        ];
        let compressed = compress_data(&data).expect("compression should succeed");
        let recovered: Vec<String> =
            decompress_data(&compressed).expect("decompression should succeed");
        assert_eq!(data, recovered);
    }

    #[test]
    fn test_compress_empty_slice() {
        let data: Vec<u32> = vec![];
        let compressed = compress_data(&data).expect("compression of empty slice should succeed");
        let recovered: Vec<u32> =
            decompress_data(&compressed).expect("decompression should succeed");
        assert_eq!(recovered, data);
    }

    #[test]
    fn test_compress_highly_compressible() {
        let data: Vec<u32> = vec![42u32; 10_000];
        let compressed = compress_data(&data).expect("compression should succeed");
        // LZ4 should compress highly repetitive data significantly
        assert!(
            compressed.len() < data.len() * 4,
            "expected compression, got {} bytes for {} elements",
            compressed.len(),
            data.len()
        );
        let recovered: Vec<u32> =
            decompress_data(&compressed).expect("decompression should succeed");
        assert_eq!(data, recovered);
    }

    #[test]
    fn test_decompress_invalid_data() {
        let bad_data = b"too short";
        let result: Result<Vec<u32>, _> = decompress_data(bad_data);
        assert!(result.is_err(), "should fail on short data");
    }
}