pipeline_modules/talker_selector.py

"""Talker-of-interest (ToI) selection heuristics."""

import logging
from typing import List, Tuple
from .data_structures import TalkerInfo

log = logging.getLogger(__name__)


def select_talker_of_interest(
    talkers: List[TalkerInfo],
    weights: dict = None,
) -> Tuple[int, TalkerInfo]:
    """
    Select the most likely talker of interest using weighted scoring.
    
    Heuristics (in order of importance if using defaults):
    1. Front-facing direction (within 45° of 0°) - hearing aid users typically face target
    2. English language - contextual heuristic
    3. Energy/loudness - louder speakers are more salient
    4. Gender clarity - non-ambiguous gender increases confidence
    
    Args:
        talkers: List of TalkerInfo objects
        weights: Optional dict with keys: 'front_bonus', 'english_bonus', 
                'energy_weight', 'gender_bonus'
                
    Returns:
        Tuple of (selected_index, selected_talker_info)
    """
    if not talkers:
        raise ValueError("No talkers to select from")
    
    # Default weights
    default_weights = {
        'front_bonus': 10.0,      # strongest factor
        'english_bonus': 5.0,     # language bonus
        'energy_weight': 5.0,     # energy multiplier
        'gender_bonus': 0.5,      # small bonus for clear gender
    }
    
    if weights:
        default_weights.update(weights)
    
    w = default_weights
    best_idx = 0
    best_score = -float('inf')
    
    scores = []
    
    for i, talker in enumerate(talkers):
        score = 0.0
        score_components = {}
        
        # Front-facing bonus
        if talker.direction_deg is not None:
            angle = talker.direction_deg
            angle_from_front = min(angle, 360 - angle)
            if angle_from_front < 45:  # Within front hemisphere
                score += w['front_bonus']
                score_components['front'] = w['front_bonus']
        
        # English language bonus
        if talker.language == 'en':
            score += w['english_bonus']
            score_components['english'] = w['english_bonus']
        
        # Energy bonus
        if talker.energy is not None:
            score += talker.energy * w['energy_weight']
            score_components['energy'] = talker.energy * w['energy_weight']
        
        # Gender clarity bonus
        if talker.gender and talker.gender != 'ambiguous' and talker.gender != 'unknown':
            score += w['gender_bonus']
            score_components['gender'] = w['gender_bonus']
        
        # Store score for later reference
        talker.selection_score = round(score, 2)
        scores.append((i, score, score_components))
        
        log.debug(f"Speaker {i+1}: score={score:.2f}, components={score_components}")
        
        if score > best_score:
            best_score = score
            best_idx = i
    
    selected_talker = talkers[best_idx]
    selected_talker.is_toi = True
    
    # Build reasoning string
    reasons = []
    best_components = scores[best_idx][2]
    if 'front' in best_components:
        reasons.append("Front-facing")
    if 'english' in best_components:
        reasons.append("English speaker")
    if 'energy' in best_components and best_components['energy'] > 0:
        reasons.append(f"High energy ({selected_talker.energy:.3f})")
    if 'gender' in best_components:
        reasons.append(f"Clear {selected_talker.gender}")
    
    selected_talker.toi_reason = "; ".join(reasons) if reasons else "Default selection"
    
    log.info(f"Selected talker: Speaker {best_idx + 1} (score: {best_score:.2f})")
    log.info(f"  Reason: {selected_talker.toi_reason}")
    
    return best_idx, selected_talker


def select_by_energy(talkers: List[TalkerInfo]) -> Tuple[int, TalkerInfo]:
    """Simple heuristic: select talker with highest energy."""
    best_idx = 0
    best_energy = -float('inf')
    
    for i, talker in enumerate(talkers):
        if talker.energy is not None and talker.energy > best_energy:
            best_energy = talker.energy
            best_idx = i
    
    talkers[best_idx].is_toi = True
    talkers[best_idx].toi_reason = f"Highest energy ({best_energy:.3f})"
    return best_idx, talkers[best_idx]


def select_by_language(talkers: List[TalkerInfo], target_language: str = 'en') -> Tuple[int, TalkerInfo]:
    """Select first talker with target language, fallback to speaker 0."""
    for i, talker in enumerate(talkers):
        if talker.language == target_language:
            talkers[i].is_toi = True
            talkers[i].toi_reason = f"{target_language.upper()} speaker"
            return i, talker
    
    # Fallback to first speaker
    talkers[0].is_toi = True
    talkers[0].toi_reason = f"No {target_language.upper()} speaker found, defaulting to speaker 1"
    return 0, talkers[0]