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

pub mod auto_batch_size;
pub mod lazy_loader;
pub mod oom_detection;
pub mod precision;
pub mod qat;
pub mod quantization;

pub use auto_batch_size::{
    detect_gpu_memory, AutoBatchSizer, BatchSizeConfig, GpuMemoryInfo, ModelPrecision,
    OptimizerType,
};
pub use lazy_loader::{LazyCheckpointLoader, LoadStrategy};
pub use oom_detection::{extract_oom_size, is_oom_error};
pub use precision::{PrecisionConverter, PrecisionType};
pub use qat::{
    compare_qat_vs_ptq_accuracy, estimate_qparams_from_tensor, fake_quantize_per_channel,
    fake_quantize_tensor, load_observer_state, save_observer_state, FakeQuantize, ObserverState,
    QATConfig, QuantizationObserver,
};
pub use quantization::{
    extract_weights_from_varmap, QuantizationConfig, QuantizationType, Quantizer,
};

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,

    /// Quantization type (none, int8, int4)
    pub quantization: QuantizationType,

    /// 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,
            quantization: QuantizationType::None,
            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()
    }
}