Trixy3/trixy_core/ml/voice_recognition/speaker_manager.py

"""
Dynamic Speaker Management system.

This module provides comprehensive speaker management functionality including
adding/removing speakers, enrollment, verification, and speaker database operations.
"""

import logging
import os
from typing import Dict, List, Optional, Tuple, Union, Any, Callable
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
import numpy as np
import torch

from .speaker_embeddings import SpeakerDatabase, SpeakerProfile, EmbeddingComparator, SpeakerEmbeddingExtractor
from .model_loader import VRModelMetadata

# Configure logging
logger = logging.getLogger(__name__)


class SpeakerManagerError(Exception):
    """Exception raised by speaker manager operations."""
    pass


@dataclass
class SpeakerInfo:
    """Information about a speaker in the system."""
    speaker_id: str
    speaker_name: str
    enrollment_status: str  # "enrolled", "pending", "rejected"
    enrollment_date: Optional[str] = None
    last_seen: Optional[str] = None
    embedding_count: int = 0
    average_quality: float = 0.0
    verification_threshold: float = 0.7
    metadata: Dict[str, Any] = None
    
    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}


@dataclass
class EnrollmentResult:
    """Result of speaker enrollment process."""
    success: bool
    speaker_id: str
    speaker_name: str
    message: str
    quality_score: float = 0.0
    embedding_count: int = 0
    verification_score: float = 0.0


@dataclass
class VerificationResult:
    """Result of speaker verification."""
    success: bool
    speaker_id: Optional[str]
    speaker_name: Optional[str]
    confidence: float
    similarity_score: float
    threshold: float
    message: str


class SpeakerManager:
    """
    Comprehensive speaker management system.
    
    Handles dynamic speaker enrollment, verification, and database management
    with support for real-time operations.
    """
    
    def __init__(self,
                 database_path: str,
                 embedding_extractor: SpeakerEmbeddingExtractor,
                 similarity_metric: str = "cosine",
                 min_enrollment_samples: int = 3,
                 max_enrollment_samples: int = 10,
                 enrollment_threshold: float = 0.8,
                 verification_threshold: float = 0.7,
                 quality_threshold: float = 0.6):
        """
        Initialize speaker manager.
        
        Args:
            database_path: Path to speaker database
            embedding_extractor: Speaker embedding extractor
            similarity_metric: Similarity metric for comparisons
            min_enrollment_samples: Minimum samples for enrollment
            max_enrollment_samples: Maximum samples for enrollment
            enrollment_threshold: Threshold for successful enrollment
            verification_threshold: Threshold for speaker verification
            quality_threshold: Minimum quality threshold for embeddings
        """
        self.database_path = database_path
        self.embedding_extractor = embedding_extractor
        self.similarity_metric = similarity_metric
        
        # Configuration
        self.min_enrollment_samples = min_enrollment_samples
        self.max_enrollment_samples = max_enrollment_samples
        self.enrollment_threshold = enrollment_threshold
        self.verification_threshold = verification_threshold
        self.quality_threshold = quality_threshold
        
        # Initialize components
        self.database = SpeakerDatabase(database_path)
        self.comparator = EmbeddingComparator(similarity_metric)
        
        # Statistics
        self.enrollment_stats = {
            'successful_enrollments': 0,
            'failed_enrollments': 0,
            'total_verifications': 0,
            'successful_verifications': 0
        }
        
        # Callbacks for events
        self.on_speaker_enrolled: Optional[Callable] = None
        self.on_speaker_removed: Optional[Callable] = None
        self.on_verification_complete: Optional[Callable] = None
        
        logger.info(f"Speaker manager initialized with {self.get_speaker_count()} speakers")
    
    def enroll_speaker(self,
                      speaker_name: str,
                      audio_features_list: List[torch.Tensor],
                      speaker_id: Optional[str] = None,
                      metadata: Optional[Dict[str, Any]] = None) -> EnrollmentResult:
        """
        Enroll a new speaker using multiple audio samples.
        
        Args:
            speaker_name: Name of the speaker
            audio_features_list: List of audio feature tensors
            speaker_id: Optional custom speaker ID
            metadata: Optional speaker metadata
            
        Returns:
            Enrollment result
        """
        try:
            # Generate speaker ID if not provided
            if speaker_id is None:
                speaker_id = self._generate_speaker_id(speaker_name)
            
            # Check if speaker already exists
            if self.database.get_speaker_by_id(speaker_id) or self.database.get_speaker_by_name(speaker_name):
                return EnrollmentResult(
                    success=False,
                    speaker_id=speaker_id,
                    speaker_name=speaker_name,
                    message=f"Speaker {speaker_name} already exists"
                )
            
            # Validate number of samples
            if len(audio_features_list) < self.min_enrollment_samples:
                return EnrollmentResult(
                    success=False,
                    speaker_id=speaker_id,
                    speaker_name=speaker_name,
                    message=f"Insufficient samples: need at least {self.min_enrollment_samples}, got {len(audio_features_list)}"
                )
            
            # Extract embeddings from all samples
            embeddings = []
            quality_scores = []
            
            for features in audio_features_list[:self.max_enrollment_samples]:
                embedding = self.embedding_extractor.extract_embedding(features)
                quality = self.embedding_extractor.compute_embedding_quality(features, embedding)
                
                # Only keep embeddings above quality threshold
                if quality >= self.quality_threshold:
                    embeddings.append(embedding)
                    quality_scores.append(quality)
            
            if len(embeddings) < self.min_enrollment_samples:
                return EnrollmentResult(
                    success=False,
                    speaker_id=speaker_id,
                    speaker_name=speaker_name,
                    message=f"Insufficient quality samples: need {self.min_enrollment_samples}, got {len(embeddings)}"
                )
            
            # Verify enrollment quality by computing intra-speaker similarities
            intra_similarities = self._compute_intra_speaker_similarities(embeddings)
            average_similarity = np.mean(intra_similarities)
            
            if average_similarity < self.enrollment_threshold:
                return EnrollmentResult(
                    success=False,
                    speaker_id=speaker_id,
                    speaker_name=speaker_name,
                    message=f"Low enrollment quality: {average_similarity:.3f} < {self.enrollment_threshold}",
                    quality_score=float(np.mean(quality_scores)),
                    verification_score=average_similarity
                )
            
            # Create speaker profile
            profile = self.database.add_speaker(speaker_id, speaker_name, metadata)
            
            # Add embeddings to profile
            enrollment_date = datetime.now().isoformat()
            for embedding, quality in zip(embeddings, quality_scores):
                profile.add_embedding(embedding, enrollment_date, quality)
            
            # Save database
            self.database.save_database()
            
            # Update statistics
            self.enrollment_stats['successful_enrollments'] += 1
            
            # Trigger callback
            if self.on_speaker_enrolled:
                self.on_speaker_enrolled(speaker_id, speaker_name)
            
            logger.info(f"Successfully enrolled speaker: {speaker_name} ({speaker_id})")
            
            return EnrollmentResult(
                success=True,
                speaker_id=speaker_id,
                speaker_name=speaker_name,
                message="Speaker enrolled successfully",
                quality_score=float(np.mean(quality_scores)),
                embedding_count=len(embeddings),
                verification_score=average_similarity
            )
            
        except Exception as e:
            self.enrollment_stats['failed_enrollments'] += 1
            logger.error(f"Enrollment failed for {speaker_name}: {e}")
            
            return EnrollmentResult(
                success=False,
                speaker_id=speaker_id or "unknown",
                speaker_name=speaker_name,
                message=f"Enrollment failed: {e}"
            )
    
    def verify_speaker(self, 
                      audio_features: torch.Tensor,
                      expected_speaker_id: Optional[str] = None,
                      threshold: Optional[float] = None) -> VerificationResult:
        """
        Verify a speaker from audio features.
        
        Args:
            audio_features: Audio features tensor
            expected_speaker_id: Expected speaker ID for verification
            threshold: Custom verification threshold
            
        Returns:
            Verification result
        """
        try:
            self.enrollment_stats['total_verifications'] += 1
            
            # Extract embedding
            query_embedding = self.embedding_extractor.extract_embedding(audio_features)
            
            # Get all speaker embeddings
            speaker_embeddings = self.database.get_all_speaker_embeddings("weighted_mean")
            
            if not speaker_embeddings:
                return VerificationResult(
                    success=False,
                    speaker_id=None,
                    speaker_name=None,
                    confidence=0.0,
                    similarity_score=0.0,
                    threshold=threshold or self.verification_threshold,
                    message="No enrolled speakers found"
                )
            
            # Find best match
            best_speaker_id, best_similarity = self.comparator.find_best_match(
                query_embedding, speaker_embeddings
            )
            
            # Get speaker profile
            speaker_profile = self.database.get_speaker_by_id(best_speaker_id)
            speaker_name = speaker_profile.speaker_name if speaker_profile else "Unknown"
            
            # Use custom threshold or default
            used_threshold = threshold or self.verification_threshold
            
            # Check if verification passes threshold
            verification_success = best_similarity >= used_threshold
            
            # If expected speaker is specified, check if it matches
            if expected_speaker_id and best_speaker_id != expected_speaker_id:
                verification_success = False
                message = f"Expected {expected_speaker_id}, got {best_speaker_id}"
            else:
                if verification_success:
                    message = f"Speaker verified: {speaker_name}"
                    self.enrollment_stats['successful_verifications'] += 1
                else:
                    message = f"Verification failed: similarity {best_similarity:.3f} < {used_threshold}"
            
            # Update last seen for verified speaker
            if verification_success and speaker_profile:
                speaker_profile.metadata['last_seen'] = datetime.now().isoformat()
                self.database.save_database()
            
            # Trigger callback
            if self.on_verification_complete:
                self.on_verification_complete(verification_success, best_speaker_id, best_similarity)
            
            return VerificationResult(
                success=verification_success,
                speaker_id=best_speaker_id if verification_success else None,
                speaker_name=speaker_name if verification_success else None,
                confidence=best_similarity,
                similarity_score=best_similarity,
                threshold=used_threshold,
                message=message
            )
            
        except Exception as e:
            logger.error(f"Verification failed: {e}")
            
            return VerificationResult(
                success=False,
                speaker_id=None,
                speaker_name=None,
                confidence=0.0,
                similarity_score=0.0,
                threshold=threshold or self.verification_threshold,
                message=f"Verification error: {e}"
            )
    
    def identify_speaker(self, audio_features: torch.Tensor) -> VerificationResult:
        """
        Identify speaker without knowing expected identity.
        
        Args:
            audio_features: Audio features tensor
            
        Returns:
            Identification result
        """
        return self.verify_speaker(audio_features, expected_speaker_id=None)
    
    def remove_speaker(self, speaker_id: str) -> bool:
        """
        Remove speaker from database.
        
        Args:
            speaker_id: Speaker ID to remove
            
        Returns:
            True if speaker was removed
        """
        speaker_profile = self.database.get_speaker_by_id(speaker_id)
        if not speaker_profile:
            return False
        
        speaker_name = speaker_profile.speaker_name
        success = self.database.remove_speaker(speaker_id)
        
        if success:
            self.database.save_database()
            
            # Trigger callback
            if self.on_speaker_removed:
                self.on_speaker_removed(speaker_id, speaker_name)
            
            logger.info(f"Removed speaker: {speaker_name} ({speaker_id})")
        
        return success
    
    def add_speaker_sample(self,
                          speaker_id: str,
                          audio_features: torch.Tensor) -> bool:
        """
        Add additional sample to existing speaker.
        
        Args:
            speaker_id: Speaker ID
            audio_features: Audio features tensor
            
        Returns:
            True if sample was added successfully
        """
        speaker_profile = self.database.get_speaker_by_id(speaker_id)
        if not speaker_profile:
            return False
        
        # Extract embedding
        embedding = self.embedding_extractor.extract_embedding(audio_features)
        quality = self.embedding_extractor.compute_embedding_quality(audio_features, embedding)
        
        # Only add if quality is sufficient
        if quality >= self.quality_threshold:
            enrollment_date = datetime.now().isoformat()
            speaker_profile.add_embedding(embedding, enrollment_date, quality)
            self.database.save_database()
            
            logger.info(f"Added sample to speaker {speaker_id} (quality: {quality:.3f})")
            return True
        else:
            logger.warning(f"Sample quality too low for speaker {speaker_id}: {quality:.3f}")
            return False
    
    def get_speaker_info(self, speaker_id: str) -> Optional[SpeakerInfo]:
        """Get detailed information about a speaker."""
        profile = self.database.get_speaker_by_id(speaker_id)
        if not profile:
            return None
        
        return SpeakerInfo(
            speaker_id=speaker_id,
            speaker_name=profile.speaker_name,
            enrollment_status="enrolled",
            enrollment_date=profile.enrollment_dates[0] if profile.enrollment_dates else None,
            last_seen=profile.metadata.get('last_seen'),
            embedding_count=len(profile.embeddings),
            average_quality=float(np.mean(profile.quality_scores)) if profile.quality_scores else 0.0,
            verification_threshold=self.verification_threshold,
            metadata=profile.metadata
        )
    
    def list_speakers(self) -> List[SpeakerInfo]:
        """Get list of all speakers with their information."""
        speakers = []
        for speaker_id, _ in self.database.get_speaker_list():
            info = self.get_speaker_info(speaker_id)
            if info:
                speakers.append(info)
        return speakers
    
    def get_speaker_count(self) -> int:
        """Get total number of enrolled speakers."""
        return self.database.get_speaker_count()
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get system statistics."""
        stats = self.enrollment_stats.copy()
        stats.update({
            'total_speakers': self.get_speaker_count(),
            'total_embeddings': self.database.get_embedding_count(),
            'verification_success_rate': (
                stats['successful_verifications'] / max(1, stats['total_verifications'])
            ),
            'enrollment_success_rate': (
                stats['successful_enrollments'] / max(1, stats['successful_enrollments'] + stats['failed_enrollments'])
            )
        })
        return stats
    
    def update_model(self, new_embedding_extractor: SpeakerEmbeddingExtractor) -> None:
        """
        Update the embedding extractor model.
        
        Args:
            new_embedding_extractor: New embedding extractor
        """
        self.embedding_extractor = new_embedding_extractor
        logger.info("Updated embedding extractor model")
    
    def validate_speaker_database(self) -> Dict[str, Any]:
        """
        Validate speaker database integrity.
        
        Returns:
            Validation report
        """
        report = {
            'is_valid': True,
            'issues': [],
            'warnings': [],
            'speaker_stats': {}
        }
        
        for speaker_id, profile in self.database.speakers.items():
            speaker_issues = []
            
            # Check embedding consistency
            if not profile.embeddings:
                speaker_issues.append("No embeddings")
            elif len(profile.embeddings) < self.min_enrollment_samples:
                speaker_issues.append(f"Too few embeddings: {len(profile.embeddings)}")
            
            # Check quality scores
            if profile.quality_scores:
                avg_quality = np.mean(profile.quality_scores)
                if avg_quality < self.quality_threshold:
                    speaker_issues.append(f"Low average quality: {avg_quality:.3f}")
            
            # Check intra-speaker consistency
            if len(profile.embeddings) > 1:
                similarities = self._compute_intra_speaker_similarities(profile.embeddings)
                avg_similarity = np.mean(similarities)
                if avg_similarity < self.enrollment_threshold:
                    speaker_issues.append(f"Low intra-speaker similarity: {avg_similarity:.3f}")
            
            if speaker_issues:
                report['issues'].append(f"Speaker {speaker_id}: {', '.join(speaker_issues)}")
                report['is_valid'] = False
            
            report['speaker_stats'][speaker_id] = {
                'embedding_count': len(profile.embeddings),
                'average_quality': float(np.mean(profile.quality_scores)) if profile.quality_scores else 0.0,
                'issues': speaker_issues
            }
        
        return report
    
    def _generate_speaker_id(self, speaker_name: str) -> str:
        """Generate unique speaker ID from name."""
        # Simple ID generation - in production, use more robust method
        base_id = "spk_" + speaker_name.lower().replace(" ", "_")
        
        # Ensure uniqueness
        counter = 1
        speaker_id = base_id
        while self.database.get_speaker_by_id(speaker_id):
            speaker_id = f"{base_id}_{counter:03d}"
            counter += 1
        
        return speaker_id
    
    def _compute_intra_speaker_similarities(self, embeddings: List[np.ndarray]) -> List[float]:
        """Compute similarities between all pairs of embeddings."""
        similarities = []
        
        for i in range(len(embeddings)):
            for j in range(i + 1, len(embeddings)):
                sim = self.comparator.compute_similarity(embeddings[i], embeddings[j])
                similarities.append(sim)
        
        return similarities


def create_speaker_manager(database_path: str,
                          embedding_extractor: SpeakerEmbeddingExtractor,
                          **kwargs) -> SpeakerManager:
    """
    Factory function to create speaker manager.
    
    Args:
        database_path: Path to speaker database
        embedding_extractor: Speaker embedding extractor
        **kwargs: Additional configuration parameters
        
    Returns:
        Configured speaker manager
    """
    return SpeakerManager(database_path, embedding_extractor, **kwargs)


def add_speaker_to_database(manager: SpeakerManager,
                           speaker_name: str,
                           audio_features_list: List[torch.Tensor],
                           **kwargs) -> EnrollmentResult:
    """
    Convenience function to add speaker to database.
    
    Args:
        manager: Speaker manager instance
        speaker_name: Name of speaker to add
        audio_features_list: List of audio features for enrollment
        **kwargs: Additional enrollment parameters
        
    Returns:
        Enrollment result
    """
    return manager.enroll_speaker(speaker_name, audio_features_list, **kwargs)


def remove_speaker_from_database(manager: SpeakerManager, speaker_id: str) -> bool:
    """
    Convenience function to remove speaker from database.
    
    Args:
        manager: Speaker manager instance
        speaker_id: Speaker ID to remove
        
    Returns:
        True if speaker was removed
    """
    return manager.remove_speaker(speaker_id)


# Example usage and testing
def test_speaker_manager():
    """Test speaker manager functionality."""
    from .model_architectures import ECADATDNNModel
    from .speaker_embeddings import SpeakerEmbeddingExtractor
    
    # Create mock model and extractor
    model = ECADATDNNModel(input_dim=40, embedding_dim=192, num_speakers=10)
    extractor = SpeakerEmbeddingExtractor(model, "cpu")
    
    # Create speaker manager
    database_path = "/tmp/test_speaker_db.pkl"
    manager = create_speaker_manager(database_path, extractor)
    
    # Test enrollment
    print("Testing speaker enrollment...")
    
    # Generate mock audio features
    audio_features_list = [torch.randn(40, 100) for _ in range(5)]
    
    # Enroll speaker
    result = manager.enroll_speaker("Test Speaker", audio_features_list)
    print(f"Enrollment result: {result.success}, {result.message}")
    
    if result.success:
        # Test verification
        print("Testing speaker verification...")
        test_features = torch.randn(40, 100)
        
        verification = manager.verify_speaker(test_features, result.speaker_id)
        print(f"Verification result: {verification.success}, confidence: {verification.confidence:.3f}")
        
        # Test identification
        identification = manager.identify_speaker(test_features)
        print(f"Identification result: {identification.speaker_name}, confidence: {identification.confidence:.3f}")
        
        # Get speaker info
        info = manager.get_speaker_info(result.speaker_id)
        print(f"Speaker info: {info.speaker_name}, embeddings: {info.embedding_count}")
        
        # Get statistics
        stats = manager.get_statistics()
        print(f"Statistics: {stats}")
    
    # Clean up
    if os.path.exists(database_path):
        os.remove(database_path)


if __name__ == "__main__":
    test_speaker_manager()