app/speaker.py

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from app.audio import rms


@dataclass
class SpeakerFocusDecision:
    should_process: bool
    enrolled: bool = False
    updated: bool = False
    similarity: float | None = None
    mixed_speaker: bool = False
    reason: str = "accept"
    profile_updates: int = 0


@dataclass
class SpeakerProfile:
    embedding: np.ndarray | None = None
    updates: int = 0


def _cosine_similarity(left: np.ndarray, right: np.ndarray) -> float:
    left_norm = float(np.linalg.norm(left))
    right_norm = float(np.linalg.norm(right))
    if left_norm <= 1e-8 or right_norm <= 1e-8:
        return 0.0
    return float(np.dot(left, right) / (left_norm * right_norm))


def _frame_audio(audio: np.ndarray, frame_samples: int, hop_samples: int) -> np.ndarray:
    if audio.size < frame_samples or frame_samples <= 0 or hop_samples <= 0:
        return np.zeros((0, max(frame_samples, 1)), dtype=np.float32)
    frames = [
        audio[start : start + frame_samples]
        for start in range(0, audio.size - frame_samples + 1, hop_samples)
    ]
    if not frames:
        return np.zeros((0, frame_samples), dtype=np.float32)
    return np.stack(frames).astype(np.float32)


def _rolloff_frequency(power: np.ndarray, freqs: np.ndarray, percentile: float) -> float:
    if power.size == 0 or freqs.size == 0:
        return 0.0
    total = float(np.sum(power))
    if total <= 1e-8:
        return 0.0
    cumulative = np.cumsum(power)
    target = total * percentile
    index = int(np.searchsorted(cumulative, target, side="left"))
    index = min(max(index, 0), freqs.size - 1)
    return float(freqs[index])


def _voiced_frame_mask(frames: np.ndarray, min_rms: float) -> np.ndarray:
    if frames.size == 0:
        return np.zeros(0, dtype=bool)
    frame_rms = np.sqrt(np.mean(np.square(frames), axis=1, dtype=np.float32))
    if frame_rms.size == 0:
        return np.zeros(0, dtype=bool)
    dynamic_floor = max(min_rms, float(np.percentile(frame_rms, 40)) * 1.15)
    return frame_rms >= dynamic_floor


def _frame_embedding(frames: np.ndarray, sample_rate: int) -> np.ndarray | None:
    if frames.size == 0:
        return None

    window = np.hanning(frames.shape[1]).astype(np.float32)
    windowed = frames * window
    spectrum = np.abs(np.fft.rfft(windowed, axis=1)).astype(np.float32)
    power = np.square(spectrum, dtype=np.float32)
    freqs = np.fft.rfftfreq(frames.shape[1], d=1.0 / sample_rate).astype(np.float32)
    total_power = np.sum(power, axis=1) + 1e-8

    centroid = np.sum(power * freqs[None, :], axis=1) / total_power
    spread = np.sqrt(np.sum(power * np.square(freqs[None, :] - centroid[:, None]), axis=1) / total_power)
    zcr = np.mean(np.abs(np.diff(np.signbit(frames), axis=1)), axis=1).astype(np.float32)
    log_energy = np.log(np.mean(np.square(frames), axis=1, dtype=np.float32) + 1e-8)

    low_band = np.logical_and(freqs >= 120.0, freqs < 700.0)
    mid_band = np.logical_and(freqs >= 700.0, freqs < 1800.0)
    high_band = np.logical_and(freqs >= 1800.0, freqs < 4200.0)

    low_ratio = np.sum(power[:, low_band], axis=1) / total_power
    mid_ratio = np.sum(power[:, mid_band], axis=1) / total_power
    high_ratio = np.sum(power[:, high_band], axis=1) / total_power

    rolloff_85 = np.array([_rolloff_frequency(row, freqs, 0.85) for row in power], dtype=np.float32)
    rolloff_95 = np.array([_rolloff_frequency(row, freqs, 0.95) for row in power], dtype=np.float32)

    features = np.stack(
        [
            centroid / max(sample_rate / 2.0, 1.0),
            spread / max(sample_rate / 2.0, 1.0),
            zcr,
            log_energy,
            low_ratio,
            mid_ratio,
            high_ratio,
            rolloff_85 / max(sample_rate / 2.0, 1.0),
            rolloff_95 / max(sample_rate / 2.0, 1.0),
        ],
        axis=1,
    ).astype(np.float32)

    feature_mean = np.mean(features, axis=0, dtype=np.float32)
    feature_std = np.std(features, axis=0, dtype=np.float32)
    embedding = np.concatenate([feature_mean, feature_std]).astype(np.float32)
    norm = float(np.linalg.norm(embedding))
    if norm <= 1e-8:
        return None
    return embedding / norm


def build_speaker_embedding(
    audio: np.ndarray,
    sample_rate: int,
    *,
    min_rms: float,
    frame_ms: int = 25,
    hop_ms: int = 10,
) -> tuple[np.ndarray | None, int]:
    if audio.size == 0 or sample_rate <= 0:
        return None, 0

    frame_samples = max(1, int(sample_rate * (frame_ms / 1000.0)))
    hop_samples = max(1, int(sample_rate * (hop_ms / 1000.0)))
    frames = _frame_audio(audio, frame_samples, hop_samples)
    if frames.shape[0] == 0:
        return None, 0

    voiced_mask = _voiced_frame_mask(frames, min_rms)
    voiced_frames = frames[voiced_mask]
    if voiced_frames.shape[0] == 0:
        return None, 0
    return _frame_embedding(voiced_frames, sample_rate), int(voiced_frames.shape[0])


def detect_mixed_speakers(
    audio: np.ndarray,
    sample_rate: int,
    *,
    min_rms: float,
    divergence_threshold: float,
) -> bool:
    if audio.size == 0 or sample_rate <= 0:
        return False

    segment_count = 3
    min_segment_samples = max(1, int(sample_rate * 0.6))
    if audio.size < min_segment_samples * segment_count:
        return False

    segment_embeddings: list[np.ndarray] = []
    boundaries = np.linspace(0, audio.size, num=segment_count + 1, dtype=int)
    for start, end in zip(boundaries[:-1], boundaries[1:]):
        segment = audio[start:end]
        embedding, voiced_frames = build_speaker_embedding(segment, sample_rate, min_rms=min_rms)
        if embedding is None or voiced_frames < 6:
            continue
        segment_embeddings.append(embedding)

    if len(segment_embeddings) < 2:
        return False

    max_divergence = 0.0
    for index, left in enumerate(segment_embeddings):
        for right in segment_embeddings[index + 1 :]:
            divergence = 1.0 - _cosine_similarity(left, right)
            max_divergence = max(max_divergence, divergence)
    return max_divergence >= divergence_threshold


def evaluate_speaker_focus(
    audio: np.ndarray,
    sample_rate: int,
    *,
    profile: SpeakerProfile,
    enabled: bool,
    min_utterance_ms: int,
    min_rms: float,
    similarity_threshold: float,
    profile_alpha: float,
    multi_speaker_threshold: float,
    reject_mixed: bool,
) -> SpeakerFocusDecision:
    if not enabled or audio.size == 0 or sample_rate <= 0:
        return SpeakerFocusDecision(should_process=True, reason="disabled", profile_updates=profile.updates)

    utterance_ms = (audio.size / sample_rate) * 1000.0
    utterance_rms = rms(audio)
    if utterance_ms < max(min_utterance_ms, 0) or utterance_rms < min_rms:
        return SpeakerFocusDecision(should_process=True, reason="insufficient_audio", profile_updates=profile.updates)

    embedding, voiced_frames = build_speaker_embedding(audio, sample_rate, min_rms=min_rms)
    if embedding is None or voiced_frames < 8:
        return SpeakerFocusDecision(should_process=True, reason="insufficient_features", profile_updates=profile.updates)

    mixed_speaker = detect_mixed_speakers(
        audio,
        sample_rate,
        min_rms=min_rms,
        divergence_threshold=multi_speaker_threshold,
    )
    if mixed_speaker and reject_mixed:
        return SpeakerFocusDecision(
            should_process=False,
            mixed_speaker=True,
            reason="mixed_speakers",
            profile_updates=profile.updates,
        )

    if profile.embedding is None:
        profile.embedding = embedding
        profile.updates = 1
        return SpeakerFocusDecision(
            should_process=True,
            enrolled=True,
            mixed_speaker=mixed_speaker,
            reason="enrolled",
            profile_updates=profile.updates,
        )

    similarity = _cosine_similarity(embedding, profile.embedding)
    if similarity < similarity_threshold:
        return SpeakerFocusDecision(
            should_process=False,
            similarity=similarity,
            mixed_speaker=mixed_speaker,
            reason="off_speaker",
            profile_updates=profile.updates,
        )

    alpha = min(max(profile_alpha, 0.0), 1.0)
    updated_embedding = ((1.0 - alpha) * profile.embedding) + (alpha * embedding)
    norm = float(np.linalg.norm(updated_embedding))
    if norm > 1e-8:
        profile.embedding = (updated_embedding / norm).astype(np.float32)
    profile.updates += 1
    return SpeakerFocusDecision(
        should_process=True,
        updated=True,
        similarity=similarity,
        mixed_speaker=mixed_speaker,
        reason="accept",
        profile_updates=profile.updates,
    )