//! Memory optimization utilities for production ML models
//!
//! Provides lazy loading and precision management for memory-constrained deployments.

pub mod auto_batch_size;
pub mod lazy_loader;
pub mod oom_detection;
pub mod precision;

pub use auto_batch_size::{
    detect_gpu_hardware, detect_gpu_memory, sm_count_from_device_name, AutoBatchSizer,
    BatchSizeConfig, GpuHardwareInfo, GpuMemoryInfo, ModelPrecision, OptimalReplayConfig,
    OptimizerType,
};
pub use lazy_loader::{LazyCheckpointLoader, LoadStrategy};
pub use oom_detection::{extract_oom_size, is_oom_error};
pub use precision::{PrecisionConverter, PrecisionType};

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Memory optimization configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryOptimizationConfig {
    /// Enable lazy checkpoint loading
    pub lazy_loading: bool,

    /// Precision type for inference (float32, float16, bfloat16)
    pub precision: PrecisionType,

    /// Maximum memory budget per model (MB)
    pub max_memory_mb: Option<f64>,

    /// Enable gradient checkpointing during training
    pub gradient_checkpointing: bool,

    /// Cache frequently used tensors
    pub tensor_caching: bool,
}

impl Default for MemoryOptimizationConfig {
    fn default() -> Self {
        Self {
            lazy_loading: true,
            precision: PrecisionType::Float32,
            max_memory_mb: None,
            gradient_checkpointing: false,
            tensor_caching: true,
        }
    }
}

/// Memory usage statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryStats {
    /// Current memory usage (MB)
    pub current_mb: f64,

    /// Peak memory usage (MB)
    pub peak_mb: f64,

    /// Memory saved by optimizations (MB)
    pub savings_mb: f64,

    /// Breakdown by component
    pub breakdown: HashMap<String, f64>,
}

impl MemoryStats {
    pub fn new() -> Self {
        Self {
            current_mb: 0.0,
            peak_mb: 0.0,
            savings_mb: 0.0,
            breakdown: HashMap::new(),
        }
    }

    pub fn update_peak(&mut self, current: f64) {
        self.current_mb = current;
        if current > self.peak_mb {
            self.peak_mb = current;
        }
    }

    pub fn add_component(&mut self, name: &str, memory_mb: f64) {
        self.breakdown.insert(name.to_string(), memory_mb);
    }
}

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