"""
Memory Optimization Configuration
Configuration for the unified memory management system.
"""
import torch
from dataclasses import dataclass
from typing import Optional

@dataclass
class MemoryOptimizationConfig:
    """Configuration for memory optimization system."""
    # Device settings
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Memory thresholds
    memory_threshold: float = 0.85  # 85% memory usage threshold
    cleanup_threshold: float = 0.75  # 75% threshold for proactive cleanup
    
    # Tensor pool settings
    max_pool_size: int = 50  # Reduced for RTX 3060
    max_tensor_size: int = 1000000  # Don't pool very large tensors
    
    # Cleanup settings
    cleanup_frequency: int = 25  # Cleanup every 25 operations
    
    # Memory optimization flags
    use_4bit_quantization: bool = True
    use_gradient_checkpointing: bool = True
    use_mixed_precision: bool = True
    
    # Shared model integration
    use_shared_model: bool = True  # Use shared Qwen model
    shared_model_name: str = "Qwen/Qwen3-0.6B"  # Default shared model
    
    # Lazy loading
    enable_lazy_loading: bool = True