Trixy3/trixy_core/ml/voice_recognition/config.py

"""
Configuration management for Voice Recognition system.

This module provides comprehensive configuration classes for the voice recognition
system, including model configuration, audio processing settings, recognition
parameters, and performance optimization options.
"""

import os
import json
import logging
from typing import Dict, List, Optional, Union, Any, Tuple
from dataclasses import dataclass, field, asdict
from pathlib import Path
from enum import Enum
import torch

# Configure logging
logger = logging.getLogger(__name__)


class ModelArchitectureType(Enum):
    """Supported model architecture types."""
    ECAPA_TDNN = "ecapa_tdnn"
    TITANET_S = "titanet_s"
    SPEAKERNET_M = "speakernet_m"


class LossType(Enum):
    """Supported loss function types."""
    ANGULAR_MARGIN = "angular_margin"  # ArcFace
    GE2E = "ge2e"
    TRIPLET = "triplet"
    CONTRASTIVE = "contrastive"


class FeatureType(Enum):
    """Supported audio feature types."""
    LOG_MEL = "log_mel"
    MFCC = "mfcc"


class VADType(Enum):
    """Supported Voice Activity Detection types."""
    WEBRTC = "webrtc"
    ENERGY_BASED = "energy_based"
    COMBINED = "combined"


@dataclass
class ModelConfig:
    """Configuration for voice recognition model."""
    model_path: str
    model_password: Optional[str] = None
    architecture_type: ModelArchitectureType = ModelArchitectureType.ECAPA_TDNN
    embedding_dim: int = 192
    num_speakers: int = 4
    speaker_mapping: Dict[str, str] = field(default_factory=dict)
    
    # Architecture-specific parameters
    ecapa_channels: int = 512
    use_attention_pooling: bool = True
    
    titanet_channels: List[int] = field(default_factory=lambda: [128, 256, 512, 1024])
    
    speakernet_hidden_dim: int = 512
    speakernet_num_layers: int = 4
    
    dropout_rate: float = 0.1
    use_batch_norm: bool = True
    
    def __post_init__(self):
        """Validate configuration after initialization."""
        if isinstance(self.architecture_type, str):
            self.architecture_type = ModelArchitectureType(self.architecture_type)
        
        if self.num_speakers < 2:
            raise ValueError("Number of speakers must be at least 2")
        if self.num_speakers > 200:
            logger.warning(f"Large number of speakers ({self.num_speakers}) may impact performance")
        
        if self.embedding_dim < 64:
            raise ValueError("Embedding dimension must be at least 64")


@dataclass
class AudioConfig:
    """Configuration for audio processing."""
    sample_rate: int = 16000
    audio_length: float = 3.0  # seconds
    feature_type: FeatureType = FeatureType.LOG_MEL
    
    # Log-Mel Spectrogram parameters
    n_mels: int = 40
    n_fft: int = 512
    hop_length: int = 160  # 10ms at 16kHz
    win_length: int = 400  # 25ms at 16kHz
    f_min: float = 80.0
    f_max: Optional[float] = None
    
    # MFCC parameters
    n_mfcc: int = 40
    dct_type: int = 2
    norm: str = "ortho"
    log_mels: bool = False
    
    # Voice Activity Detection
    use_vad: bool = True
    vad_type: VADType = VADType.WEBRTC
    vad_aggressiveness: int = 2  # 0-3 for WebRTC VAD
    vad_frame_duration: int = 30  # ms
    energy_threshold: float = 0.01
    
    # Data augmentation
    use_augmentation: bool = True
    noise_factor: float = 0.05
    speed_factor: float = 0.05
    pitch_factor: float = 0.02
    volume_factor: float = 0.1
    reverb_factor: float = 0.03
    
    def __post_init__(self):
        """Validate configuration after initialization."""
        if isinstance(self.feature_type, str):
            self.feature_type = FeatureType(self.feature_type)
        if isinstance(self.vad_type, str):
            self.vad_type = VADType(self.vad_type)
        
        if self.sample_rate not in [8000, 16000, 22050, 44100, 48000]:
            logger.warning(f"Unusual sample rate: {self.sample_rate}Hz")
        
        if self.audio_length < 0.5:
            raise ValueError("Audio length must be at least 0.5 seconds")
        if self.audio_length > 10.0:
            logger.warning(f"Long audio length ({self.audio_length}s) may impact performance")


@dataclass
class RecognitionConfig:
    """Configuration for voice recognition engine."""
    confidence_threshold: float = 0.7
    verification_threshold: float = 0.6
    identification_threshold: float = 0.5
    
    # Loss function configuration
    loss_type: LossType = LossType.ANGULAR_MARGIN
    angular_margin: float = 0.5
    angular_scale: float = 64.0
    ge2e_init_w: float = 10.0
    ge2e_init_b: float = -5.0
    contrastive_margin: float = 2.0
    triplet_margin: float = 0.3
    
    # Similarity calculation
    similarity_metric: str = "cosine"  # "cosine", "euclidean", "dot_product"
    use_score_normalization: bool = True
    
    # Temporal filtering
    use_temporal_smoothing: bool = True
    smoothing_window_size: int = 5
    confidence_decay_factor: float = 0.9
    
    # Anti-spoofing
    use_anti_spoofing: bool = False
    spoofing_threshold: float = 0.5
    
    def __post_init__(self):
        """Validate configuration after initialization."""
        if isinstance(self.loss_type, str):
            self.loss_type = LossType(self.loss_type)
        
        if not 0.0 <= self.confidence_threshold <= 1.0:
            raise ValueError("Confidence threshold must be between 0.0 and 1.0")
        if not 0.0 <= self.verification_threshold <= 1.0:
            raise ValueError("Verification threshold must be between 0.0 and 1.0")
        if not 0.0 <= self.identification_threshold <= 1.0:
            raise ValueError("Identification threshold must be between 0.0 and 1.0")


@dataclass
class PerformanceConfig:
    """Configuration for performance optimization."""
    device: str = "auto"  # "auto", "cpu", "cuda", "mps"
    batch_size: int = 1
    use_half_precision: bool = False
    use_torch_compile: bool = False
    
    # Memory management
    max_memory_usage_mb: int = 512
    clear_cache_interval: int = 100
    
    # Latency optimization
    target_latency_ms: float = 100.0
    use_streaming_inference: bool = True
    chunk_size_ms: float = 30.0
    
    # Threading
    num_workers: int = 2
    pin_memory: bool = True
    
    # Model optimization
    use_quantization: bool = False
    quantization_backend: str = "fbgemm"  # "fbgemm", "qnnpack"
    use_tensorrt: bool = False
    
    def __post_init__(self):
        """Validate and set device configuration."""
        if self.device == "auto":
            if torch.cuda.is_available():
                self.device = "cuda"
            elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                self.device = "mps"
            else:
                self.device = "cpu"
        
        if self.batch_size < 1:
            raise ValueError("Batch size must be at least 1")
        if self.batch_size > 32:
            logger.warning(f"Large batch size ({self.batch_size}) may cause memory issues")


@dataclass
class VoiceRecognitionConfig:
    """Main configuration class for voice recognition system."""
    model_config: ModelConfig
    audio_config: AudioConfig = field(default_factory=AudioConfig)
    recognition_config: RecognitionConfig = field(default_factory=RecognitionConfig)
    performance_config: PerformanceConfig = field(default_factory=PerformanceConfig)
    
    # System configuration
    debug_mode: bool = False
    log_level: str = "INFO"
    save_debug_audio: bool = False
    debug_audio_dir: str = "./debug/voice_recognition"
    
    def __post_init__(self):
        """Validate complete configuration."""
        # Ensure audio length matches expected input for model
        expected_samples = int(self.audio_config.audio_length * self.audio_config.sample_rate)
        if expected_samples < 8000:  # Minimum for reliable recognition
            logger.warning(f"Short audio length may impact recognition accuracy")
        
        # Validate feature extraction parameters
        if self.audio_config.feature_type == FeatureType.LOG_MEL:
            if self.audio_config.n_mels != 40:
                logger.warning("Standard n_mels is 40 for voice recognition")
        elif self.audio_config.feature_type == FeatureType.MFCC:
            if self.audio_config.n_mfcc != 40:
                logger.warning("Standard n_mfcc is 40 for voice recognition")
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert configuration to dictionary."""
        return asdict(self)
    
    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]) -> 'VoiceRecognitionConfig':
        """Create configuration from dictionary."""
        model_config = ModelConfig(**config_dict.get('model_config', {}))
        audio_config = AudioConfig(**config_dict.get('audio_config', {}))
        recognition_config = RecognitionConfig(**config_dict.get('recognition_config', {}))
        performance_config = PerformanceConfig(**config_dict.get('performance_config', {}))
        
        # Remove nested configs from main dict
        main_config = {k: v for k, v in config_dict.items() 
                      if k not in ['model_config', 'audio_config', 'recognition_config', 'performance_config']}
        
        return cls(
            model_config=model_config,
            audio_config=audio_config,
            recognition_config=recognition_config,
            performance_config=performance_config,
            **main_config
        )
    
    def save(self, file_path: str) -> None:
        """Save configuration to JSON file."""
        config_dict = self.to_dict()
        
        # Convert enums to strings for JSON serialization
        def convert_enums(obj):
            if isinstance(obj, dict):
                return {k: convert_enums(v) for k, v in obj.items()}
            elif isinstance(obj, list):
                return [convert_enums(item) for item in obj]
            elif isinstance(obj, Enum):
                return obj.value
            return obj
        
        config_dict = convert_enums(config_dict)
        
        Path(file_path).parent.mkdir(parents=True, exist_ok=True)
        with open(file_path, 'w') as f:
            json.dump(config_dict, f, indent=2)
        
        logger.info(f"Configuration saved to {file_path}")
    
    @classmethod
    def load(cls, file_path: str) -> 'VoiceRecognitionConfig':
        """Load configuration from JSON file."""
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Configuration file not found: {file_path}")
        
        with open(file_path, 'r') as f:
            config_dict = json.load(f)
        
        logger.info(f"Configuration loaded from {file_path}")
        return cls.from_dict(config_dict)


def create_default_config(model_path: str, model_password: Optional[str] = None) -> VoiceRecognitionConfig:
    """Create default configuration for voice recognition."""
    model_config = ModelConfig(
        model_path=model_path,
        model_password=model_password,
        architecture_type=ModelArchitectureType.ECAPA_TDNN,
        embedding_dim=192
    )
    
    return VoiceRecognitionConfig(model_config=model_config)


def create_titanet_config(model_path: str, model_password: Optional[str] = None) -> VoiceRecognitionConfig:
    """Create TitaNet-S optimized configuration."""
    model_config = ModelConfig(
        model_path=model_path,
        model_password=model_password,
        architecture_type=ModelArchitectureType.TITANET_S,
        embedding_dim=192,
        titanet_channels=[128, 256, 512, 1024],
        dropout_rate=0.05
    )
    
    audio_config = AudioConfig(
        audio_length=2.5,
        use_augmentation=True,
        noise_factor=0.03
    )
    
    recognition_config = RecognitionConfig(
        loss_type=LossType.GE2E,
        confidence_threshold=0.6,
        verification_threshold=0.6
    )
    
    performance_config = PerformanceConfig(
        batch_size=64,
        use_half_precision=True
    )
    
    return VoiceRecognitionConfig(
        model_config=model_config,
        audio_config=audio_config,
        recognition_config=recognition_config,
        performance_config=performance_config
    )


def create_speakernet_config(model_path: str, model_password: Optional[str] = None) -> VoiceRecognitionConfig:
    """Create SpeakerNet-M optimized configuration."""
    model_config = ModelConfig(
        model_path=model_path,
        model_password=model_password,
        architecture_type=ModelArchitectureType.SPEAKERNET_M,
        embedding_dim=256,
        speakernet_hidden_dim=512,
        speakernet_num_layers=4,
        dropout_rate=0.1
    )
    
    audio_config = AudioConfig(
        feature_type=FeatureType.MFCC,
        n_mfcc=40,
        audio_length=3.0
    )
    
    recognition_config = RecognitionConfig(
        loss_type=LossType.TRIPLET,
        triplet_margin=0.3,
        confidence_threshold=0.65
    )
    
    return VoiceRecognitionConfig(
        model_config=model_config,
        audio_config=audio_config,
        recognition_config=recognition_config
    )


def auto_tune_config(config: VoiceRecognitionConfig, 
                    target_latency_ms: float = 100.0,
                    target_accuracy: float = 0.9) -> VoiceRecognitionConfig:
    """Automatically tune configuration for performance and accuracy targets."""
    # Adjust batch size based on target latency
    if target_latency_ms < 50:
        config.performance_config.batch_size = 1
        config.performance_config.use_half_precision = True
        config.audio_config.audio_length = min(config.audio_config.audio_length, 2.0)
    elif target_latency_ms < 100:
        config.performance_config.batch_size = min(config.performance_config.batch_size, 4)
        config.performance_config.use_half_precision = True
    
    # Adjust thresholds based on target accuracy
    if target_accuracy >= 0.95:
        config.recognition_config.confidence_threshold = max(0.8, config.recognition_config.confidence_threshold)
        config.recognition_config.verification_threshold = max(0.75, config.recognition_config.verification_threshold)
        config.audio_config.use_augmentation = True
    elif target_accuracy <= 0.8:
        config.recognition_config.confidence_threshold = min(0.6, config.recognition_config.confidence_threshold)
        config.recognition_config.verification_threshold = min(0.55, config.recognition_config.verification_threshold)
    
    config.performance_config.target_latency_ms = target_latency_ms
    
    logger.info(f"Configuration auto-tuned for {target_latency_ms}ms latency and {target_accuracy} accuracy")
    return config


def validate_config_compatibility(config: VoiceRecognitionConfig) -> List[str]:
    """Validate configuration compatibility and return list of warnings/errors."""
    warnings = []
    
    # Check device compatibility
    if config.performance_config.device == "cuda" and not torch.cuda.is_available():
        warnings.append("CUDA device specified but not available, falling back to CPU")
        config.performance_config.device = "cpu"
    
    # Check memory requirements
    estimated_memory = estimate_memory_usage(config)
    if estimated_memory > config.performance_config.max_memory_usage_mb:
        warnings.append(f"Estimated memory usage ({estimated_memory}MB) exceeds limit ({config.performance_config.max_memory_usage_mb}MB)")
    
    # Check audio processing compatibility
    if config.audio_config.use_vad and config.audio_config.vad_type == VADType.WEBRTC:
        try:
            import webrtcvad
        except ImportError:
            warnings.append("WebRTC VAD specified but webrtcvad package not available")
    
    # Check feature extraction compatibility
    hop_length_ms = config.audio_config.hop_length / config.audio_config.sample_rate * 1000
    if hop_length_ms > 15:
        warnings.append(f"Large hop length ({hop_length_ms:.1f}ms) may reduce temporal resolution")
    
    return warnings


def estimate_memory_usage(config: VoiceRecognitionConfig) -> int:
    """Estimate memory usage in MB for the given configuration."""
    # Audio buffer memory
    audio_samples = int(config.audio_config.audio_length * config.audio_config.sample_rate)
    audio_memory = audio_samples * 4 * config.performance_config.batch_size  # float32
    
    # Feature memory
    if config.audio_config.feature_type == FeatureType.LOG_MEL:
        feature_frames = audio_samples // config.audio_config.hop_length
        feature_memory = feature_frames * config.audio_config.n_mels * 4 * config.performance_config.batch_size
    else:  # MFCC
        feature_frames = audio_samples // config.audio_config.hop_length
        feature_memory = feature_frames * config.audio_config.n_mfcc * 4 * config.performance_config.batch_size
    
    # Model memory (rough estimate)
    model_memory = config.model_config.embedding_dim * config.model_config.num_speakers * 4
    if config.model_config.architecture_type == ModelArchitectureType.ECAPA_TDNN:
        model_memory += config.model_config.ecapa_channels * 1000  # Rough estimate
    elif config.model_config.architecture_type == ModelArchitectureType.TITANET_S:
        model_memory += sum(config.model_config.titanet_channels) * 500
    
    total_memory_bytes = audio_memory + feature_memory + model_memory
    return int(total_memory_bytes / (1024 * 1024))  # Convert to MB


# Example configurations for testing
def get_example_configs() -> Dict[str, VoiceRecognitionConfig]:
    """Get example configurations for different use cases."""
    return {
        "default": create_default_config("/path/to/model.pth"),
        "titanet": create_titanet_config("/path/to/titanet_model.pth"),
        "speakernet": create_speakernet_config("/path/to/speakernet_model.pth"),
        "lightweight": auto_tune_config(create_default_config("/path/to/model.pth"), target_latency_ms=50),
        "high_accuracy": auto_tune_config(create_default_config("/path/to/model.pth"), target_accuracy=0.95)
    }