Trixy3/trixy_core/ml/wakeword/config.py

"""
Configuration management for wakeword detection system.

This module provides comprehensive configuration management for all aspects
of the wakeword detection pipeline including model loading, audio processing,
detection parameters, and performance tuning.
"""

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

from .audio_features import SpectrogramConfig
from .detection_engine import DetectionConfig


@dataclass
class ModelConfig:
    """Configuration for model loading and management."""
    
    # Model file paths
    model_path: str = ""
    model_password: str = ""
    backup_model_paths: List[str] = field(default_factory=list)
    
    # Model selection
    model_type: str = "auto"  # "auto", "standard", "improved", "lightweight"
    use_model_optimization: bool = True
    verify_model_hash: bool = True
    
    # Device configuration
    device: str = "auto"  # "auto", "cpu", "cuda", "mps"
    use_half_precision: bool = False
    enable_model_compilation: bool = True
    
    # Fallback behavior
    fallback_on_error: bool = True
    fallback_model_path: str = ""
    
    def validate(self) -> bool:
        """Validate model configuration."""
        if not self.model_path:
            return False
        
        # Check if model file exists
        if not Path(self.model_path).exists():
            return False
        
        # Validate device setting
        valid_devices = ["auto", "cpu", "cuda", "mps"]
        if self.device not in valid_devices:
            return False
        
        return True


@dataclass  
class AudioConfig:
    """Configuration for audio processing."""
    
    # Audio input parameters
    sample_rate: int = 16000
    channels: int = 1  # Mono audio
    bit_depth: int = 16
    
    # Buffer configuration
    buffer_duration: float = 2.0  # seconds
    chunk_duration: float = 1.0   # seconds for processing
    overlap_ratio: float = 0.5
    
    # Feature extraction
    spectrogram_config: SpectrogramConfig = field(default_factory=SpectrogramConfig)
    
    # Audio preprocessing
    enable_noise_reduction: bool = False
    enable_gain_control: bool = True
    enable_voice_activity_detection: bool = False
    
    # Streaming parameters
    max_chunks_in_memory: int = 10
    use_circular_buffer: bool = True
    
    def validate(self) -> bool:
        """Validate audio configuration."""
        if self.sample_rate <= 0:
            return False
        if self.channels not in [1, 2]:
            return False
        if self.bit_depth not in [16, 24, 32]:
            return False
        if self.buffer_duration <= 0 or self.chunk_duration <= 0:
            return False
        if not (0.0 <= self.overlap_ratio <= 1.0):
            return False
        
        return self.spectrogram_config and hasattr(self.spectrogram_config, 'validate')


@dataclass
class PerformanceConfig:
    """Configuration for performance optimization."""
    
    # Threading and concurrency
    enable_threading: bool = True
    max_worker_threads: int = 2
    thread_priority: str = "normal"  # "low", "normal", "high"
    
    # Memory management
    max_memory_usage_mb: float = 512.0
    enable_memory_pool: bool = True
    gc_frequency: int = 100  # Garbage collection every N chunks
    
    # Inference optimization
    enable_batch_inference: bool = False
    max_batch_size: int = 8
    inference_timeout_ms: float = 100.0
    
    # Monitoring and logging
    enable_performance_monitoring: bool = True
    log_performance_stats: bool = False
    stats_update_interval: int = 10  # seconds
    
    def validate(self) -> bool:
        """Validate performance configuration."""
        if self.max_worker_threads < 1:
            return False
        if self.max_memory_usage_mb <= 0:
            return False
        if self.max_batch_size < 1:
            return False
        if self.inference_timeout_ms <= 0:
            return False
        
        return True


@dataclass
class WakewordConfig:
    """
    Complete configuration for wakeword detection system.
    
    This is the main configuration class that encompasses all aspects
    of the wakeword detection pipeline.
    """
    
    # Sub-configurations
    model_config: ModelConfig = field(default_factory=ModelConfig)
    audio_config: AudioConfig = field(default_factory=AudioConfig)
    detection_config: DetectionConfig = field(default_factory=DetectionConfig)
    performance_config: PerformanceConfig = field(default_factory=PerformanceConfig)
    
    # System integration
    enable_event_system: bool = True
    event_handler: Optional[Any] = None
    satellite_id: str = ""
    
    # Debugging and development
    debug_mode: bool = False
    save_debug_audio: bool = False
    debug_output_path: str = "/tmp/trixy_wakeword_debug"
    
    # Feature flags
    enable_voice_recognition_integration: bool = False
    enable_conversation_context: bool = True
    enable_adaptive_thresholds: bool = False
    
    def validate(self) -> bool:
        """Validate complete configuration."""
        if not self.model_config.validate():
            return False
        if not self.audio_config.validate():
            return False
        if not self.detection_config.validate():
            return False
        if not self.performance_config.validate():
            return False
        
        return True
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert configuration to dictionary."""
        return {
            'model_config': asdict(self.model_config),
            'audio_config': asdict(self.audio_config),
            'detection_config': asdict(self.detection_config),
            'performance_config': asdict(self.performance_config),
            'enable_event_system': self.enable_event_system,
            'satellite_id': self.satellite_id,
            'debug_mode': self.debug_mode,
            'save_debug_audio': self.save_debug_audio,
            'debug_output_path': self.debug_output_path,
            'enable_voice_recognition_integration': self.enable_voice_recognition_integration,
            'enable_conversation_context': self.enable_conversation_context,
            'enable_adaptive_thresholds': self.enable_adaptive_thresholds
        }
    
    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]) -> 'WakewordConfig':
        """Create configuration from dictionary."""
        # Extract sub-configurations
        model_config = ModelConfig(**config_dict.get('model_config', {}))
        audio_config = AudioConfig(**config_dict.get('audio_config', {}))
        detection_config = DetectionConfig(**config_dict.get('detection_config', {}))
        performance_config = PerformanceConfig(**config_dict.get('performance_config', {}))
        
        # Create main config
        return cls(
            model_config=model_config,
            audio_config=audio_config,
            detection_config=detection_config,
            performance_config=performance_config,
            enable_event_system=config_dict.get('enable_event_system', True),
            satellite_id=config_dict.get('satellite_id', ''),
            debug_mode=config_dict.get('debug_mode', False),
            save_debug_audio=config_dict.get('save_debug_audio', False),
            debug_output_path=config_dict.get('debug_output_path', '/tmp/trixy_wakeword_debug'),
            enable_voice_recognition_integration=config_dict.get('enable_voice_recognition_integration', False),
            enable_conversation_context=config_dict.get('enable_conversation_context', True),
            enable_adaptive_thresholds=config_dict.get('enable_adaptive_thresholds', False)
        )
    
    def save_to_file(self, file_path: Union[str, Path]):
        """Save configuration to JSON file."""
        file_path = Path(file_path)
        file_path.parent.mkdir(parents=True, exist_ok=True)
        
        with open(file_path, 'w') as f:
            json.dump(self.to_dict(), f, indent=2)
    
    @classmethod
    def load_from_file(cls, file_path: Union[str, Path]) -> 'WakewordConfig':
        """Load configuration from JSON file."""
        file_path = Path(file_path)
        
        if not file_path.exists():
            raise FileNotFoundError(f"Configuration file not found: {file_path}")
        
        with open(file_path, 'r') as f:
            config_dict = json.load(f)
        
        return cls.from_dict(config_dict)
    
    def update_from_env(self, prefix: str = "TRIXY_WAKEWORD_"):
        """Update configuration from environment variables."""
        env_mappings = {
            f"{prefix}MODEL_PATH": ("model_config", "model_path"),
            f"{prefix}MODEL_PASSWORD": ("model_config", "model_password"),
            f"{prefix}DEVICE": ("model_config", "device"),
            f"{prefix}SAMPLE_RATE": ("audio_config", "sample_rate"),
            f"{prefix}CUSTOM_THRESHOLD": ("detection_config", "custom_threshold"),
            f"{prefix}SYSTEM_THRESHOLD": ("detection_config", "system_command_threshold"),
            f"{prefix}DEBUG_MODE": ("debug_mode", None),
            f"{prefix}SATELLITE_ID": ("satellite_id", None),
        }
        
        for env_var, (section, key) in env_mappings.items():
            value = os.getenv(env_var)
            if value is not None:
                if section == "debug_mode" or section == "satellite_id":
                    # Direct attribute
                    if section == "debug_mode":
                        setattr(self, section, value.lower() in ['true', '1', 'yes'])
                    else:
                        setattr(self, section, value)
                else:
                    # Sub-configuration attribute
                    config_obj = getattr(self, section)
                    if hasattr(config_obj, key):
                        current_value = getattr(config_obj, key)
                        # Type conversion based on current value type
                        if isinstance(current_value, bool):
                            setattr(config_obj, key, value.lower() in ['true', '1', 'yes'])
                        elif isinstance(current_value, int):
                            setattr(config_obj, key, int(value))
                        elif isinstance(current_value, float):
                            setattr(config_obj, key, float(value))
                        else:
                            setattr(config_obj, key, value)


def create_default_config() -> WakewordConfig:
    """Create a default wakeword configuration."""
    return WakewordConfig(
        model_config=ModelConfig(
            model_path="models/wakeword/default/model.pth",
            model_password="",
            device="auto",
            use_model_optimization=True
        ),
        audio_config=AudioConfig(
            sample_rate=16000,
            buffer_duration=2.0,
            chunk_duration=1.0,
            overlap_ratio=0.5
        ),
        detection_config=DetectionConfig(
            custom_threshold=0.7,
            system_command_threshold=0.8,
            use_temporal_smoothing=True,
            smoothing_window_size=5
        ),
        performance_config=PerformanceConfig(
            enable_threading=True,
            max_worker_threads=2,
            enable_performance_monitoring=True
        )
    )


def create_lightweight_config() -> WakewordConfig:
    """Create a lightweight configuration for resource-constrained devices."""
    config = create_default_config()
    
    # Optimize for low resource usage
    config.model_config.use_half_precision = True
    config.model_config.enable_model_compilation = False
    
    config.audio_config.buffer_duration = 1.5
    config.audio_config.max_chunks_in_memory = 5
    
    config.detection_config.use_temporal_smoothing = False
    config.detection_config.batch_inference = False
    
    config.performance_config.max_worker_threads = 1
    config.performance_config.max_memory_usage_mb = 128.0
    config.performance_config.enable_memory_pool = True
    
    return config


def create_high_performance_config() -> WakewordConfig:
    """Create a high-performance configuration for powerful devices."""
    config = create_default_config()
    
    # Optimize for high performance
    config.model_config.use_model_optimization = True
    config.model_config.enable_model_compilation = True
    
    config.audio_config.buffer_duration = 3.0
    config.audio_config.max_chunks_in_memory = 20
    
    config.detection_config.use_temporal_smoothing = True
    config.detection_config.batch_inference = True
    config.detection_config.enable_model_optimization = True
    
    config.performance_config.max_worker_threads = 4
    config.performance_config.max_memory_usage_mb = 1024.0
    config.performance_config.enable_batch_inference = True
    config.performance_config.max_batch_size = 16
    
    return config


def validate_config_compatibility(config: WakewordConfig) -> List[str]:
    """
    Validate configuration for compatibility and potential issues.
    
    Args:
        config: Configuration to validate
        
    Returns:
        List of warning/error messages
    """
    warnings = []
    
    # Check model and audio compatibility
    if config.model_config.use_half_precision and config.model_config.device == "cpu":
        warnings.append("Half precision not supported on CPU, will be disabled")
    
    # Check buffer and chunk size compatibility
    if config.audio_config.chunk_duration >= config.audio_config.buffer_duration:
        warnings.append("Chunk duration should be smaller than buffer duration")
    
    # Check memory constraints
    estimated_memory = estimate_memory_usage(config)
    if estimated_memory > config.performance_config.max_memory_usage_mb:
        warnings.append(f"Estimated memory usage ({estimated_memory:.1f}MB) exceeds limit")
    
    # Check performance settings
    if config.detection_config.batch_inference and not config.performance_config.enable_batch_inference:
        warnings.append("Batch inference enabled in detection but disabled in performance config")
    
    # Check temporal smoothing compatibility
    if config.detection_config.use_temporal_smoothing and config.performance_config.max_worker_threads < 2:
        warnings.append("Temporal smoothing may need multiple threads for optimal performance")
    
    return warnings


def estimate_memory_usage(config: WakewordConfig) -> float:
    """
    Estimate memory usage in MB for the given configuration.
    
    Args:
        config: Configuration to analyze
        
    Returns:
        Estimated memory usage in MB
    """
    # Audio buffer memory
    samples_per_second = config.audio_config.sample_rate
    buffer_samples = int(config.audio_config.buffer_duration * samples_per_second)
    chunk_samples = int(config.audio_config.chunk_duration * samples_per_second)
    
    # Assuming 32-bit float samples
    buffer_memory = buffer_samples * 4 / (1024 * 1024)  # MB
    chunk_memory = chunk_samples * config.audio_config.max_chunks_in_memory * 4 / (1024 * 1024)
    
    # Feature extraction memory (spectrograms)
    spec_config = config.audio_config.spectrogram_config
    spec_memory = spec_config.n_mels * spec_config.time_frames * 4 / (1024 * 1024)
    
    # Model memory (rough estimate)
    model_memory = 50.0  # MB - rough estimate for RepCNN models
    
    # Additional overhead
    overhead = 20.0  # MB
    
    total_memory = buffer_memory + chunk_memory + spec_memory + model_memory + overhead
    
    return total_memory


def auto_tune_config(config: WakewordConfig, target_latency_ms: float = 100.0,
                    available_memory_mb: float = 512.0) -> WakewordConfig:
    """
    Automatically tune configuration parameters for target performance.
    
    Args:
        config: Base configuration to tune
        target_latency_ms: Target inference latency in milliseconds
        available_memory_mb: Available memory in MB
        
    Returns:
        Tuned configuration
    """
    tuned_config = WakewordConfig.from_dict(config.to_dict())
    
    # Adjust for latency target
    if target_latency_ms < 50.0:
        # Very low latency requirements
        tuned_config.audio_config.chunk_duration = 0.5
        tuned_config.detection_config.use_temporal_smoothing = False
        tuned_config.model_config.use_half_precision = True
        tuned_config.performance_config.max_worker_threads = 1
    elif target_latency_ms > 200.0:
        # Can afford higher latency for better accuracy
        tuned_config.audio_config.chunk_duration = 1.5
        tuned_config.detection_config.use_temporal_smoothing = True
        tuned_config.detection_config.smoothing_window_size = 7
    
    # Adjust for memory constraints
    estimated_memory = estimate_memory_usage(tuned_config)
    if estimated_memory > available_memory_mb:
        # Reduce memory usage
        tuned_config.audio_config.buffer_duration = max(1.0, tuned_config.audio_config.buffer_duration * 0.7)
        tuned_config.audio_config.max_chunks_in_memory = max(3, tuned_config.audio_config.max_chunks_in_memory // 2)
        tuned_config.performance_config.max_memory_usage_mb = available_memory_mb * 0.8
    
    return tuned_config