"""Feature extractor for the ECAPA-TDNN speaker encoder.

Converts raw audio waveforms into log-mel spectrograms suitable for the
ECAPA-TDNN speaker encoder model.
"""

import numpy as np
import torch
from transformers.feature_extraction_utils import BatchFeature, FeatureExtractionMixin


class EcapaTdnnFeatureExtractor(FeatureExtractionMixin):
    r"""
    Feature extractor for ECAPA-TDNN speaker encoder models.

    Converts raw audio waveforms to 128-bin log-mel spectrograms matching the
    Qwen3-TTS preprocessing pipeline.

    Args:
        sample_rate (`int`, defaults to 24000):
            Target sample rate in Hz. Audio will be resampled if needed.
        n_fft (`int`, defaults to 1024):
            FFT window size.
        hop_length (`int`, defaults to 256):
            Hop length between STFT frames.
        n_mels (`int`, defaults to 128):
            Number of mel-frequency bins.
        fmin (`float`, defaults to 0):
            Minimum frequency for mel filterbank.
        fmax (`float`, defaults to 12000):
            Maximum frequency for mel filterbank.
    """

    model_input_names = ["input_values"]

    def __init__(
        self,
        sample_rate=24000,
        n_fft=1024,
        hop_length=256,
        n_mels=128,
        fmin=0,
        fmax=12000,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.sample_rate = sample_rate
        self.sampling_rate = sample_rate  # alias for HF pipeline compatibility
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.n_mels = n_mels
        self.fmin = fmin
        self.fmax = fmax

    def __call__(self, raw_speech, sampling_rate=None, return_tensors="pt", **kwargs):
        """
        Process raw audio waveform(s) into log-mel spectrogram features.

        Args:
            raw_speech (`np.ndarray`, `list[np.ndarray]`, or file path `str`):
                Raw audio waveform(s) as float32 numpy array(s), or a file path.
            sampling_rate (`int`, *optional*):
                Sample rate of the input audio. Resampled to ``self.sample_rate``
                if different.
            return_tensors (`str`, defaults to ``"pt"``):
                Return type — ``"pt"`` for PyTorch tensors.

        Returns:
            ``BatchFeature`` with ``input_values`` key containing the log-mel
            spectrogram tensor of shape ``(batch, time, n_mels)``.
        """
        # Handle single input
        if isinstance(raw_speech, str):
            import librosa
            raw_speech, sampling_rate = librosa.load(raw_speech, sr=None, mono=True)

        if isinstance(raw_speech, np.ndarray) and raw_speech.ndim == 1:
            raw_speech = [raw_speech]

        features = []
        for audio in raw_speech:
            if isinstance(audio, str):
                import librosa
                audio, sampling_rate = librosa.load(audio, sr=None, mono=True)

            mel = self._compute_mel(audio, sampling_rate or self.sample_rate)
            features.append(mel)

        # Pad to same length
        max_len = max(f.shape[1] for f in features)
        padded = []
        for f in features:
            if f.shape[1] < max_len:
                f = torch.nn.functional.pad(f, (0, 0, 0, max_len - f.shape[1]))
            padded.append(f)

        input_values = torch.cat(padded, dim=0)
        return BatchFeature({"input_values": input_values})

    def _compute_mel(self, audio, sr):
        """Compute 128-bin log-mel spectrogram matching Qwen3-TTS requirements."""
        import librosa
        from librosa.filters import mel as librosa_mel_fn

        if isinstance(audio, torch.Tensor):
            audio = audio.numpy()

        if sr != self.sample_rate:
            audio = librosa.resample(
                audio.astype(np.float32), orig_sr=sr, target_sr=self.sample_rate
            )

        y = torch.from_numpy(audio).unsqueeze(0).float()
        mel_basis = torch.from_numpy(
            librosa_mel_fn(
                sr=self.sample_rate,
                n_fft=self.n_fft,
                n_mels=self.n_mels,
                fmin=self.fmin,
                fmax=self.fmax,
            )
        ).float()

        padding = (self.n_fft - self.hop_length) // 2
        y = torch.nn.functional.pad(
            y.unsqueeze(1), (padding, padding), mode="reflect"
        ).squeeze(1)
        hann = torch.hann_window(self.n_fft)
        spec = torch.stft(
            y,
            self.n_fft,
            hop_length=self.hop_length,
            win_length=self.n_fft,
            window=hann,
            center=False,
            return_complex=True,
        )
        spec = torch.abs(spec)
        mel = torch.matmul(mel_basis, spec)
        mel = torch.log(torch.clamp(mel, min=1e-5))
        return mel.transpose(1, 2)  # (1, time, n_mels)