Trixy3/trixy_core/ml/voice_recognition/recognition_engine.py

"""
Voice Recognition Engine with real-time processing and confidence scoring.

This module provides the core voice recognition engine that integrates all
components for real-time speaker identification and verification.
"""

import logging
import time
import threading
from typing import Dict, List, Optional, Tuple, Union, Any, Callable
from dataclasses import dataclass, field
from collections import deque
from datetime import datetime
import numpy as np
import torch

from .config import VoiceRecognitionConfig
from .audio_features import VRAudioFeatureExtractor
from .vad import VoiceActivityDetector, VADResult
from .speaker_manager import SpeakerManager, VerificationResult
from .speaker_embeddings import SpeakerEmbeddingExtractor, EmbeddingComparator

# Configure logging
logger = logging.getLogger(__name__)


@dataclass
class RecognitionResult:
    """Result of voice recognition process."""
    # Speaker identification
    speaker_id: Optional[str]
    speaker_name: Optional[str]
    confidence: float
    similarity_score: float
    
    # Processing information
    is_speech: bool
    vad_confidence: float
    processing_time_ms: float
    audio_quality: float
    
    # Temporal information
    start_time: float
    end_time: float
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
    
    # Additional details
    threshold_used: float = 0.0
    alternative_speakers: List[Tuple[str, float]] = field(default_factory=list)
    debug_info: Dict[str, Any] = field(default_factory=dict)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert result to dictionary."""
        return {
            'speaker_id': self.speaker_id,
            'speaker_name': self.speaker_name,
            'confidence': self.confidence,
            'similarity_score': self.similarity_score,
            'is_speech': self.is_speech,
            'vad_confidence': self.vad_confidence,
            'processing_time_ms': self.processing_time_ms,
            'audio_quality': self.audio_quality,
            'start_time': self.start_time,
            'end_time': self.end_time,
            'timestamp': self.timestamp,
            'threshold_used': self.threshold_used,
            'alternative_speakers': self.alternative_speakers,
            'debug_info': self.debug_info
        }


@dataclass 
class IdentificationResult(RecognitionResult):
    """Result specifically for speaker identification (who is speaking)."""
    pass


@dataclass
class VerificationResult(RecognitionResult):
    """Result specifically for speaker verification (is this person X)."""
    expected_speaker_id: Optional[str] = None
    verification_success: bool = False


class ConfidenceCalibrator:
    """Calibrate confidence scores for better reliability."""
    
    def __init__(self, 
                 temperature: float = 1.0,
                 use_platt_scaling: bool = True):
        """
        Initialize confidence calibrator.
        
        Args:
            temperature: Temperature scaling parameter
            use_platt_scaling: Whether to use Platt scaling
        """
        self.temperature = temperature
        self.use_platt_scaling = use_platt_scaling
        
        # Platt scaling parameters (will be learned)
        self.platt_a = 1.0
        self.platt_b = 0.0
        
        # Statistics for calibration
        self.score_history = deque(maxlen=1000)
        self.label_history = deque(maxlen=1000)
    
    def calibrate_confidence(self, raw_score: float, additional_info: Dict[str, Any] = None) -> float:
        """
        Calibrate raw similarity score to confidence.
        
        Args:
            raw_score: Raw similarity score
            additional_info: Additional information for calibration
            
        Returns:
            Calibrated confidence score
        """
        # Temperature scaling
        scaled_score = raw_score / self.temperature
        
        # Platt scaling if enabled
        if self.use_platt_scaling:
            # Sigmoid function: 1 / (1 + exp(A * score + B))
            calibrated = 1.0 / (1.0 + np.exp(self.platt_a * scaled_score + self.platt_b))
        else:
            # Simple clipping to [0, 1]
            calibrated = np.clip(scaled_score, 0.0, 1.0)
        
        # Additional adjustments based on context
        if additional_info:
            # Adjust based on VAD confidence
            vad_confidence = additional_info.get('vad_confidence', 1.0)
            calibrated *= vad_confidence
            
            # Adjust based on audio quality
            audio_quality = additional_info.get('audio_quality', 1.0)
            calibrated *= (0.5 + 0.5 * audio_quality)  # Don't penalize too much
            
            # Adjust based on temporal consistency
            temporal_consistency = additional_info.get('temporal_consistency', 1.0)
            calibrated *= (0.8 + 0.2 * temporal_consistency)
        
        return float(np.clip(calibrated, 0.0, 1.0))
    
    def update_calibration(self, score: float, is_correct: bool):
        """Update calibration parameters based on feedback."""
        self.score_history.append(score)
        self.label_history.append(1.0 if is_correct else 0.0)
        
        # Simple adaptation of temperature
        if len(self.score_history) >= 50:
            scores = np.array(list(self.score_history))
            labels = np.array(list(self.label_history))
            
            # Compute accuracy at different temperature values
            temperatures = np.linspace(0.5, 2.0, 20)
            best_temp = 1.0
            best_accuracy = 0.0
            
            for temp in temperatures:
                calibrated_scores = scores / temp
                predictions = calibrated_scores > 0.5
                accuracy = np.mean(predictions == labels)
                
                if accuracy > best_accuracy:
                    best_accuracy = accuracy
                    best_temp = temp
            
            # Update temperature with momentum
            self.temperature = 0.9 * self.temperature + 0.1 * best_temp


class VoiceRecognitionEngine:
    """
    Core voice recognition engine with real-time processing capabilities.
    
    Integrates VAD, feature extraction, speaker recognition, and confidence scoring
    for robust real-time speaker identification and verification.
    """
    
    def __init__(self,
                 config: VoiceRecognitionConfig,
                 feature_extractor: VRAudioFeatureExtractor,
                 vad_detector: VoiceActivityDetector,
                 speaker_manager: SpeakerManager,
                 embedding_extractor: SpeakerEmbeddingExtractor):
        """
        Initialize voice recognition engine.
        
        Args:
            config: Voice recognition configuration
            feature_extractor: Audio feature extractor
            vad_detector: Voice activity detector
            speaker_manager: Speaker management system
            embedding_extractor: Speaker embedding extractor
        """
        self.config = config
        self.feature_extractor = feature_extractor
        self.vad_detector = vad_detector
        self.speaker_manager = speaker_manager
        self.embedding_extractor = embedding_extractor
        
        # Initialize confidence calibrator
        self.calibrator = ConfidenceCalibrator()
        
        # Temporal smoothing
        self.temporal_window_size = config.recognition_config.smoothing_window_size
        self.confidence_decay = config.recognition_config.confidence_decay_factor
        self.recent_results = deque(maxlen=self.temporal_window_size)
        
        # Performance monitoring
        self.processing_times = deque(maxlen=100)
        self.quality_scores = deque(maxlen=100)
        
        # Thread safety
        self._lock = threading.Lock()
        
        # Statistics
        self.stats = {
            'total_processed': 0,
            'speech_detected': 0,
            'speakers_identified': 0,
            'average_processing_time_ms': 0.0,
            'average_confidence': 0.0
        }
        
        logger.info("Voice recognition engine initialized")
    
    def process_audio(self, audio_data: np.ndarray, sample_rate: int = None) -> RecognitionResult:
        """
        Process audio data for speaker recognition.
        
        Args:
            audio_data: Raw audio data
            sample_rate: Audio sample rate (uses config default if None)
            
        Returns:
            Recognition result
        """
        start_time = time.time()
        
        try:
            with self._lock:
                # Convert to torch tensor
                if isinstance(audio_data, np.ndarray):
                    audio_tensor = torch.from_numpy(audio_data).float()
                else:
                    audio_tensor = audio_data.float()
                
                # Ensure correct sample rate
                if sample_rate is None:
                    sample_rate = self.config.audio_config.sample_rate
                
                # Step 1: Voice Activity Detection
                vad_result = self._perform_vad(audio_tensor, sample_rate)
                
                if not vad_result.is_speech:
                    # No speech detected
                    processing_time = (time.time() - start_time) * 1000
                    return self._create_no_speech_result(vad_result, processing_time)
                
                # Step 2: Feature Extraction
                features = self._extract_features(audio_tensor)
                
                # Step 3: Speaker Recognition
                recognition_result = self._recognize_speaker(features, vad_result)
                
                # Step 4: Temporal Smoothing
                smoothed_result = self._apply_temporal_smoothing(recognition_result)
                
                # Step 5: Final Processing
                final_result = self._finalize_result(smoothed_result, start_time, len(audio_data) / sample_rate)
                
                # Update statistics
                self._update_statistics(final_result)
                
                return final_result
                
        except Exception as e:
            logger.error(f"Error processing audio: {e}")
            processing_time = (time.time() - start_time) * 1000
            
            return RecognitionResult(
                speaker_id=None,
                speaker_name=None,
                confidence=0.0,
                similarity_score=0.0,
                is_speech=False,
                vad_confidence=0.0,
                processing_time_ms=processing_time,
                audio_quality=0.0,
                start_time=0.0,
                end_time=len(audio_data) / (sample_rate or 16000),
                debug_info={'error': str(e)}
            )
    
    def identify_speaker(self, audio_data: np.ndarray, sample_rate: int = None) -> IdentificationResult:
        """
        Identify who is speaking in the audio.
        
        Args:
            audio_data: Raw audio data
            sample_rate: Audio sample rate
            
        Returns:
            Identification result
        """
        base_result = self.process_audio(audio_data, sample_rate)
        
        return IdentificationResult(
            speaker_id=base_result.speaker_id,
            speaker_name=base_result.speaker_name,
            confidence=base_result.confidence,
            similarity_score=base_result.similarity_score,
            is_speech=base_result.is_speech,
            vad_confidence=base_result.vad_confidence,
            processing_time_ms=base_result.processing_time_ms,
            audio_quality=base_result.audio_quality,
            start_time=base_result.start_time,
            end_time=base_result.end_time,
            timestamp=base_result.timestamp,
            threshold_used=base_result.threshold_used,
            alternative_speakers=base_result.alternative_speakers,
            debug_info=base_result.debug_info
        )
    
    def verify_speaker(self, 
                      audio_data: np.ndarray, 
                      expected_speaker_id: str,
                      sample_rate: int = None) -> VerificationResult:
        """
        Verify if the speaker in audio matches expected speaker.
        
        Args:
            audio_data: Raw audio data
            expected_speaker_id: Expected speaker ID
            sample_rate: Audio sample rate
            
        Returns:
            Verification result
        """
        base_result = self.process_audio(audio_data, sample_rate)
        
        # Check if identified speaker matches expected
        verification_success = (
            base_result.speaker_id == expected_speaker_id and
            base_result.confidence >= self.config.recognition_config.verification_threshold
        )
        
        return VerificationResult(
            speaker_id=base_result.speaker_id,
            speaker_name=base_result.speaker_name,
            confidence=base_result.confidence,
            similarity_score=base_result.similarity_score,
            is_speech=base_result.is_speech,
            vad_confidence=base_result.vad_confidence,
            processing_time_ms=base_result.processing_time_ms,
            audio_quality=base_result.audio_quality,
            start_time=base_result.start_time,
            end_time=base_result.end_time,
            timestamp=base_result.timestamp,
            threshold_used=base_result.threshold_used,
            alternative_speakers=base_result.alternative_speakers,
            debug_info=base_result.debug_info,
            expected_speaker_id=expected_speaker_id,
            verification_success=verification_success
        )
    
    def _perform_vad(self, audio_tensor: torch.Tensor, sample_rate: int) -> VADResult:
        """Perform voice activity detection."""
        # Convert to numpy for VAD
        audio_np = audio_tensor.numpy()
        
        # Ensure correct sample rate for VAD
        if sample_rate != self.config.audio_config.sample_rate:
            # Simple resampling (in production, use proper resampling)
            target_length = int(len(audio_np) * self.config.audio_config.sample_rate / sample_rate)
            audio_np = np.interp(
                np.linspace(0, len(audio_np), target_length),
                np.arange(len(audio_np)),
                audio_np
            )
        
        return self.vad_detector.detect(audio_np)
    
    def _extract_features(self, audio_tensor: torch.Tensor) -> torch.Tensor:
        """Extract audio features."""
        return self.feature_extractor(audio_tensor)
    
    def _recognize_speaker(self, features: torch.Tensor, vad_result: VADResult) -> RecognitionResult:
        """Perform speaker recognition on features."""
        # Extract speaker embedding
        embedding = self.embedding_extractor.extract_embedding(features)
        
        # Compute audio quality
        audio_quality = self._compute_audio_quality(features, embedding)
        
        # Get all speaker embeddings for comparison
        speaker_embeddings = self.speaker_manager.database.get_all_speaker_embeddings("weighted_mean")
        
        if not speaker_embeddings:
            return RecognitionResult(
                speaker_id=None,
                speaker_name=None,
                confidence=0.0,
                similarity_score=0.0,
                is_speech=vad_result.is_speech,
                vad_confidence=vad_result.confidence,
                processing_time_ms=0.0,
                audio_quality=audio_quality,
                start_time=vad_result.start_time,
                end_time=vad_result.end_time
            )
        
        # Compute similarities to all speakers
        similarities = {}
        comparator = self.speaker_manager.comparator
        
        for speaker_id, speaker_embedding in speaker_embeddings.items():
            similarity = comparator.compute_similarity(embedding, speaker_embedding)
            similarities[speaker_id] = similarity
        
        # Find best match
        best_speaker_id = max(similarities, key=similarities.get)
        best_similarity = similarities[best_speaker_id]
        
        # Get speaker name
        speaker_profile = self.speaker_manager.database.get_speaker_by_id(best_speaker_id)
        speaker_name = speaker_profile.speaker_name if speaker_profile else "Unknown"
        
        # Calibrate confidence
        calibration_info = {
            'vad_confidence': vad_result.confidence,
            'audio_quality': audio_quality,
            'temporal_consistency': self._compute_temporal_consistency(best_speaker_id)
        }
        
        confidence = self.calibrator.calibrate_confidence(best_similarity, calibration_info)
        
        # Get alternative speakers (top 3)
        sorted_similarities = sorted(similarities.items(), key=lambda x: x[1], reverse=True)
        alternative_speakers = [(spk_id, sim) for spk_id, sim in sorted_similarities[1:4]]
        
        return RecognitionResult(
            speaker_id=best_speaker_id,
            speaker_name=speaker_name,
            confidence=confidence,
            similarity_score=best_similarity,
            is_speech=vad_result.is_speech,
            vad_confidence=vad_result.confidence,
            processing_time_ms=0.0,  # Will be set later
            audio_quality=audio_quality,
            start_time=vad_result.start_time,
            end_time=vad_result.end_time,
            threshold_used=self.config.recognition_config.identification_threshold,
            alternative_speakers=alternative_speakers,
            debug_info={
                'all_similarities': similarities,
                'embedding_norm': float(np.linalg.norm(embedding))
            }
        )
    
    def _compute_audio_quality(self, features: torch.Tensor, embedding: np.ndarray) -> float:
        """Compute overall audio quality score."""
        quality_factors = []
        
        # Feature-based quality
        feature_energy = torch.mean(features ** 2).item()
        energy_quality = min(1.0, feature_energy * 5)  # Scale appropriately
        quality_factors.append(energy_quality)
        
        # Feature variance (too low might indicate poor audio)
        feature_var = torch.var(features).item()
        var_quality = min(1.0, feature_var * 10)
        quality_factors.append(var_quality)
        
        # Embedding quality (from extractor)
        embedding_quality = self.embedding_extractor.compute_embedding_quality(features, embedding)
        quality_factors.append(embedding_quality)
        
        # Signal-to-noise ratio estimate
        signal_power = torch.mean(features ** 2, dim=-1)
        noise_estimate = torch.min(signal_power)
        snr_estimate = torch.mean(signal_power) / (noise_estimate + 1e-8)
        snr_quality = min(1.0, snr_estimate.item() / 10)
        quality_factors.append(snr_quality)
        
        return float(np.mean(quality_factors))
    
    def _compute_temporal_consistency(self, speaker_id: str) -> float:
        """Compute temporal consistency with recent predictions."""
        if not self.recent_results:
            return 1.0
        
        # Count recent predictions for this speaker
        recent_count = sum(1 for result in self.recent_results if result.speaker_id == speaker_id)
        consistency = recent_count / len(self.recent_results)
        
        return consistency
    
    def _apply_temporal_smoothing(self, current_result: RecognitionResult) -> RecognitionResult:
        """Apply temporal smoothing to recognition results."""
        if not self.config.recognition_config.use_temporal_smoothing:
            return current_result
        
        # Add current result to history
        self.recent_results.append(current_result)
        
        if len(self.recent_results) < 2:
            return current_result
        
        # Compute weighted average of recent confidences for same speaker
        same_speaker_results = [r for r in self.recent_results if r.speaker_id == current_result.speaker_id]
        
        if same_speaker_results:
            # Apply exponential decay to weights
            weights = [self.confidence_decay ** i for i in range(len(same_speaker_results))]
            weights.reverse()  # Most recent gets highest weight
            
            weighted_confidence = sum(w * r.confidence for w, r in zip(weights, same_speaker_results))
            total_weight = sum(weights)
            
            smoothed_confidence = weighted_confidence / total_weight
            
            # Update current result
            current_result.confidence = smoothed_confidence
        
        return current_result
    
    def _finalize_result(self, result: RecognitionResult, start_time: float, audio_duration: float) -> RecognitionResult:
        """Finalize recognition result with timing and quality checks."""
        # Set processing time
        processing_time = (time.time() - start_time) * 1000
        result.processing_time_ms = processing_time
        
        # Apply final confidence thresholding
        if result.confidence < self.config.recognition_config.identification_threshold:
            result.speaker_id = None
            result.speaker_name = None
            result.confidence = 0.0
        
        # Set audio timing
        if result.start_time == 0.0 and result.end_time == 0.0:
            result.start_time = 0.0
            result.end_time = audio_duration
        
        return result
    
    def _create_no_speech_result(self, vad_result: VADResult, processing_time: float) -> RecognitionResult:
        """Create result for when no speech is detected."""
        return RecognitionResult(
            speaker_id=None,
            speaker_name=None,
            confidence=0.0,
            similarity_score=0.0,
            is_speech=False,
            vad_confidence=vad_result.confidence,
            processing_time_ms=processing_time,
            audio_quality=0.0,
            start_time=vad_result.start_time,
            end_time=vad_result.end_time,
            debug_info={'vad_reason': 'No speech detected'}
        )
    
    def _update_statistics(self, result: RecognitionResult):
        """Update engine statistics."""
        self.stats['total_processed'] += 1
        
        if result.is_speech:
            self.stats['speech_detected'] += 1
            
            if result.speaker_id:
                self.stats['speakers_identified'] += 1
        
        # Update rolling averages
        self.processing_times.append(result.processing_time_ms)
        if result.is_speech:
            self.quality_scores.append(result.audio_quality)
        
        self.stats['average_processing_time_ms'] = np.mean(list(self.processing_times))
        
        if self.quality_scores:
            self.stats['average_confidence'] = np.mean([r.confidence for r in self.recent_results if r.is_speech])
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get engine statistics."""
        stats = self.stats.copy()
        
        if stats['speech_detected'] > 0:
            stats['identification_rate'] = stats['speakers_identified'] / stats['speech_detected']
        else:
            stats['identification_rate'] = 0.0
        
        if stats['total_processed'] > 0:
            stats['speech_detection_rate'] = stats['speech_detected'] / stats['total_processed']
        else:
            stats['speech_detection_rate'] = 0.0
        
        return stats
    
    def reset_statistics(self):
        """Reset engine statistics."""
        self.stats = {
            'total_processed': 0,
            'speech_detected': 0,
            'speakers_identified': 0,
            'average_processing_time_ms': 0.0,
            'average_confidence': 0.0
        }
        self.processing_times.clear()
        self.quality_scores.clear()
        self.recent_results.clear()


def create_recognition_engine(config: VoiceRecognitionConfig,
                             feature_extractor: VRAudioFeatureExtractor,
                             vad_detector: VoiceActivityDetector,
                             speaker_manager: SpeakerManager,
                             embedding_extractor: SpeakerEmbeddingExtractor) -> VoiceRecognitionEngine:
    """
    Factory function to create voice recognition engine.
    
    Args:
        config: Voice recognition configuration
        feature_extractor: Audio feature extractor
        vad_detector: Voice activity detector
        speaker_manager: Speaker management system
        embedding_extractor: Speaker embedding extractor
        
    Returns:
        Configured voice recognition engine
    """
    return VoiceRecognitionEngine(
        config, feature_extractor, vad_detector, speaker_manager, embedding_extractor
    )


# Example usage and testing
def test_recognition_engine():
    """Test voice recognition engine functionality."""
    from .config import create_default_config
    from .audio_features import create_feature_extractor
    from .vad import create_vad_detector, VADConfig
    from .model_architectures import ECADATDNNModel
    from .speaker_embeddings import SpeakerEmbeddingExtractor
    from .speaker_manager import create_speaker_manager
    
    # Create components
    config = create_default_config("/tmp/test_model.pth")
    
    feature_extractor = create_feature_extractor("log_mel", sample_rate=16000)
    
    vad_config = VADConfig()
    vad_detector = create_vad_detector("energy", vad_config)
    
    model = ECADATDNNModel(input_dim=40, embedding_dim=192, num_speakers=4)
    embedding_extractor = SpeakerEmbeddingExtractor(model, "cpu")
    
    speaker_manager = create_speaker_manager("/tmp/test_speakers.pkl", embedding_extractor)
    
    # Create recognition engine
    engine = create_recognition_engine(
        config, feature_extractor, vad_detector, speaker_manager, embedding_extractor
    )
    
    # Test with mock audio
    print("Testing voice recognition engine...")
    
    # Generate test audio (sine wave + noise)
    sample_rate = 16000
    duration = 3.0
    t = np.linspace(0, duration, int(sample_rate * duration))
    audio = np.sin(2 * np.pi * 440 * t) + 0.1 * np.random.randn(len(t))
    
    # Test identification
    result = engine.identify_speaker(audio, sample_rate)
    print(f"Identification result: {result.speaker_name}, confidence: {result.confidence:.3f}")
    print(f"Processing time: {result.processing_time_ms:.1f}ms")
    print(f"Audio quality: {result.audio_quality:.3f}")
    
    # Test verification (will fail since no speakers enrolled)
    verification = engine.verify_speaker(audio, "test_speaker", sample_rate)
    print(f"Verification result: {verification.verification_success}, confidence: {verification.confidence:.3f}")
    
    # Get statistics
    stats = engine.get_statistics()
    print(f"Engine statistics: {stats}")


if __name__ == "__main__":
    test_recognition_engine()