"""
Speaker embedding extraction with efficient two-phase processing.

This module implements a two-phase API for extracting speaker embeddings:
1. preprocess(): Computes filterbank (fbank) features once for the full audio
2. extract(): Extracts embeddings for multiple overlapping segments using cached fbank

This design avoids redundant fbank computation for overlapping segments, which would
otherwise increase processing time by 9-10x due to the sliding window approach used
in speaker diarization pipelines.

Example usage:
    model = WeSpeakerEmbeddingModel(session, fbank_extractor)

    # Phase 1: Preprocess full audio once
    fbank_cache = model.preprocess(waveform)

    # Phase 2: Extract embeddings for multiple segments
    segments = [
        EmbeddingSegment(start=0.0, end=10.0, mask=mask1),
        EmbeddingSegment(start=1.0, end=11.0, mask=mask2),
        ...
    ]
    embeddings = model.extract(fbank_cache, segments, batch_size=128)
"""

import logging
import time
from typing import TypedDict

import numpy as np
import numpy.typing as npt
import onnxruntime as ort

from onnx_diarization.fbank import FbankExtractor

logger = logging.getLogger(__name__)

# Default batch size for embedding extraction inference
# Balances memory usage and throughput for typical use cases
DEFAULT_EMBEDDING_BATCH_SIZE = 128


class FbankData(TypedDict):
    """
    Cached filterbank features and metadata from preprocessing phase.

    This cache allows efficient extraction of embeddings from multiple overlapping
    segments without recomputing fbank features for each segment.

    Attributes:
        fbank: Filterbank features for the full audio with shape (num_frames, feature_dim)
        sample_rate: Audio sample rate in Hz (typically 16000)
        frame_step_seconds: Time step between consecutive fbank frames in seconds
    """

    fbank: npt.NDArray[np.float32]
    sample_rate: int
    frame_step_seconds: float


class EmbeddingSegment(TypedDict):
    """
    A time segment with speaker activity mask for embedding extraction.

    Represents a specific time window in the audio along with a mask indicating
    which frames contain the target speaker. Used to extract speaker-specific
    embeddings from longer audio segments.

    Attributes:
        start: Start time of segment in seconds
        end: End time of segment in seconds
        mask: Speaker activity mask with shape (num_frames,) where values indicate
              speaker presence (0.0 = absent, 1.0 = present, intermediate values allowed)
    """

    start: float
    end: float
    mask: npt.NDArray[np.float32]


def pad_features(
    features_list: list[npt.NDArray[np.float32]], max_len: int, feature_dim: int
) -> npt.NDArray[np.float32]:
    """Pad feature arrays to uniform length for batched inference.

    Args:
        features_list: List of feature arrays with varying lengths
        max_len: Maximum length to pad all arrays to
        feature_dim: Feature dimension for padding

    Returns:
        Padded array with shape (batch_size, max_len, feature_dim)
    """
    padded = np.zeros((len(features_list), max_len, feature_dim), dtype=np.float32)
    for i, feats in enumerate(features_list):
        padded[i, : len(feats), :] = feats
    return padded


class WeSpeakerEmbeddingModel:
    """
    WeSpeaker ONNX embedding model with two-phase processing for efficiency.

    This model extracts 256-dimensional speaker embeddings using a WeSpeaker model
    via ONNX Runtime. It implements a two-phase workflow to optimize performance
    when processing multiple overlapping segments:

    Phase 1 (preprocess): Compute filterbank features once for the full audio
    Phase 2 (extract): Extract embeddings for multiple segments using cached features

    The two-phase design is critical for performance in diarization pipelines that
    use sliding windows with significant overlap (typically 90% overlap). Without
    caching, fbank features would be recomputed 9-10x for overlapping regions.

    Attributes:
        session: ONNX Runtime inference session for the embedding model
        dimension: Output embedding dimension (fixed at 256)
        metric: Distance metric for comparing embeddings (cosine distance)
        fbank: Filterbank feature extractor

    Example:
        # Initialize model
        model = WeSpeakerEmbeddingModel(session, FbankExtractor(sample_rate=16000))

        # Phase 1: Preprocess audio
        fbank_data = model.preprocess(waveform)

        # Phase 2: Extract embeddings for segments with speaker masks
        segments = [
            EmbeddingSegment(start=0.0, end=10.0, mask=speaker1_mask),
            EmbeddingSegment(start=1.0, end=11.0, mask=speaker2_mask),
        ]
        embeddings = model.extract(fbank_data, segments, batch_size=128)
    """

    def __init__(self, session: ort.InferenceSession, fbank_extractor: FbankExtractor) -> None:
        self.session = session
        self.dimension = 256
        self.metric = "cosine"
        self.fbank = fbank_extractor

    def preprocess(self, waveform: npt.NDArray[np.float32]) -> FbankData:
        """
        Phase 1: Compute and cache filterbank features for the full audio.

        This method computes filterbank (fbank) features once for the entire audio
        waveform. The resulting FbankData can then be reused to extract embeddings
        for multiple overlapping segments efficiently in Phase 2 (extract).

        This preprocessing step is essential for performance: in typical diarization
        pipelines with 90% overlap between segments, computing fbank on-demand would
        result in 9-10x redundant computation.

        Args:
            waveform: Raw audio waveform with shape (num_samples,), typically at 16kHz

        Returns:
            FbankData containing cached filterbank features and metadata for use in extract()
        """
        t_start = time.perf_counter()
        fbank = self.fbank.compute_fbank(waveform)
        audio_duration = len(waveform) / self.fbank.sample_rate
        frame_step_seconds = self.fbank.fbank_opts.frame_opts.frame_shift_ms / 1000.0

        logger.info(f"Computed fbank for audio ({audio_duration:.2f}s) in {time.perf_counter() - t_start:.3f}s")

        return FbankData(fbank=fbank, sample_rate=self.fbank.sample_rate, frame_step_seconds=frame_step_seconds)

    def extract(
        self,
        fbank_data: FbankData,
        segments: list[EmbeddingSegment] | None = None,
        batch_size: int = DEFAULT_EMBEDDING_BATCH_SIZE,
    ) -> npt.NDArray[np.float32]:
        """
        Phase 2: Extract speaker embeddings for multiple segments using cached fbank.

        This method efficiently extracts embeddings for multiple time segments by:
        1. Extracting relevant fbank slices from the cached full-audio fbank (no recomputation)
        2. Applying speaker activity masks to isolate target speaker frames
        3. Running batched ONNX inference to generate embeddings

        The batching strategy processes segments in sub-batches, grouping by similar
        lengths to minimize padding waste and optimize memory usage.

        Args:
            fbank_data: Cached filterbank features from preprocess()
            segments: List of time segments with speaker activity masks to extract embeddings for.
                     If None, extracts a single embedding for the entire audio with uniform mask
                     (useful for single-speaker audio).
            batch_size: Maximum number of segments to process in a single inference batch.
                       Larger values increase throughput but require more memory.

        Returns:
            Array of embeddings with shape (num_segments, 256), one 256-dim embedding per segment.
            If segments is None, returns shape (1, 256) for the full audio.
        """
        t_start = time.perf_counter()

        # Handle single-speaker case: create a segment for the full audio
        if segments is None:
            num_frames = fbank_data["fbank"].shape[0]
            audio_duration = num_frames * fbank_data["frame_step_seconds"]
            mask = np.ones(num_frames, dtype=np.float32)
            segments = [EmbeddingSegment(start=0.0, end=audio_duration, mask=mask)]

        # Extract fbank features for each segment's time range from the cached full fbank
        fbank_segments = self._extract_fbank_segments(fbank_data, segments)
        masks = np.array([seg["mask"] for seg in segments], dtype=np.float32)

        # Apply speaker masks to isolate frames where target speaker is active
        masked_items = self._apply_masks(masks, fbank_segments)

        # Run batched ONNX inference on masked features
        embeddings = self._run_batched_inference(
            masked_items, len(segments), fbank_segments.shape[2], inference_batch_size=batch_size
        )

        logger.info(f"Extracted {len(segments)} embeddings in {time.perf_counter() - t_start:.3f}s")
        return embeddings

    def _extract_fbank_segments(
        self, fbank_data: FbankData, segments: list[EmbeddingSegment]
    ) -> npt.NDArray[np.float32]:
        """
        Extract fbank feature slices for each segment from the cached full-audio fbank.

        Converts segment time boundaries (in seconds) to frame indices, then extracts
        the corresponding slice from the cached fbank array. Ensures all segments have
        uniform frame counts through padding or trimming.

        This avoids recomputing fbank for overlapping segments - we simply slice
        different regions from the same cached fbank array.
        """
        fbank_full = fbank_data["fbank"]
        frame_step_seconds = fbank_data["frame_step_seconds"]

        # Calculate expected frame count based on segment duration
        # Assumes all segments have the same duration (typical in diarization)
        expected_num_frames = int((segments[0]["end"] - segments[0]["start"]) / frame_step_seconds)

        fbank_segments = []
        for seg in segments:
            # Convert time boundaries to frame indices
            start_frame = int(seg["start"] / frame_step_seconds)
            end_frame = min(int(seg["end"] / frame_step_seconds), fbank_full.shape[0])

            # Extract the fbank slice for this segment's time range
            segment_fbank = fbank_full[start_frame:end_frame]

            # Ensure uniform length (handles edge cases at audio boundaries)
            segment_fbank = self._pad_or_trim(segment_fbank, expected_num_frames)
            fbank_segments.append(segment_fbank)

        return np.array(fbank_segments)

    def _pad_or_trim(self, array: npt.NDArray[np.float32], target_length: int) -> npt.NDArray[np.float32]:
        if array.shape[0] < target_length:
            padding = target_length - array.shape[0]
            return np.pad(array, ((0, padding), (0, 0)), mode="constant")
        return array[:target_length]

    def _apply_masks(
        self, masks: npt.NDArray[np.float32], fbank_features: npt.NDArray[np.float32]
    ) -> list[npt.NDArray[np.float32] | None]:
        """Apply speaker masks to filterbank features to isolate speaker-specific frames.

        For each item in the batch, upsamples the mask to match fbank frame count,
        binarizes it, and filters fbank features to keep only frames where the
        target speaker is active.

        The mask upsampling is necessary because segmentation masks typically have
        lower temporal resolution than fbank features (e.g., 589 mask frames vs
        ~1000 fbank frames for a 10-second segment).

        Args:
            masks: Speaker activity masks with shape (batch_size, num_frames)
            fbank_features: Filterbank features with shape (batch_size, num_frames, feature_dim)

        Returns:
            List of masked feature arrays (or None for items with empty masks)
        """
        masked_items = []

        for i in range(len(masks)):
            feats = fbank_features[i]
            mask = masks[i]

            # Skip segments where speaker is never active
            if mask.sum() == 0:
                masked_items.append(None)
                continue

            # Upsample mask to match fbank frame count using linear interpolation
            # This aligns the coarse segmentation mask with the higher-resolution fbank
            mask_upsampled = np.interp(
                np.arange(len(feats)),
                np.linspace(0, len(feats) - 1, len(mask)),
                mask,
            )

            # Binarize: keep frames where upsampled mask value > 0.5
            mask_binary = mask_upsampled > 0.5

            if mask_binary.sum() == 0:
                masked_items.append(None)
                continue

            # Extract only the frames where target speaker is active
            # This filters out frames with other speakers or silence
            feats_masked = feats[mask_binary]
            masked_items.append(feats_masked)

        return masked_items

    def _run_batched_inference(
        self,
        masked_items: list[npt.NDArray[np.float32] | None],
        batch_size: int,
        feature_dim: int,
        inference_batch_size: int,
    ) -> npt.NDArray[np.float32]:
        """Run ONNX inference in sub-batches for memory efficiency.

        Groups items by similar length to minimize padding waste, then processes
        each sub-batch through the ONNX model.

        The batching strategy:
        1. Filter out None items (inactive speakers)
        2. Sort by length to group similar lengths together
        3. Process in sub-batches with minimal padding
        4. Map results back to original indices

        Args:
            masked_items: List of masked feature arrays (or None for empty items)
            batch_size: Original batch size for output array
            feature_dim: Feature dimension for padding
            inference_batch_size: Maximum batch size for inference

        Returns:
            Embeddings array with shape (batch_size, embedding_dim)
        """
        # Initialize output array with zeros
        embeddings_output = np.zeros((batch_size, self.dimension), dtype=np.float32)

        # Filter out None items (segments with no speaker activity)
        valid_indices = [i for i, item in enumerate(masked_items) if item is not None]
        if not valid_indices:
            return embeddings_output

        # Sort by length to group similar lengths together
        # This minimizes padding waste: shorter items are batched together,
        # longer items are batched together, reducing unnecessary padding
        sorted_indices = sorted(valid_indices, key=lambda i: len(masked_items[i]))

        # Process in sub-batches for memory efficiency
        for batch_start in range(0, len(sorted_indices), inference_batch_size):
            batch_indices = sorted_indices[batch_start : batch_start + inference_batch_size]

            # Pad all items in this sub-batch to the max length within the sub-batch
            max_len = max(len(masked_items[idx]) for idx in batch_indices)
            padded_feats = pad_features([masked_items[idx] for idx in batch_indices], max_len, feature_dim)

            # Run ONNX inference on this sub-batch
            embeddings_batch = self.session.run(None, {"feats": padded_feats})[0]

            # Map embeddings back to original indices in the output array
            for k, original_idx in enumerate(batch_indices):
                embeddings_output[original_idx] = embeddings_batch[k]

        return embeddings_output
