scirs2-core 0.1.0-alpha.3

//! # Memory Management
//!
//! This module provides efficient memory management utilities for scientific computing.
//!
//! ## Features
//!
//! * Chunk processing for large datasets
//! * Buffer pool for memory reuse
//! * Zero-copy transformations
//! * Memory usage tracking and metrics
//! * Detailed memory allocation analysis
//!
//! ## Usage
//!
//! ```rust,no_run
//! use scirs2_core::memory::{ChunkProcessor2D, BufferPool};
//! use ndarray::Array2;
//!
//! // Process a large array in chunks
//! let large_array = Array2::<f64>::zeros((10000, 10000));
//! let mut processor = ChunkProcessor2D::new(&large_array, (1000, 1000));
//!
//! processor.process_chunks(|chunk, coords| {
//!     // Process each chunk (e.g., compute statistics, apply transformations)
//!     println!("Processing chunk at position {:?}", coords);
//! });
//!
//! // Use a buffer pool to reuse memory
//! let mut pool = BufferPool::<f64>::new();
//!
//! // Acquire a buffer from the pool
//! let mut buffer = pool.acquire_vec(1000);
//!
//! // Use the buffer for some computation
//! for i in 0..buffer.len() {
//!     buffer[i] = i as f64;
//! }
//!
//! // Release the buffer back to the pool when done
//! pool.release_vec(buffer);
//!
//! // Track memory usage with the metrics system
//! use scirs2_core::memory::metrics::{track_allocation, track_deallocation, format_memory_report};
//!
//! // Record an allocation
//! track_allocation("MyComponent", 1024, 0x1000);
//!
//! // Record a deallocation
//! track_deallocation("MyComponent", 1024, 0x1000);
//!
//! // Print a memory usage report
//! println!("{}", format_memory_report());
//! ```

use ndarray::{ArrayBase, Data, Dimension, Ix2, ViewRepr};
use std::any::TypeId;
use std::collections::HashMap;
use std::marker::PhantomData;
use std::sync::{Arc, Mutex};

/// A processor for working with large arrays in manageable chunks
pub struct ChunkProcessor<'a, A, S, D>
where
    S: Data<Elem = A>,
    D: Dimension,
{
    array: &'a ArrayBase<S, D>,
    // Chunk shape is necessary for the full implementation of chunked processing
    #[allow(dead_code)]
    chunk_shape: D,
    // Current position needed for iterative processing
    #[allow(dead_code)]
    position: D,
}

impl<'a, A, S, D> ChunkProcessor<'a, A, S, D>
where
    S: Data<Elem = A>,
    D: Dimension,
{
    /// Create a new chunk processor for the given array and chunk shape
    pub fn new(array: &'a ArrayBase<S, D>, chunk_shape: D) -> Self {
        let position = D::zeros(array.ndim());
        Self {
            array,
            chunk_shape,
            position,
        }
    }

    /// Process the array in chunks, calling the provided function for each chunk
    pub fn process_chunks<F>(&mut self, mut f: F)
    where
        F: FnMut(&ArrayBase<ViewRepr<&A>, D>, D),
    {
        // This is a simplified implementation that just iterates through chunks
        // A real implementation would be more sophisticated and handle edge cases
        let position = D::zeros(self.array.ndim());

        // TODO: Implement actual chunk iteration logic
        // For now, just call the function on the entire array
        f(&self.array.view(), position);
    }

    /// Get the total number of chunks
    pub fn num_chunks(&self) -> usize {
        // This is a simplified placeholder
        // A real implementation would calculate the actual number of chunks
        1
    }
}

/// A specialized chunk processor for 2D arrays
pub struct ChunkProcessor2D<'a, A, S>
where
    S: Data<Elem = A>,
{
    array: &'a ArrayBase<S, Ix2>,
    chunk_shape: (usize, usize),
    // Current position tracking for iterator implementation
    #[allow(dead_code)]
    current_row: usize,
    #[allow(dead_code)]
    current_col: usize,
}

impl<'a, A, S> ChunkProcessor2D<'a, A, S>
where
    S: Data<Elem = A>,
{
    /// Create a new 2D chunk processor
    pub fn new(array: &'a ArrayBase<S, Ix2>, chunk_shape: (usize, usize)) -> Self {
        Self {
            array,
            chunk_shape,
            current_row: 0,
            current_col: 0,
        }
    }

    /// Process the 2D array in chunks
    pub fn process_chunks<F>(&mut self, mut f: F)
    where
        F: FnMut(&ArrayBase<ViewRepr<&A>, Ix2>, (usize, usize)),
    {
        let (rows, cols) = self.array.dim();
        let (chunk_rows, chunk_cols) = self.chunk_shape;

        for row_start in (0..rows).step_by(chunk_rows) {
            for col_start in (0..cols).step_by(chunk_cols) {
                let row_end = (row_start + chunk_rows).min(rows);
                let col_end = (col_start + chunk_cols).min(cols);

                // Get a view of the current chunk
                let chunk = self
                    .array
                    .slice(ndarray::s![row_start..row_end, col_start..col_end]);

                // Call the processing function
                f(&chunk, (row_start, col_start));
            }
        }
    }
}

/// Memory buffer pool for reusing allocated memory
pub struct BufferPool<T: Clone + Default> {
    // In a real implementation, this would be thread-safe with proper locking
    // and would handle different buffer sizes more efficiently
    vectors: Vec<Vec<T>>,
    arrays: Vec<ndarray::Array1<T>>,
}

impl<T: Clone + Default> BufferPool<T> {
    /// Create a new buffer pool
    pub fn new() -> Self {
        Self {
            vectors: Vec::new(),
            arrays: Vec::new(),
        }
    }

    /// Acquire a vector from the pool, or create a new one if none are available
    pub fn acquire_vec(&mut self, capacity: usize) -> Vec<T> {
        // Find a suitable vector in the pool
        for i in 0..self.vectors.len() {
            if self.vectors[i].capacity() >= capacity {
                // Found a suitable vector, remove it from the pool and return it
                let mut vec = self.vectors.swap_remove(i);
                vec.clear();
                vec.resize(capacity, T::default());
                return vec;
            }
        }

        // No suitable vector found, create a new one
        vec![T::default(); capacity]
    }

    /// Release a vector back to the pool
    pub fn release_vec(&mut self, vec: Vec<T>) {
        // Add the vector back to the pool for reuse
        self.vectors.push(vec);
    }

    /// Acquire an ndarray::Array1 from the pool, or create a new one if none are available
    pub fn acquire_array(&mut self, size: usize) -> ndarray::Array1<T> {
        // Find a suitable array in the pool
        for i in 0..self.arrays.len() {
            if self.arrays[i].len() >= size {
                // Found a suitable array, remove it from the pool and return it
                let mut array = self.arrays.swap_remove(i);
                // Resize the array (this will truncate or extend)
                if array.len() != size {
                    array = ndarray::Array1::from_elem(size, T::default());
                }
                return array;
            }
        }

        // No suitable array found, create a new one
        ndarray::Array1::from_elem(size, T::default())
    }

    /// Release an ndarray::Array1 back to the pool
    pub fn release_array(&mut self, array: ndarray::Array1<T>) {
        // Add the array back to the pool for reuse
        self.arrays.push(array);
    }

    /// Clear the pool, releasing all memory
    pub fn clear(&mut self) {
        self.vectors.clear();
        self.arrays.clear();
    }
}

impl<T: Clone + Default> Default for BufferPool<T> {
    fn default() -> Self {
        Self::new()
    }
}

/// Shared global buffer pool for different types
pub struct GlobalBufferPool {
    // Use TypeId to store pools for different types
    pools: Mutex<HashMap<TypeId, Box<dyn std::any::Any + Send + Sync>>>,
}

impl GlobalBufferPool {
    /// Create a new global buffer pool
    pub fn new() -> Self {
        Self {
            pools: Mutex::new(HashMap::new()),
        }
    }

    /// Get or create a buffer pool for a specific type
    pub fn get_pool<T: Clone + Default + 'static + Send + Sync>(
        &self,
    ) -> Arc<Mutex<BufferPool<T>>> {
        let type_id = TypeId::of::<T>();
        let mut pools = self.pools.lock().unwrap();

        use std::collections::hash_map::Entry;
        match pools.entry(type_id) {
            Entry::Vacant(entry) => {
                // Create a new pool for this type
                let pool = Arc::new(Mutex::new(BufferPool::<T>::new()));
                entry.insert(Box::new(pool.clone()));
                pool
            }
            Entry::Occupied(entry) => {
                // Return the existing pool
                match entry.get().downcast_ref::<Arc<Mutex<BufferPool<T>>>>() {
                    Some(pool) => pool.clone(),
                    None => panic!("Type mismatch in global buffer pool"),
                }
            }
        }
    }

    /// Clear all pools, releasing all memory
    pub fn clear_all(&self) {
        let mut pools = self.pools.lock().unwrap();
        pools.clear();
    }
}

/// Implementation of Default for GlobalBufferPool
impl Default for GlobalBufferPool {
    fn default() -> Self {
        Self::new()
    }
}

/// Static global buffer pool instance
pub fn global_buffer_pool() -> &'static GlobalBufferPool {
    use once_cell::sync::Lazy;
    static GLOBAL_POOL: Lazy<GlobalBufferPool> = Lazy::new(GlobalBufferPool::new);
    &GLOBAL_POOL
}

/// Zero-copy array view for efficient data transformations
pub struct ZeroCopyView<'a, T, D>
where
    D: Dimension,
{
    _phantom: PhantomData<T>,
    inner: ndarray::ArrayView<'a, T, D>,
}

impl<'a, T, D> ZeroCopyView<'a, T, D>
where
    D: Dimension,
{
    /// Create a new zero-copy view from an array
    pub fn new(array: &'a ndarray::Array<T, D>) -> Self {
        Self {
            _phantom: PhantomData,
            inner: array.view(),
        }
    }

    /// Get the underlying array view
    pub fn view(&self) -> ndarray::ArrayView<'a, T, D> {
        self.inner.clone()
    }

    /// Transform the view using a mapping function
    pub fn transform<F, U>(&self, f: F) -> ndarray::Array<U, D>
    where
        F: Fn(&T) -> U,
        U: Clone,
    {
        self.inner.map(f)
    }
}

/// Memory usage tracker for monitoring memory consumption
pub struct MemoryTracker {
    allocations: Mutex<HashMap<String, usize>>,
}

impl MemoryTracker {
    /// Create a new memory tracker
    pub fn new() -> Self {
        Self {
            allocations: Mutex::new(HashMap::new()),
        }
    }
}

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

impl MemoryTracker {
    /// Track a memory allocation
    pub fn track_allocation(&self, name: &str, size: usize) {
        let mut allocations = self.allocations.lock().unwrap();
        *allocations.entry(name.to_string()).or_insert(0) += size;
    }

    /// Track a memory deallocation
    pub fn track_deallocation(&self, name: &str, size: usize) {
        let mut allocations = self.allocations.lock().unwrap();
        if let Some(current) = allocations.get_mut(name) {
            *current = current.saturating_sub(size);
        }
    }

    /// Get the current memory usage for a specific allocation
    pub fn get_usage(&self, name: &str) -> usize {
        let allocations = self.allocations.lock().unwrap();
        allocations.get(name).copied().unwrap_or_default()
    }

    /// Get the total memory usage across all tracked allocations
    pub fn get_total_usage(&self) -> usize {
        let allocations = self.allocations.lock().unwrap();
        allocations.values().sum()
    }

    /// Reset all tracking data
    pub fn reset(&self) {
        let mut allocations = self.allocations.lock().unwrap();
        allocations.clear();
    }
}

/// Static global memory tracker instance
pub fn global_memory_tracker() -> &'static MemoryTracker {
    use once_cell::sync::Lazy;
    static GLOBAL_TRACKER: Lazy<MemoryTracker> = Lazy::new(MemoryTracker::new);
    &GLOBAL_TRACKER
}

/// Advanced memory metrics system
pub mod metrics;

// Re-export key metric functions for convenient usage
pub use metrics::{
    format_memory_report, generate_memory_report, track_allocation, track_deallocation,
    track_resize,
};