config.py

#!/usr/bin/env python3
"""
Configuration Management System
Handles model configuration, training parameters, and deployment settings
"""

import os
import json
import yaml
from pathlib import Path
from typing import Dict, Any, Optional, Union
from dataclasses import dataclass, asdict, field
import logging

# ============================================================
# Configuration Data Classes
# ============================================================

@dataclass
class ModelConfig:
    """Model architecture configuration"""
    vocab_size: int = 39
    hidden_dim: int = 256
    num_layers: int = 4
    nhead: int = 8
    use_pe: bool = True
    dropout: float = 0.1
    activation: str = "relu"
    
    # CNN backbone settings
    cnn_channels: list = field(default_factory=lambda: [64, 128, 256])
    cnn_strides: list = field(default_factory=lambda: [2, 2, 2])
    use_residual: bool = True
    
    # Transformer settings
    dim_feedforward: int = 1024
    attention_dropout: float = 0.1
    layer_norm_eps: float = 1e-5

@dataclass
class TrainingConfig:
    """Training configuration"""
    # Data settings
    csv_file: str = "clean_plates.csv"
    img_dir: str = "src2/resized_plates"
    val_split: float = 0.1
    batch_size: int = 16
    num_workers: int = 4
    pin_memory: bool = True
    
    # Training parameters
    epochs: int = 20
    learning_rate: float = 1e-4
    weight_decay: float = 1e-5
    label_smoothing: float = 0.1
    gradient_clip: float = 1.0
    accumulation_steps: int = 1
    
    # Optimizer settings
    optimizer: str = "AdamW"
    betas: tuple = (0.9, 0.999)
    eps: float = 1e-8
    
    # Scheduler settings
    scheduler: str = "OneCycleLR"
    scheduler_params: Dict[str, Any] = field(default_factory=lambda: {
        "pct_start": 0.1,
        "anneal_strategy": "cos",
        "div_factor": 25,
        "final_div_factor": 10000
    })
    
    # Early stopping
    early_stopping: bool = True
    patience: int = 10
    min_delta: float = 0.001
    
    # Mixed precision
    use_amp: bool = True
    amp_opt_level: str = "O1"

@dataclass
class DataConfig:
    """Data preprocessing configuration"""
    # Image settings
    target_height: int = 96
    target_width: int = 512
    normalize: bool = True
    mean: list = field(default_factory=lambda: [0.485, 0.456, 0.406])
    std: list = field(default_factory=lambda: [0.229, 0.224, 0.225])
    
    # Augmentation settings
    use_augmentation: bool = True
    augmentation_prob: float = 0.5
    
    # Augmentation parameters
    rotation_range: float = 5.0
    translation_range: float = 0.1
    brightness_range: float = 0.2
    contrast_range: float = 0.2
    saturation_range: float = 0.1
    perspective_prob: float = 0.3
    perspective_distortion: float = 0.1
    noise_prob: float = 0.2
    noise_std: float = 10.0
    blur_prob: float = 0.15
    blur_kernel_size: int = 3

@dataclass
class InferenceConfig:
    """Inference configuration"""
    # Model settings
    model_path: str = "results/best.pth"
    vocab_path: str = "vocab.json"
    device: str = "auto"
    compile_model: bool = True
    
    # Decoding settings
    default_method: str = "beam"
    beam_width: int = 5
    max_length: int = 32
    sos_id: int = 1
    eos_id: int = 2
    
    # Quality settings
    min_confidence: float = 0.3
    quality_check: bool = True
    min_quality_score: float = 0.3
    
    # Performance settings
    batch_size: int = 1
    num_workers: int = 2
    pin_memory: bool = True

@dataclass
class LoggingConfig:
    """Logging configuration"""
    level: str = "INFO"
    format: str = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    file_logging: bool = True
    console_logging: bool = True
    log_dir: str = "logs"
    max_file_size: int = 10 * 1024 * 1024  # 10MB
    backup_count: int = 5

@dataclass
class ExportConfig:
    """Model export configuration"""
    output_dir: str = "exported_model"
    formats: list = field(default_factory=lambda: ["torchscript", "onnx"])
    optimize: bool = True
    quantize: bool = False
    
    # TorchScript settings
    torchscript_optimize: bool = True
    
    # ONNX settings
    onnx_opset_version: int = 11
    onnx_dynamic_axes: bool = True
    
    # Package settings
    create_package: bool = True
    include_examples: bool = True
    include_docs: bool = True

@dataclass
class Config:
    """Main configuration class"""
    model: ModelConfig = field(default_factory=ModelConfig)
    training: TrainingConfig = field(default_factory=TrainingConfig)
    data: DataConfig = field(default_factory=DataConfig)
    inference: InferenceConfig = field(default_factory=InferenceConfig)
    logging: LoggingConfig = field(default_factory=LoggingConfig)
    export: ExportConfig = field(default_factory=ExportConfig)
    
    # Global settings
    seed: int = 42
    deterministic: bool = True
    cudnn_benchmark: bool = True
    cudnn_deterministic: bool = False

# ============================================================
# Configuration Manager
# ============================================================

class ConfigManager:
    """Configuration manager for loading, saving, and validating configs"""
    
    def __init__(self, config_path: Optional[str] = None):
        """
        Initialize configuration manager
        
        Args:
            config_path: Path to configuration file
        """
        self.config_path = config_path
        self.config = Config()
        self.logger = self._setup_logger()
        
        if config_path and Path(config_path).exists():
            self.load_config(config_path)
    
    def _setup_logger(self) -> logging.Logger:
        """Setup logger"""
        logger = logging.getLogger("ConfigManager")
        logger.setLevel(logging.INFO)
        
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)
        
        return logger
    
    def load_config(self, config_path: str) -> 'ConfigManager':
        """
        Load configuration from file
        
        Args:
            config_path: Path to configuration file
            
        Returns:
            Self for method chaining
        """
        config_path = Path(config_path)
        
        if not config_path.exists():
            raise FileNotFoundError(f"Configuration file not found: {config_path}")
        
        try:
            with open(config_path, 'r') as f:
                if config_path.suffix.lower() in ['.yaml', '.yml']:
                    config_dict = yaml.safe_load(f)
                elif config_path.suffix.lower() == '.json':
                    config_dict = json.load(f)
                else:
                    raise ValueError(f"Unsupported config format: {config_path.suffix}")
            
            # Update configuration
            self._update_config_from_dict(config_dict)
            self.config_path = str(config_path)
            
            self.logger.info(f"✅ Configuration loaded from: {config_path}")
            
        except Exception as e:
            self.logger.error(f"❌ Failed to load configuration: {e}")
            raise
        
        return self
    
    def save_config(self, config_path: str, format: str = "yaml") -> 'ConfigManager':
        """
        Save configuration to file
        
        Args:
            config_path: Path to save configuration
            format: File format ('yaml' or 'json')
            
        Returns:
            Self for method chaining
        """
        config_path = Path(config_path)
        config_path.parent.mkdir(parents=True, exist_ok=True)
        
        try:
            config_dict = self._config_to_dict()
            
            with open(config_path, 'w') as f:
                if format.lower() == 'yaml':
                    yaml.dump(config_dict, f, default_flow_style=False, indent=2)
                elif format.lower() == 'json':
                    json.dump(config_dict, f, indent=2)
                else:
                    raise ValueError(f"Unsupported format: {format}")
            
            self.logger.info(f"💾 Configuration saved to: {config_path}")
            
        except Exception as e:
            self.logger.error(f"❌ Failed to save configuration: {e}")
            raise
        
        return self
    
    def _update_config_from_dict(self, config_dict: Dict[str, Any]):
        """Update configuration from dictionary"""
        for section_name, section_config in config_dict.items():
            if hasattr(self.config, section_name):
                section = getattr(self.config, section_name)
                
                if isinstance(section, dict):
                    section.update(section_config)
                else:
                    # Update dataclass fields
                    for key, value in section_config.items():
                        if hasattr(section, key):
                            setattr(section, key, value)
                        else:
                            self.logger.warning(f"Unknown config key: {section_name}.{key}")
    
    def _config_to_dict(self) -> Dict[str, Any]:
        """Convert configuration to dictionary"""
        config_dict = {}
        
        for field_name in self.config.__dataclass_fields__:
            field_value = getattr(self.config, field_name)
            
            if hasattr(field_value, '__dataclass_fields__'):
                # Nested dataclass
                config_dict[field_name] = asdict(field_value)
            else:
                config_dict[field_name] = field_value
        
        return config_dict
    
    def validate_config(self) -> bool:
        """
        Validate configuration parameters
        
        Returns:
            True if configuration is valid
        """
        errors = []
        
        # Validate model config
        if self.config.model.vocab_size <= 0:
            errors.append("vocab_size must be positive")
        
        if self.config.model.hidden_dim <= 0:
            errors.append("hidden_dim must be positive")
        
        if self.config.model.num_layers <= 0:
            errors.append("num_layers must be positive")
        
        if self.config.model.nhead <= 0:
            errors.append("nhead must be positive")
        
        # Validate training config
        if self.config.training.epochs <= 0:
            errors.append("epochs must be positive")
        
        if self.config.training.batch_size <= 0:
            errors.append("batch_size must be positive")
        
        if self.config.training.learning_rate <= 0:
            errors.append("learning_rate must be positive")
        
        # Validate data config
        if self.config.data.target_height <= 0:
            errors.append("target_height must be positive")
        
        if self.config.data.target_width <= 0:
            errors.append("target_width must be positive")
        
        if len(self.config.data.mean) != 3:
            errors.append("mean must have 3 values (RGB)")
        
        if len(self.config.data.std) != 3:
            errors.append("std must have 3 values (RGB)")
        
        # Report errors
        if errors:
            self.logger.error("❌ Configuration validation failed:")
            for error in errors:
                self.logger.error(f"   • {error}")
            return False
        
        self.logger.info("✅ Configuration validation passed")
        return True
    
    def get_config(self) -> Config:
        """Get current configuration"""
        return self.config
    
    def update_config(self, **kwargs) -> 'ConfigManager':
        """
        Update configuration parameters
        
        Args:
            **kwargs: Configuration parameters to update
            
        Returns:
            Self for method chaining
        """
        for key, value in kwargs.items():
            if '.' in key:
                # Nested parameter (e.g., 'model.hidden_dim')
                section_name, param_name = key.split('.', 1)
                if hasattr(self.config, section_name):
                    section = getattr(self.config, section_name)
                    if hasattr(section, param_name):
                        setattr(section, param_name, value)
                    else:
                        self.logger.warning(f"Unknown parameter: {key}")
                else:
                    self.logger.warning(f"Unknown section: {section_name}")
            else:
                # Top-level parameter
                if hasattr(self.config, key):
                    setattr(self.config, key, value)
                else:
                    self.logger.warning(f"Unknown parameter: {key}")
        
        return self
    
    def print_config(self):
        """Print current configuration"""
        print("\n" + "="*60)
        print("📋 CURRENT CONFIGURATION")
        print("="*60)
        
        config_dict = self._config_to_dict()
        self._print_dict(config_dict, indent=0)
        
        print("="*60)
    
    def _print_dict(self, d: Dict[str, Any], indent: int = 0):
        """Recursively print dictionary"""
        for key, value in d.items():
            if isinstance(value, dict):
                print("  " * indent + f"{key}:")
                self._print_dict(value, indent + 1)
            else:
                print("  " * indent + f"{key}: {value}")

# ============================================================
# Predefined Configurations
# ============================================================

def get_default_config() -> Config:
    """Get default configuration"""
    return Config()

def get_training_config() -> Config:
    """Get configuration optimized for training"""
    config = Config()
    
    # Training optimizations
    config.training.batch_size = 32
    config.training.learning_rate = 5e-4
    config.training.use_amp = True
    config.training.accumulation_steps = 2
    
    # Data augmentation
    config.data.use_augmentation = True
    config.data.augmentation_prob = 0.7
    
    return config

def get_inference_config() -> Config:
    """Get configuration optimized for inference"""
    config = Config()
    
    # Inference optimizations
    config.inference.compile_model = True
    config.inference.batch_size = 1
    config.inference.beam_width = 5
    
    # Disable augmentation for inference
    config.data.use_augmentation = False
    
    return config

def get_lightweight_config() -> Config:
    """Get lightweight configuration for resource-constrained environments"""
    config = Config()
    
    # Smaller model
    config.model.hidden_dim = 128
    config.model.num_layers = 2
    config.model.nhead = 4
    
    # Smaller batch size
    config.training.batch_size = 8
    config.inference.batch_size = 1
    
    # Disable expensive features
    config.training.use_amp = False
    config.inference.compile_model = False
    
    return config

def get_high_accuracy_config() -> Config:
    """Get configuration optimized for high accuracy"""
    config = Config()
    
    # Larger model
    config.model.hidden_dim = 512
    config.model.num_layers = 6
    config.model.nhead = 16
    
    # More training
    config.training.epochs = 50
    config.training.learning_rate = 1e-4
    
    # Better decoding
    config.inference.beam_width = 10
    config.inference.default_method = "beam"
    
    return config

# ============================================================
# Configuration Factory
# ============================================================

class ConfigFactory:
    """Factory for creating configurations"""
    
    @staticmethod
    def create_config(config_type: str = "default", **kwargs) -> Config:
        """
        Create configuration by type
        
        Args:
            config_type: Type of configuration ('default', 'training', 'inference', 'lightweight', 'high_accuracy')
            **kwargs: Additional parameters to override
            
        Returns:
            Configuration object
        """
        configs = {
            "default": get_default_config,
            "training": get_training_config,
            "inference": get_inference_config,
            "lightweight": get_lightweight_config,
            "high_accuracy": get_high_accuracy_config
        }
        
        if config_type not in configs:
            raise ValueError(f"Unknown config type: {config_type}")
        
        config = configs[config_type]()
        
        # Apply overrides
        for key, value in kwargs.items():
            if '.' in key:
                section_name, param_name = key.split('.', 1)
                if hasattr(config, section_name):
                    section = getattr(config, section_name)
                    if hasattr(section, param_name):
                        setattr(section, param_name, value)
            else:
                if hasattr(config, key):
                    setattr(config, key, value)
        
        return config
    
    @staticmethod
    def save_preset_configs(output_dir: str = "configs"):
        """Save all preset configurations to files"""
        output_dir = Path(output_dir)
        output_dir.mkdir(exist_ok=True)
        
        configs = {
            "default": get_default_config,
            "training": get_training_config,
            "inference": get_inference_config,
            "lightweight": get_lightweight_config,
            "high_accuracy": get_high_accuracy_config
        }
        
        for name, config_func in configs.items():
            config = config_func()
            config_dict = asdict(config)
            
            # Save as YAML
            yaml_path = output_dir / f"{name}.yaml"
            with open(yaml_path, 'w') as f:
                yaml.dump(config_dict, f, default_flow_style=False, indent=2)
            
            # Save as JSON
            json_path = output_dir / f"{name}.json"
            with open(json_path, 'w') as f:
                json.dump(config_dict, f, indent=2)
            
            print(f"💾 Saved {name} config to {yaml_path} and {json_path}")

# ============================================================
# Utility Functions
# ============================================================

def load_config_from_file(config_path: str) -> Config:
    """Load configuration from file"""
    manager = ConfigManager(config_path)
    return manager.get_config()

def save_config_to_file(config: Config, config_path: str, format: str = "yaml"):
    """Save configuration to file"""
    manager = ConfigManager()
    manager.config = config
    manager.save_config(config_path, format)

def merge_configs(base_config: Config, override_config: Dict[str, Any]) -> Config:
    """Merge configuration with overrides"""
    manager = ConfigManager()
    manager.config = base_config
    
    for key, value in override_config.items():
        manager.update_config(**{key: value})
    
    return manager.get_config()

# ============================================================
# Testing Functions
# ============================================================

def test_config_system():
    """Test configuration system"""
    print("🧪 Testing Configuration System...")
    
    # Test default config
    config = get_default_config()
    print("✅ Default config created")
    
    # Test config manager
    manager = ConfigManager()
    manager.config = config
    print("✅ Config manager created")
    
    # Test validation
    is_valid = manager.validate_config()
    print(f"✅ Config validation: {'PASSED' if is_valid else 'FAILED'}")
    
    # Test saving and loading
    test_config_path = "test_config.yaml"
    manager.save_config(test_config_path)
    print("✅ Config saved")
    
    manager2 = ConfigManager(test_config_path)
    print("✅ Config loaded")
    
    # Test factory
    training_config = ConfigFactory.create_config("training")
    print("✅ Training config created")
    
    # Cleanup
    if Path(test_config_path).exists():
        Path(test_config_path).unlink()
    
    print("🎉 All configuration tests passed!")

if __name__ == "__main__":
    test_config_system()