Trixy3/trixy_core/ml/wakeword/detection_engine.py

"""
Wakeword detection engine with confidence scoring and inference optimization.

This module provides the core detection engine that processes audio features
and generates wakeword predictions with confidence scores and threshold management.
"""

import time
import threading
from typing import Optional, Dict, List, Tuple, Callable, Union
from dataclasses import dataclass
from enum import Enum
import warnings

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from .audio_features import AudioFeatureExtractor, SpectrogramConfig
from .audio_buffer import AudioChunk
from .model_loader import ModelMetadata


class WakewordType(Enum):
    """Supported wakeword types."""
    CUSTOM = "custom"
    SYSTEM_COMMAND = "system_command"
    NEGATIVE = "negative"  # Not a wakeword


@dataclass
class DetectionResult:
    """Result of wakeword detection."""
    
    wakeword_type: WakewordType
    confidence: float
    raw_scores: Dict[str, float]
    processing_time_ms: float
    chunk_id: int
    timestamp: float
    features_shape: Tuple[int, ...]
    
    @property
    def is_wakeword_detected(self) -> bool:
        """Check if a valid wakeword was detected."""
        return self.wakeword_type in [WakewordType.CUSTOM, WakewordType.SYSTEM_COMMAND]
    
    def to_dict(self) -> Dict:
        """Convert result to dictionary."""
        return {
            'wakeword_type': self.wakeword_type.value,
            'confidence': self.confidence,
            'raw_scores': self.raw_scores,
            'processing_time_ms': self.processing_time_ms,
            'chunk_id': self.chunk_id,
            'timestamp': self.timestamp,
            'features_shape': list(self.features_shape),
            'is_wakeword_detected': self.is_wakeword_detected
        }


@dataclass
class DetectionConfig:
    """Configuration for wakeword detection."""
    
    # Confidence thresholds
    custom_threshold: float = 0.7
    system_command_threshold: float = 0.8  # Higher threshold for admin commands
    negative_threshold: float = 0.3
    
    # Temporal filtering
    use_temporal_smoothing: bool = True
    smoothing_window_size: int = 5
    min_detections_in_window: int = 3
    
    # Performance optimization
    use_half_precision: bool = False
    enable_model_optimization: bool = True
    batch_inference: bool = False
    max_batch_size: int = 8
    
    # Debugging
    log_all_predictions: bool = False
    save_detection_features: bool = False
    
    def validate(self) -> bool:
        """Validate configuration parameters."""
        if not (0.0 <= self.custom_threshold <= 1.0):
            return False
        if not (0.0 <= self.system_command_threshold <= 1.0):
            return False
        if not (0.0 <= self.negative_threshold <= 1.0):
            return False
        if self.smoothing_window_size < 1:
            return False
        if self.min_detections_in_window < 1:
            return False
        if self.max_batch_size < 1:
            return False
        return True


class ConfidenceCalibrator:
    """
    Calibrate model outputs to well-calibrated probability scores.
    
    Uses temperature scaling and Platt scaling for confidence calibration.
    """
    
    def __init__(self):
        self.temperature = 1.0
        self.platt_a = 1.0
        self.platt_b = 0.0
        self.is_calibrated = False
    
    def calibrate_with_temperature(self, logits: torch.Tensor) -> torch.Tensor:
        """Apply temperature scaling to logits."""
        return logits / self.temperature
    
    def calibrate_with_platt(self, scores: torch.Tensor) -> torch.Tensor:
        """Apply Platt scaling to probability scores."""
        return torch.sigmoid(self.platt_a * scores + self.platt_b)
    
    def set_temperature(self, temperature: float):
        """Set temperature parameter."""
        self.temperature = max(0.1, temperature)
        self.is_calibrated = True
    
    def set_platt_parameters(self, a: float, b: float):
        """Set Platt scaling parameters."""
        self.platt_a = a
        self.platt_b = b
        self.is_calibrated = True


class TemporalFilter:
    """
    Temporal filtering for reducing false positives.
    
    Implements various temporal smoothing strategies.
    """
    
    def __init__(self, window_size: int, min_detections: int):
        self.window_size = window_size
        self.min_detections = min_detections
        self.detection_history = []
        self._lock = threading.Lock()
    
    def add_detection(self, result: DetectionResult) -> bool:
        """
        Add detection result and check if temporal condition is met.
        
        Args:
            result: Detection result
            
        Returns:
            True if temporal filtering criteria are met
        """
        with self._lock:
            # Add current detection
            self.detection_history.append(result)
            
            # Keep only recent detections
            if len(self.detection_history) > self.window_size:
                self.detection_history.pop(0)
            
            # Count positive detections in window
            positive_detections = sum(
                1 for r in self.detection_history 
                if r.is_wakeword_detected and r.wakeword_type == result.wakeword_type
            )
            
            return positive_detections >= self.min_detections
    
    def clear(self):
        """Clear detection history."""
        with self._lock:
            self.detection_history.clear()
    
    def get_history(self) -> List[DetectionResult]:
        """Get detection history."""
        with self._lock:
            return self.detection_history.copy()


class WakewordDetectionEngine:
    """
    Core wakeword detection engine with optimized inference and confidence scoring.
    
    This class handles the complete detection pipeline from audio features
    to wakeword classification with confidence estimation.
    """
    
    def __init__(self, model: nn.Module, metadata: ModelMetadata,
                 feature_extractor: AudioFeatureExtractor,
                 config: DetectionConfig,
                 device: Optional[torch.device] = None):
        """
        Initialize wakeword detection engine.
        
        Args:
            model: Trained PyTorch model
            metadata: Model metadata
            feature_extractor: Audio feature extractor
            config: Detection configuration
            device: PyTorch device
        """
        self.model = model
        self.metadata = metadata
        self.feature_extractor = feature_extractor
        self.config = config
        self.device = device or torch.device("cpu")
        
        # Validate configuration
        if not config.validate():
            raise ValueError("Invalid detection configuration")
        
        # Move model to device and set to eval mode
        self.model.to(self.device)
        self.model.eval()
        
        # Optimize model if requested
        if config.enable_model_optimization:
            self._optimize_model()
        
        # Set up half precision if requested
        if config.use_half_precision and self.device.type == "cuda":
            self.model = self.model.half()
            self._use_half_precision = True
        else:
            self._use_half_precision = False
        
        # Initialize confidence calibrator
        self.confidence_calibrator = ConfidenceCalibrator()
        
        # Initialize temporal filter if enabled
        if config.use_temporal_smoothing:
            self.temporal_filter = TemporalFilter(
                config.smoothing_window_size,
                config.min_detections_in_window
            )
        else:
            self.temporal_filter = None
        
        # Class mapping
        self.class_labels = metadata.class_labels
        self.class_to_wakeword = {
            "custom": WakewordType.CUSTOM,
            "system_command": WakewordType.SYSTEM_COMMAND,
            "negative": WakewordType.NEGATIVE
        }
        
        # Performance tracking
        self._inference_stats = {
            'total_inferences': 0,
            'total_time_ms': 0.0,
            'avg_time_ms': 0.0,
            'min_time_ms': float('inf'),
            'max_time_ms': 0.0,
            'detections_count': {wt.value: 0 for wt in WakewordType}
        }
        
        self._lock = threading.Lock()
    
    def _optimize_model(self):
        """Apply model optimizations for inference."""
        try:
            # Try to use TorchScript if possible
            if hasattr(self.model, 'reparameterize'):
                # RepCNN-specific optimization
                self.model.reparameterize()
            
            # Compile model for better performance (PyTorch 2.0+)
            if hasattr(torch, 'compile'):
                self.model = torch.compile(self.model, mode="reduce-overhead")
                
        except Exception as e:
            warnings.warn(f"Model optimization failed: {e}")
    
    def detect_wakeword(self, audio_chunk: AudioChunk) -> DetectionResult:
        """
        Detect wakeword in audio chunk.
        
        Args:
            audio_chunk: Audio chunk to process
            
        Returns:
            Detection result with confidence scores
        """
        start_time = time.time()
        
        try:
            # Extract features
            features = self.feature_extractor.extract_features(
                audio_chunk.to_tensor(self.device),
                normalize_audio=True
            )
            
            # Run inference
            with torch.no_grad():
                if self._use_half_precision:
                    features = features.half()
                
                # Forward pass
                logits = self.model(features)
                
                # Apply softmax to get probabilities
                probabilities = F.softmax(logits, dim=-1)
                
                # Apply confidence calibration if available
                if self.confidence_calibrator.is_calibrated:
                    probabilities = self.confidence_calibrator.calibrate_with_platt(probabilities)
                
                # Convert to numpy for processing
                probs = probabilities.cpu().float().numpy().squeeze()
            
            # Create raw scores dictionary
            raw_scores = {}
            for i, label in enumerate(self.class_labels):
                raw_scores[label] = float(probs[i])
            
            # Determine predicted class and confidence
            predicted_class_idx = np.argmax(probs)
            predicted_label = self.class_labels[predicted_class_idx]
            confidence = float(probs[predicted_class_idx])
            
            # Map to wakeword type
            wakeword_type = self.class_to_wakeword.get(predicted_label, WakewordType.NEGATIVE)
            
            # Apply threshold-based filtering
            is_above_threshold = self._check_confidence_threshold(wakeword_type, confidence)
            if not is_above_threshold:
                wakeword_type = WakewordType.NEGATIVE
                confidence = max(raw_scores.get("negative", 0.0), 1.0 - confidence)
            
            # Calculate processing time
            processing_time_ms = (time.time() - start_time) * 1000
            
            # Create result
            result = DetectionResult(
                wakeword_type=wakeword_type,
                confidence=confidence,
                raw_scores=raw_scores,
                processing_time_ms=processing_time_ms,
                chunk_id=audio_chunk.chunk_id,
                timestamp=audio_chunk.timestamp,
                features_shape=features.shape
            )
            
            # Apply temporal filtering if enabled
            if self.temporal_filter:
                temporal_pass = self.temporal_filter.add_detection(result)
                if not temporal_pass and result.is_wakeword_detected:
                    # Override detection if temporal filter doesn't pass
                    result.wakeword_type = WakewordType.NEGATIVE
                    result.confidence = raw_scores.get("negative", 0.5)
            
            # Update statistics
            self._update_stats(result)
            
            return result
            
        except Exception as e:
            # Return negative result on error
            processing_time_ms = (time.time() - start_time) * 1000
            return DetectionResult(
                wakeword_type=WakewordType.NEGATIVE,
                confidence=0.0,
                raw_scores={"error": 1.0},
                processing_time_ms=processing_time_ms,
                chunk_id=audio_chunk.chunk_id,
                timestamp=audio_chunk.timestamp,
                features_shape=(0,)
            )
    
    def detect_batch(self, audio_chunks: List[AudioChunk]) -> List[DetectionResult]:
        """
        Batch detection for multiple audio chunks.
        
        Args:
            audio_chunks: List of audio chunks
            
        Returns:
            List of detection results
        """
        if not self.config.batch_inference or len(audio_chunks) == 1:
            # Process individually
            return [self.detect_wakeword(chunk) for chunk in audio_chunks]
        
        # Batch processing
        batch_size = min(len(audio_chunks), self.config.max_batch_size)
        results = []
        
        for i in range(0, len(audio_chunks), batch_size):
            batch = audio_chunks[i:i + batch_size]
            batch_results = self._process_batch(batch)
            results.extend(batch_results)
        
        return results
    
    def _process_batch(self, audio_chunks: List[AudioChunk]) -> List[DetectionResult]:
        """Process a batch of audio chunks."""
        start_time = time.time()
        
        try:
            # Extract features for all chunks
            batch_features = []
            for chunk in audio_chunks:
                features = self.feature_extractor.extract_features(
                    chunk.to_tensor(self.device),
                    normalize_audio=True
                )
                batch_features.append(features)
            
            # Stack into batch tensor
            batch_tensor = torch.cat(batch_features, dim=0)
            
            # Run batch inference
            with torch.no_grad():
                if self._use_half_precision:
                    batch_tensor = batch_tensor.half()
                
                batch_logits = self.model(batch_tensor)
                batch_probabilities = F.softmax(batch_logits, dim=-1)
                
                if self.confidence_calibrator.is_calibrated:
                    batch_probabilities = self.confidence_calibrator.calibrate_with_platt(batch_probabilities)
                
                batch_probs = batch_probabilities.cpu().float().numpy()
            
            # Process results for each chunk
            results = []
            total_processing_time = (time.time() - start_time) * 1000
            processing_time_per_chunk = total_processing_time / len(audio_chunks)
            
            for i, (chunk, probs) in enumerate(zip(audio_chunks, batch_probs)):
                # Create raw scores
                raw_scores = {}
                for j, label in enumerate(self.class_labels):
                    raw_scores[label] = float(probs[j])
                
                # Determine prediction
                predicted_class_idx = np.argmax(probs)
                predicted_label = self.class_labels[predicted_class_idx]
                confidence = float(probs[predicted_class_idx])
                
                wakeword_type = self.class_to_wakeword.get(predicted_label, WakewordType.NEGATIVE)
                
                # Apply threshold filtering
                if not self._check_confidence_threshold(wakeword_type, confidence):
                    wakeword_type = WakewordType.NEGATIVE
                    confidence = max(raw_scores.get("negative", 0.0), 1.0 - confidence)
                
                result = DetectionResult(
                    wakeword_type=wakeword_type,
                    confidence=confidence,
                    raw_scores=raw_scores,
                    processing_time_ms=processing_time_per_chunk,
                    chunk_id=chunk.chunk_id,
                    timestamp=chunk.timestamp,
                    features_shape=batch_features[i].shape
                )
                
                results.append(result)
                self._update_stats(result)
            
            return results
            
        except Exception as e:
            # Return negative results on error
            processing_time_per_chunk = (time.time() - start_time) * 1000 / len(audio_chunks)
            return [
                DetectionResult(
                    wakeword_type=WakewordType.NEGATIVE,
                    confidence=0.0,
                    raw_scores={"error": 1.0},
                    processing_time_ms=processing_time_per_chunk,
                    chunk_id=chunk.chunk_id,
                    timestamp=chunk.timestamp,
                    features_shape=(0,)
                )
                for chunk in audio_chunks
            ]
    
    def _check_confidence_threshold(self, wakeword_type: WakewordType, confidence: float) -> bool:
        """Check if confidence meets threshold for wakeword type."""
        if wakeword_type == WakewordType.CUSTOM:
            return confidence >= self.config.custom_threshold
        elif wakeword_type == WakewordType.SYSTEM_COMMAND:
            return confidence >= self.config.system_command_threshold
        elif wakeword_type == WakewordType.NEGATIVE:
            return confidence >= self.config.negative_threshold
        return False
    
    def _update_stats(self, result: DetectionResult):
        """Update performance statistics."""
        with self._lock:
            self._inference_stats['total_inferences'] += 1
            self._inference_stats['total_time_ms'] += result.processing_time_ms
            self._inference_stats['avg_time_ms'] = (
                self._inference_stats['total_time_ms'] / self._inference_stats['total_inferences']
            )
            self._inference_stats['min_time_ms'] = min(
                self._inference_stats['min_time_ms'], result.processing_time_ms
            )
            self._inference_stats['max_time_ms'] = max(
                self._inference_stats['max_time_ms'], result.processing_time_ms
            )
            self._inference_stats['detections_count'][result.wakeword_type.value] += 1
    
    def update_thresholds(self, custom_threshold: Optional[float] = None,
                         system_command_threshold: Optional[float] = None,
                         negative_threshold: Optional[float] = None):
        """Update confidence thresholds."""
        if custom_threshold is not None:
            self.config.custom_threshold = custom_threshold
        if system_command_threshold is not None:
            self.config.system_command_threshold = system_command_threshold
        if negative_threshold is not None:
            self.config.negative_threshold = negative_threshold
    
    def calibrate_confidence(self, validation_data: List[Tuple[torch.Tensor, int]]):
        """
        Calibrate confidence scores using validation data.
        
        Args:
            validation_data: List of (features, label) tuples
        """
        if not validation_data:
            return
        
        # Extract features and labels
        features_list = []
        labels_list = []
        
        for features, label in validation_data:
            features_list.append(features)
            labels_list.append(label)
        
        # Stack tensors
        all_features = torch.stack(features_list).to(self.device)
        all_labels = torch.tensor(labels_list).to(self.device)
        
        # Get model predictions
        with torch.no_grad():
            logits = self.model(all_features)
            probabilities = F.softmax(logits, dim=-1)
        
        # Implement temperature scaling (simplified)
        # This is a basic implementation - for production use, implement proper optimization
        temperatures = [0.5, 1.0, 1.5, 2.0, 2.5]
        best_temperature = 1.0
        best_score = float('inf')
        
        for temp in temperatures:
            scaled_probs = F.softmax(logits / temp, dim=-1)
            loss = F.cross_entropy(scaled_probs, all_labels)
            
            if loss < best_score:
                best_score = loss
                best_temperature = temp
        
        self.confidence_calibrator.set_temperature(best_temperature)
    
    def get_performance_stats(self) -> Dict:
        """Get performance statistics."""
        with self._lock:
            stats = self._inference_stats.copy()
            
        # Add additional metrics
        stats.update({
            'model_name': self.metadata.model_name,
            'model_type': self.metadata.model_type,
            'device': str(self.device),
            'use_half_precision': self._use_half_precision,
            'temporal_filtering_enabled': self.temporal_filter is not None,
            'confidence_calibrated': self.confidence_calibrator.is_calibrated
        })
        
        return stats
    
    def reset_stats(self):
        """Reset performance statistics."""
        with self._lock:
            self._inference_stats = {
                'total_inferences': 0,
                'total_time_ms': 0.0,
                'avg_time_ms': 0.0,
                'min_time_ms': float('inf'),
                'max_time_ms': 0.0,
                'detections_count': {wt.value: 0 for wt in WakewordType}
            }
        
        if self.temporal_filter:
            self.temporal_filter.clear()
    
    def warm_up(self, num_warmup_iterations: int = 5):
        """
        Warm up the model for consistent performance measurements.
        
        Args:
            num_warmup_iterations: Number of warmup iterations
        """
        # Create dummy input matching expected shape
        dummy_input = torch.randn(
            1, 
            self.metadata.input_shape[0],
            self.metadata.input_shape[1], 
            self.metadata.input_shape[2]
        ).to(self.device)
        
        if self._use_half_precision:
            dummy_input = dummy_input.half()
        
        # Run warmup iterations
        with torch.no_grad():
            for _ in range(num_warmup_iterations):
                _ = self.model(dummy_input)
                
                # Synchronize CUDA if using GPU
                if self.device.type == "cuda":
                    torch.cuda.synchronize()


def create_detection_engine(model: nn.Module, metadata: ModelMetadata,
                          feature_extractor: AudioFeatureExtractor,
                          config: Optional[DetectionConfig] = None,
                          device: Optional[torch.device] = None) -> WakewordDetectionEngine:
    """
    Factory function to create detection engine.
    
    Args:
        model: Trained PyTorch model
        metadata: Model metadata
        feature_extractor: Audio feature extractor
        config: Detection configuration (uses default if None)
        device: PyTorch device
        
    Returns:
        Detection engine instance
    """
    if config is None:
        config = DetectionConfig()
    
    return WakewordDetectionEngine(model, metadata, feature_extractor, config, device)