Trixy3/trixy_core/ml/voice_recognition/examples.py

"""
Examples and testing framework for Voice Recognition system.

This module provides comprehensive examples, testing utilities, and mock
components for developing and testing the voice recognition system.
"""

import asyncio
import logging
import time
import threading
from typing import Dict, List, Optional, Any, Callable, Tuple
from dataclasses import dataclass
import numpy as np
import torch

from .config import VoiceRecognitionConfig, create_default_config, ModelArchitectureType
from .detector import VoiceRecognitionDetector, SpeakerIdentificationResult, AudioChunk

# Configure logging
logger = logging.getLogger(__name__)


class VRMockEventHandler:
    """Mock event handler for testing voice recognition system."""
    
    def __init__(self):
        """Initialize mock event handler."""
        self.events = []
        self.event_callbacks = {}
        self._lock = threading.Lock()
    
    def trigger(self, event_name: str, event_data: Dict[str, Any]):
        """
        Trigger an event.
        
        Args:
            event_name: Name of the event
            event_data: Event data dictionary
        """
        with self._lock:
            event_entry = {
                'name': event_name,
                'data': event_data,
                'timestamp': time.time()
            }
            self.events.append(event_entry)
            
            # Call registered callbacks
            if event_name in self.event_callbacks:
                for callback in self.event_callbacks[event_name]:
                    try:
                        callback(event_name, event_data)
                    except Exception as e:
                        logger.error(f"Error in event callback for {event_name}: {e}")
        
        logger.debug(f"Event triggered: {event_name}")
    
    def register_callback(self, event_name: str, callback: Callable):
        """Register a callback for an event."""
        if event_name not in self.event_callbacks:
            self.event_callbacks[event_name] = []
        self.event_callbacks[event_name].append(callback)
    
    def get_events(self, event_name: Optional[str] = None) -> List[Dict[str, Any]]:
        """Get recorded events."""
        with self._lock:
            if event_name:
                return [e for e in self.events if e['name'] == event_name]
            return self.events.copy()
    
    def clear_events(self):
        """Clear recorded events."""
        with self._lock:
            self.events.clear()


class VRAudioSimulator:
    """Audio simulator for testing voice recognition."""
    
    def __init__(self, sample_rate: int = 16000):
        """
        Initialize audio simulator.
        
        Args:
            sample_rate: Audio sample rate
        """
        self.sample_rate = sample_rate
        self._speaker_profiles = {}
        self._background_noise_level = 0.01
    
    def add_speaker_profile(self, 
                           speaker_id: str,
                           fundamental_freq: float = 150.0,
                           formant_freqs: List[float] = None):
        """
        Add a speaker profile for audio simulation.
        
        Args:
            speaker_id: Speaker identifier
            fundamental_freq: Fundamental frequency (F0)
            formant_freqs: Formant frequencies
        """
        if formant_freqs is None:
            formant_freqs = [800, 1200, 2500]  # Default formants
        
        self._speaker_profiles[speaker_id] = {
            'f0': fundamental_freq,
            'formants': formant_freqs
        }
    
    def generate_speaker_audio(self, 
                              speaker_id: str,
                              duration_seconds: float = 3.0,
                              add_noise: bool = True) -> np.ndarray:
        """
        Generate synthetic audio for a speaker.
        
        Args:
            speaker_id: Speaker to generate audio for
            duration_seconds: Audio duration
            add_noise: Whether to add background noise
            
        Returns:
            Generated audio samples
        """
        if speaker_id not in self._speaker_profiles:
            raise ValueError(f"Speaker {speaker_id} not found in profiles")
        
        profile = self._speaker_profiles[speaker_id]
        n_samples = int(duration_seconds * self.sample_rate)
        t = np.linspace(0, duration_seconds, n_samples)
        
        # Generate fundamental frequency
        f0_signal = np.sin(2 * np.pi * profile['f0'] * t)
        
        # Add formants
        audio = f0_signal.copy()
        for formant_freq in profile['formants']:
            formant_signal = 0.3 * np.sin(2 * np.pi * formant_freq * t)
            audio += formant_signal
        
        # Add harmonics
        for harmonic in range(2, 6):
            harmonic_freq = profile['f0'] * harmonic
            harmonic_amp = 1.0 / harmonic
            audio += harmonic_amp * np.sin(2 * np.pi * harmonic_freq * t)
        
        # Normalize
        audio = audio / np.max(np.abs(audio)) * 0.8
        
        # Add envelope (speech-like amplitude modulation)
        envelope = self._generate_speech_envelope(duration_seconds, self.sample_rate)
        audio = audio * envelope
        
        # Add background noise
        if add_noise:
            noise = np.random.normal(0, self._background_noise_level, n_samples)
            audio += noise
        
        return audio.astype(np.float32)
    
    def generate_background_audio(self, duration_seconds: float = 3.0) -> np.ndarray:
        """Generate background/noise audio."""
        n_samples = int(duration_seconds * self.sample_rate)
        
        # Generate colored noise
        white_noise = np.random.normal(0, 1, n_samples)
        
        # Apply simple low-pass filter for more realistic noise
        b = np.ones(5) / 5  # Simple moving average
        colored_noise = np.convolve(white_noise, b, mode='same')
        
        # Normalize
        colored_noise = colored_noise / np.max(np.abs(colored_noise)) * 0.1
        
        return colored_noise.astype(np.float32)
    
    def _generate_speech_envelope(self, duration: float, sample_rate: int) -> np.ndarray:
        """Generate speech-like amplitude envelope."""
        n_samples = int(duration * sample_rate)
        t = np.linspace(0, duration, n_samples)
        
        # Create speech-like envelope with pauses
        envelope = np.ones(n_samples)
        
        # Add some speech pauses
        pause_starts = np.random.uniform(0.2, duration - 0.4, size=max(1, int(duration / 2)))
        for pause_start in pause_starts:
            pause_duration = np.random.uniform(0.1, 0.3)
            start_idx = int(pause_start * sample_rate)
            end_idx = int((pause_start + pause_duration) * sample_rate)
            end_idx = min(end_idx, n_samples)
            
            # Gradual fade in/out
            fade_samples = int(0.05 * sample_rate)  # 50ms fade
            if end_idx - start_idx > 2 * fade_samples:
                envelope[start_idx:start_idx + fade_samples] *= np.linspace(1, 0.1, fade_samples)
                envelope[start_idx + fade_samples:end_idx - fade_samples] *= 0.1
                envelope[end_idx - fade_samples:end_idx] *= np.linspace(0.1, 1, fade_samples)
        
        return envelope


class VoiceRecognitionTester:
    """Comprehensive testing framework for voice recognition system."""
    
    def __init__(self, config: VoiceRecognitionConfig):
        """
        Initialize voice recognition tester.
        
        Args:
            config: Voice recognition configuration
        """
        self.config = config
        self.event_handler = VRMockEventHandler()
        self.audio_simulator = VRAudioSimulator(config.audio_config.sample_rate)
        
        # Test results storage
        self.test_results = {}
        self._current_test = None
    
    def setup_test_speakers(self, speaker_configs: Dict[str, Dict[str, Any]]):
        """
        Setup test speakers for simulation.
        
        Args:
            speaker_configs: Dictionary of speaker configurations
        """
        for speaker_id, config in speaker_configs.items():
            self.audio_simulator.add_speaker_profile(
                speaker_id=speaker_id,
                fundamental_freq=config.get('f0', 150.0),
                formant_freqs=config.get('formants', [800, 1200, 2500])
            )
    
    def run_identification_test(self, 
                               test_speakers: List[str],
                               samples_per_speaker: int = 10,
                               add_noise: bool = True) -> Dict[str, Any]:
        """
        Run speaker identification accuracy test.
        
        Args:
            test_speakers: List of speaker IDs to test
            samples_per_speaker: Number of test samples per speaker
            add_noise: Whether to add noise to test samples
            
        Returns:
            Test results dictionary
        """
        logger.info(f"Running identification test with {len(test_speakers)} speakers")
        
        # Create detector (this would normally use a real model)
        detector = self._create_mock_detector()
        
        results = {
            'test_type': 'identification',
            'speakers': test_speakers,
            'samples_per_speaker': samples_per_speaker,
            'total_samples': len(test_speakers) * samples_per_speaker,
            'correct_identifications': 0,
            'total_attempts': 0,
            'accuracy': 0.0,
            'per_speaker_accuracy': {},
            'confusion_matrix': {},
            'avg_confidence': 0.0,
            'avg_processing_time_ms': 0.0
        }
        
        processing_times = []
        confidences = []
        
        for speaker_id in test_speakers:
            correct_for_speaker = 0
            
            for sample_idx in range(samples_per_speaker):
                # Generate test audio
                audio = self.audio_simulator.generate_speaker_audio(
                    speaker_id=speaker_id,
                    duration_seconds=self.config.audio_config.audio_length,
                    add_noise=add_noise
                )
                
                # Run identification
                start_time = time.time()
                result = detector.identify_speaker(audio)
                processing_time = (time.time() - start_time) * 1000
                
                processing_times.append(processing_time)
                confidences.append(result.confidence)
                
                # Check accuracy
                results['total_attempts'] += 1
                if result.speaker_id == speaker_id:
                    results['correct_identifications'] += 1
                    correct_for_speaker += 1
                
                # Update confusion matrix
                predicted = result.speaker_id or 'unknown'
                if speaker_id not in results['confusion_matrix']:
                    results['confusion_matrix'][speaker_id] = {}
                if predicted not in results['confusion_matrix'][speaker_id]:
                    results['confusion_matrix'][speaker_id][predicted] = 0
                results['confusion_matrix'][speaker_id][predicted] += 1
            
            # Calculate per-speaker accuracy
            results['per_speaker_accuracy'][speaker_id] = correct_for_speaker / samples_per_speaker
        
        # Calculate overall metrics
        results['accuracy'] = results['correct_identifications'] / results['total_attempts']
        results['avg_confidence'] = np.mean(confidences)
        results['avg_processing_time_ms'] = np.mean(processing_times)
        
        logger.info(f"Identification test completed: {results['accuracy']:.2%} accuracy")
        
        return results
    
    def run_verification_test(self,
                             test_speakers: List[str],
                             impostor_samples: int = 20,
                             genuine_samples: int = 20) -> Dict[str, Any]:
        """
        Run speaker verification test (genuine vs impostor).
        
        Args:
            test_speakers: List of speaker IDs to test
            impostor_samples: Number of impostor samples per speaker
            genuine_samples: Number of genuine samples per speaker
            
        Returns:
            Test results dictionary
        """
        logger.info(f"Running verification test with {len(test_speakers)} speakers")
        
        detector = self._create_mock_detector()
        
        results = {
            'test_type': 'verification',
            'speakers': test_speakers,
            'genuine_samples': genuine_samples,
            'impostor_samples': impostor_samples,
            'true_positives': 0,
            'false_positives': 0,
            'true_negatives': 0,
            'false_negatives': 0,
            'far': 0.0,  # False Acceptance Rate
            'frr': 0.0,  # False Rejection Rate
            'eer': 0.0,  # Equal Error Rate
            'accuracy': 0.0
        }
        
        for speaker_id in test_speakers:
            # Test genuine samples
            for _ in range(genuine_samples):
                audio = self.audio_simulator.generate_speaker_audio(speaker_id)
                result = detector.verify_speaker(audio, speaker_id)
                
                if result.confidence >= self.config.recognition_config.verification_threshold:
                    results['true_positives'] += 1
                else:
                    results['false_negatives'] += 1
            
            # Test impostor samples
            for _ in range(impostor_samples):
                # Generate audio from a different speaker or background
                if len(test_speakers) > 1:
                    impostor_speakers = [s for s in test_speakers if s != speaker_id]
                    impostor_id = np.random.choice(impostor_speakers)
                    audio = self.audio_simulator.generate_speaker_audio(impostor_id)
                else:
                    audio = self.audio_simulator.generate_background_audio()
                
                result = detector.verify_speaker(audio, speaker_id)
                
                if result.confidence >= self.config.recognition_config.verification_threshold:
                    results['false_positives'] += 1
                else:
                    results['true_negatives'] += 1
        
        # Calculate metrics
        total_genuine = len(test_speakers) * genuine_samples
        total_impostor = len(test_speakers) * impostor_samples
        
        results['far'] = results['false_positives'] / total_impostor if total_impostor > 0 else 0
        results['frr'] = results['false_negatives'] / total_genuine if total_genuine > 0 else 0
        results['accuracy'] = (results['true_positives'] + results['true_negatives']) / (total_genuine + total_impostor)
        
        # Estimate EER (simplified)
        results['eer'] = (results['far'] + results['frr']) / 2
        
        logger.info(f"Verification test completed: FAR={results['far']:.2%}, FRR={results['frr']:.2%}")
        
        return results
    
    def run_performance_test(self, 
                           duration_minutes: float = 5.0,
                           target_fps: float = 33.3) -> Dict[str, Any]:
        """
        Run performance stress test.
        
        Args:
            duration_minutes: Test duration in minutes
            target_fps: Target processing rate (frames per second)
            
        Returns:
            Performance test results
        """
        logger.info(f"Running performance test for {duration_minutes} minutes")
        
        detector = self._create_mock_detector()
        
        results = {
            'test_type': 'performance',
            'duration_minutes': duration_minutes,
            'target_fps': target_fps,
            'frames_processed': 0,
            'frames_dropped': 0,
            'avg_processing_time_ms': 0.0,
            'max_processing_time_ms': 0.0,
            'achieved_fps': 0.0,
            'memory_usage_mb': 0.0,
            'errors': 0
        }
        
        processing_times = []
        start_time = time.time()
        end_time = start_time + (duration_minutes * 60)
        frame_interval = 1.0 / target_fps
        
        # Generate a test speaker
        test_speaker = 'performance_test_speaker'
        self.audio_simulator.add_speaker_profile(test_speaker)
        
        while time.time() < end_time:
            frame_start = time.time()
            
            try:
                # Generate test audio
                audio = self.audio_simulator.generate_speaker_audio(
                    test_speaker,
                    duration_seconds=self.config.audio_config.audio_length
                )
                
                # Process audio
                result = detector.identify_speaker(audio)
                processing_time = (time.time() - frame_start) * 1000
                processing_times.append(processing_time)
                
                results['frames_processed'] += 1
                
                # Simulate real-time constraint
                elapsed = time.time() - frame_start
                if elapsed < frame_interval:
                    time.sleep(frame_interval - elapsed)
                else:
                    results['frames_dropped'] += 1
                
            except Exception as e:
                logger.error(f"Error in performance test: {e}")
                results['errors'] += 1
        
        # Calculate results
        total_duration = time.time() - start_time
        results['achieved_fps'] = results['frames_processed'] / total_duration
        results['avg_processing_time_ms'] = np.mean(processing_times) if processing_times else 0
        results['max_processing_time_ms'] = np.max(processing_times) if processing_times else 0
        
        logger.info(f"Performance test completed: {results['achieved_fps']:.1f} FPS achieved")
        
        return results
    
    def _create_mock_detector(self) -> VoiceRecognitionDetector:
        """Create a mock detector for testing."""
        # This is a simplified mock - in real testing you'd use actual models
        # For now, we'll create a basic detector that simulates behavior
        
        class MockDetector:
            def __init__(self, config, event_handler):
                self.config = config
                self.event_handler = event_handler
                self._speakers = ['speaker_1', 'speaker_2', 'speaker_3', 'speaker_4']
            
            def identify_speaker(self, audio_data):
                # Simulate processing delay
                time.sleep(0.01)  # 10ms processing time
                
                # Mock identification based on audio characteristics
                audio_energy = np.mean(audio_data ** 2)
                speaker_idx = int(audio_energy * 1000) % len(self._speakers)
                
                return SpeakerIdentificationResult(
                    speaker_id=self._speakers[speaker_idx],
                    speaker_name=f"Speaker {speaker_idx + 1}",
                    confidence=0.7 + np.random.uniform(-0.2, 0.2),
                    processing_time_ms=10.0 + np.random.uniform(-2, 2)
                )
            
            def verify_speaker(self, audio_data, claimed_speaker_id):
                result = self.identify_speaker(audio_data)
                # Mock verification logic
                if result.speaker_id == claimed_speaker_id:
                    result.confidence = min(0.9, result.confidence + 0.1)
                else:
                    result.confidence = max(0.1, result.confidence - 0.3)
                return result
        
        return MockDetector(self.config, self.event_handler)


def example_speaker_identification():
    """Example of basic speaker identification."""
    logger.info("Running speaker identification example...")
    
    # Create configuration
    config = create_default_config(
        model_path="/mock/path/to/model.pth",
        model_password="test_password"
    )
    
    # Create mock event handler
    event_handler = VRMockEventHandler()
    
    # Create audio simulator
    audio_sim = VRAudioSimulator(config.audio_config.sample_rate)
    
    # Setup test speakers
    speakers = {
        'alice': {'f0': 180, 'formants': [900, 1300, 2800]},
        'bob': {'f0': 120, 'formants': [700, 1100, 2300]},
        'charlie': {'f0': 150, 'formants': [800, 1200, 2500]}
    }
    
    for speaker_id, params in speakers.items():
        audio_sim.add_speaker_profile(speaker_id, **params)
    
    # Generate and process audio samples
    results = []
    for speaker_id in speakers:
        audio = audio_sim.generate_speaker_audio(speaker_id)
        
        # In a real scenario, you would use VoiceRecognitionDetector here
        # For the example, we'll simulate the result
        result = {
            'speaker_id': speaker_id,
            'confidence': 0.85,
            'processing_time_ms': 25.3
        }
        results.append(result)
        
        logger.info(f"Identified {speaker_id} with confidence {result['confidence']:.2f}")
    
    logger.info("Speaker identification example completed")
    return results


def example_speaker_verification():
    """Example of speaker verification."""
    logger.info("Running speaker verification example...")
    
    # Create configuration
    config = create_default_config("/mock/path/to/model.pth")
    config.recognition_config.verification_threshold = 0.6
    
    # Create audio simulator
    audio_sim = VRAudioSimulator()
    audio_sim.add_speaker_profile('target_speaker', f0=160)
    
    # Test genuine speaker
    genuine_audio = audio_sim.generate_speaker_audio('target_speaker')
    genuine_confidence = 0.82  # Mock result
    
    # Test impostor
    impostor_audio = audio_sim.generate_background_audio()
    impostor_confidence = 0.35  # Mock result
    
    results = {
        'genuine_verification': {
            'confidence': genuine_confidence,
            'verified': genuine_confidence >= config.recognition_config.verification_threshold
        },
        'impostor_verification': {
            'confidence': impostor_confidence,
            'verified': impostor_confidence >= config.recognition_config.verification_threshold
        }
    }
    
    logger.info(f"Genuine speaker verification: {results['genuine_verification']}")
    logger.info(f"Impostor verification: {results['impostor_verification']}")
    
    return results


def example_dynamic_speaker_management():
    """Example of dynamic speaker management."""
    logger.info("Running dynamic speaker management example...")
    
    # Create configuration
    config = create_default_config("/mock/path/to/model.pth")
    
    # Create audio simulator
    audio_sim = VRAudioSimulator()
    
    # Simulate adding new speakers
    new_speakers = ['david', 'eva', 'frank']
    training_samples = []
    
    for speaker_id in new_speakers:
        # Add profile and generate training samples
        audio_sim.add_speaker_profile(speaker_id, f0=140 + len(training_samples) * 20)
        
        speaker_samples = []
        for _ in range(5):  # 5 training samples per speaker
            audio = audio_sim.generate_speaker_audio(speaker_id)
            speaker_samples.append(audio)
        
        training_samples.append({
            'speaker_id': speaker_id,
            'samples': speaker_samples
        })
    
    # Simulate speaker addition process
    results = {
        'speakers_added': [],
        'training_samples_per_speaker': 5,
        'success_rate': 1.0  # Mock 100% success
    }
    
    for speaker_data in training_samples:
        # In real implementation, this would call detector.add_speaker()
        results['speakers_added'].append(speaker_data['speaker_id'])
        logger.info(f"Added speaker: {speaker_data['speaker_id']}")
    
    logger.info("Dynamic speaker management example completed")
    return results


def run_all_vr_examples() -> Dict[str, Any]:
    """Run all voice recognition examples."""
    logger.info("Running all voice recognition examples...")
    
    results = {
        'speaker_identification': example_speaker_identification(),
        'speaker_verification': example_speaker_verification(),
        'dynamic_speaker_management': example_dynamic_speaker_management()
    }
    
    logger.info("All voice recognition examples completed")
    return results


if __name__ == "__main__":
    # Configure logging
    logging.basicConfig(level=logging.INFO)
    
    # Run examples
    all_results = run_all_vr_examples()
    print("\nExample Results:")
    for example_name, result in all_results.items():
        print(f"{example_name}: {result}")