import math
from typing import Dict, List, Optional

import mlx.core as mx
import mlx.nn as nn


def hann_window_periodic(size: int) -> mx.array:
    """
    Create periodic Hann window (fftbins=True), matching scipy.signal.get_window('hann', size, fftbins=True).

    Uses size in denominator instead of size-1, which is the correct window for STFT.
    """
    return mx.array([0.5 * (1 - math.cos(2 * math.pi * n / size)) for n in range(size)])


def get_padding(kernel_size: int, dilation: int = 1) -> int:
    return int((kernel_size * dilation - dilation) / 2)


class Snake(nn.Module):

    def __init__(
        self,
        in_features: int,
        alpha: float = 1.0,
        alpha_trainable: bool = True,
        alpha_logscale: bool = False,
    ):
        super().__init__()
        self.in_features = in_features
        self.alpha_logscale = alpha_logscale

        if alpha_logscale:
            self.alpha = mx.zeros(in_features) * alpha
        else:
            self.alpha = mx.ones(in_features) * alpha

        self.no_div_by_zero = 1e-9

    def __call__(self, x: mx.array) -> mx.array:

        alpha = mx.reshape(self.alpha, (1, -1, 1))

        if self.alpha_logscale:
            alpha = mx.exp(alpha)

        no_div_by_zero = 1e-9
        min_alpha = 1e-4

        alpha_sign = mx.sign(alpha)
        alpha_abs = mx.abs(alpha)

        alpha_clamped = alpha_sign * mx.maximum(alpha_abs, min_alpha)

        alpha_clamped = mx.where(alpha_abs < no_div_by_zero, min_alpha, alpha_clamped)

        return x + (1.0 / alpha_clamped) * mx.power(mx.sin(x * alpha), 2)


class ResBlock(nn.Module):

    def __init__(
        self,
        channels: int = 512,
        kernel_size: int = 3,
        dilations: List[int] = [1, 3, 5],
    ):
        super().__init__()
        self.channels = channels
        self.convs1 = []
        self.convs2 = []
        self.activations1 = []
        self.activations2 = []

        for dilation in dilations:
            # First conv in each dilation block
            self.convs1.append(
                nn.Conv1d(
                    channels,
                    channels,
                    kernel_size,
                    stride=1,
                    padding=get_padding(kernel_size, dilation),
                    dilation=dilation,
                )
            )
            # Second conv (dilation=1)
            self.convs2.append(
                nn.Conv1d(
                    channels,
                    channels,
                    kernel_size,
                    stride=1,
                    padding=get_padding(kernel_size, 1),
                )
            )
            # Snake activations
            self.activations1.append(Snake(channels, alpha_logscale=False))
            self.activations2.append(Snake(channels, alpha_logscale=False))

    def __call__(self, x: mx.array) -> mx.array:
        for i in range(len(self.convs1)):
            xt = self.activations1[i](x)
            xt = mx.swapaxes(xt, 1, 2)  # (B, C, T) -> (B, T, C)
            xt = self.convs1[i](xt)
            xt = mx.swapaxes(xt, 1, 2)  # (B, T, C) -> (B, C, T)
            xt = self.activations2[i](xt)
            xt = mx.swapaxes(xt, 1, 2)  # (B, C, T) -> (B, T, C)
            xt = self.convs2[i](xt)
            xt = mx.swapaxes(xt, 1, 2)  # (B, T, C) -> (B, C, T)
            x = xt + x  # Residual connection
        return x


def _linear_interpolate_1d_to_size(x: mx.array, new_size: int) -> mx.array:

    T = x.shape[-1]
    if new_size == T:
        return x

    new_positions = mx.linspace(0, T - 1, new_size)
    indices_low = mx.floor(new_positions).astype(mx.int32)
    indices_high = mx.minimum(indices_low + 1, T - 1)
    weights = new_positions - indices_low.astype(x.dtype)

    low_vals = mx.take(x, indices_low, axis=-1)
    high_vals = mx.take(x, indices_high, axis=-1)

    return low_vals + weights * (high_vals - low_vals)


class SineGen(nn.Module):

    def __init__(
        self,
        samp_rate: int,
        harmonic_num: int = 0,
        sine_amp: float = 0.1,
        noise_std: float = 0.003,
        voiced_threshold: float = 0,
        use_interpolation: bool = False,
        upsample_scale: int = 1,
    ):
        super().__init__()
        self.sine_amp = sine_amp
        self.noise_std = noise_std
        self.harmonic_num = harmonic_num
        self.sampling_rate = samp_rate
        self.voiced_threshold = voiced_threshold
        self.use_interpolation = use_interpolation
        self.upsample_scale = upsample_scale

    def _f02uv(self, f0: mx.array) -> mx.array:
        return (f0 > self.voiced_threshold).astype(mx.float32)

    def _f02sine_interpolation(self, f0_values: mx.array) -> mx.array:

        B, T, H = f0_values.shape

        rad_values = (f0_values / self.sampling_rate) % 1

        rand_ini = mx.random.uniform(shape=(B, H))
        rand_ini = mx.concatenate([mx.zeros((B, 1)), rand_ini[:, 1:]], axis=1)
        rad_values = rad_values.at[:, 0, :].add(rand_ini)

        # Downsample rad_values: (B, T, H) -> (B, T', H)
        # Transpose to (B, H, T) for interpolation along T axis
        rad_values_t = mx.swapaxes(rad_values, 1, 2)  # (B, H, T)
        T_down = max(1, T // self.upsample_scale)
        rad_values_down = _linear_interpolate_1d_to_size(
            rad_values_t, new_size=T_down
        )  # (B, H, T')
        rad_values_down = mx.swapaxes(rad_values_down, 1, 2)  # (B, T', H)

        # Cumsum at lower rate
        phase = mx.cumsum(rad_values_down, axis=1) * 2 * math.pi

        # Upsample phase back to original length T: (B, T', H) -> (B, T, H)
        # Scale phase by upsample_scale before interpolation
        phase_t = mx.swapaxes(phase, 1, 2) * self.upsample_scale  # (B, H, T')
        phase_up = _linear_interpolate_1d_to_size(phase_t, new_size=T)  # (B, H, T)
        phase_up = mx.swapaxes(phase_up, 1, 2)  # (B, T, H)

        return mx.sin(phase_up)

    def __call__(self, f0: mx.array) -> tuple:

        B, _, T = f0.shape

        harmonic_multipliers = mx.arange(1, self.harmonic_num + 2).reshape(1, -1, 1)
        F_mat = f0 * harmonic_multipliers / self.sampling_rate

        if self.use_interpolation:
            fn = mx.swapaxes(f0, 1, 2) * mx.arange(
                1, self.harmonic_num + 2
            )  # (B, T, H)
            sine_waves = self._f02sine_interpolation(fn) * self.sine_amp
            sine_waves = mx.swapaxes(sine_waves, 1, 2)
        else:
            theta_mat = 2 * math.pi * (mx.cumsum(F_mat, axis=-1) % 1)

            phase_vec = mx.random.uniform(
                low=-math.pi, high=math.pi, shape=(B, self.harmonic_num + 1, 1)
            )
            mask = mx.arange(self.harmonic_num + 1).reshape(1, -1, 1) > 0
            phase_vec = mx.where(mask, phase_vec, 0.0)

            sine_waves = self.sine_amp * mx.sin(theta_mat + phase_vec)

        uv = self._f02uv(f0)

        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
        noise = noise_amp * mx.random.normal(shape=sine_waves.shape)

        sine_waves = sine_waves * uv + noise

        return sine_waves, uv, noise


class SourceModuleHnNSF(nn.Module):

    def __init__(
        self,
        sampling_rate: int,
        upsample_scale: int,
        harmonic_num: int = 0,
        sine_amp: float = 0.1,
        add_noise_std: float = 0.003,
        voiced_threshod: float = 0,
        use_interpolation: bool = False,
    ):

        super().__init__()
        self.sine_amp = sine_amp
        self.noise_std = add_noise_std

        self.l_sin_gen = SineGen(
            sampling_rate,
            harmonic_num,
            sine_amp,
            add_noise_std,
            voiced_threshod,
            use_interpolation=use_interpolation,
            upsample_scale=upsample_scale,
        )

        self.l_linear = nn.Linear(harmonic_num + 1, 1)

    def __call__(self, x: mx.array) -> tuple:

        # Generate sine harmonics
        sine_wavs, uv, _ = self.l_sin_gen(mx.swapaxes(x, 1, 2))
        sine_wavs = mx.swapaxes(sine_wavs, 1, 2)
        uv = mx.swapaxes(uv, 1, 2)

        # Merge harmonics
        sine_merge = mx.tanh(self.l_linear(sine_wavs))

        # Source for noise branch, in the same shape as uv
        noise = mx.random.normal(shape=uv.shape) * self.sine_amp / 3

        return sine_merge, noise, uv


def stft(x: mx.array, n_fft: int, hop_length: int, window: mx.array) -> tuple:
    """
    Short-Time Fourier Transform.

    Args:
        x: Input signal (B, T)
        n_fft: FFT size
        hop_length: Hop length
        window: Window function

    Returns:
        Tuple of (real, imag) each of shape (B, n_fft//2+1, num_frames)
    """
    B, T = x.shape

    # Pad signal with reflect padding to match PyTorch's torch.stft behavior
    pad_length = n_fft // 2
    # Reflect padding: mirror the signal at the edges
    # Left pad: reverse of x[:, 1:pad_length+1]
    # Right pad: reverse of x[:, -(pad_length+1):-1]
    left_pad = x[:, 1 : pad_length + 1][:, ::-1]
    right_pad = x[:, -(pad_length + 1) : -1][:, ::-1]
    x_padded = mx.concatenate([left_pad, x, right_pad], axis=1)

    # Calculate number of frames
    num_frames = (x_padded.shape[1] - n_fft) // hop_length + 1

    # Create frames using vectorized slicing (avoid Python loop)
    # Generate all frame start indices at once
    frame_starts = mx.arange(num_frames) * hop_length  # (num_frames,)
    sample_offsets = mx.arange(n_fft)  # (n_fft,)
    # All indices: (num_frames, n_fft) where indices[f, s] = frame_starts[f] + s
    all_indices = frame_starts[:, None] + sample_offsets[None, :]  # (num_frames, n_fft)
    # Gather frames: x_padded is (B, T), gather along axis 1
    # Result shape: (B, num_frames, n_fft)
    frames = mx.take(x_padded, all_indices.flatten(), axis=1).reshape(
        B, num_frames, n_fft
    )
    # Transpose to (B, n_fft, num_frames)
    frames = mx.swapaxes(frames, 1, 2)

    # Apply window
    window_expanded = mx.reshape(window, (1, -1, 1))
    frames = frames * window_expanded

    # FFT
    # MLX doesn't have a direct rfft, so we use fft and take the first half
    fft_result = mx.fft.fft(frames, axis=1)

    # Take positive frequencies
    real = mx.real(fft_result[:, : n_fft // 2 + 1, :])
    imag = mx.imag(fft_result[:, : n_fft // 2 + 1, :])

    return real, imag


def istft(
    magnitude: mx.array, phase: mx.array, n_fft: int, hop_length: int, window: mx.array
) -> mx.array:
    """
    Inverse Short-Time Fourier Transform (pure MLX implementation).

    Args:
        magnitude: Magnitude (B, n_fft//2+1, num_frames)
        phase: Phase (B, n_fft//2+1, num_frames)
        n_fft: FFT size
        hop_length: Hop length
        window: Window function

    Returns:
        Reconstructed signal (B, T)
    """
    # Clip magnitude
    magnitude = mx.clip(magnitude, a_min=None, a_max=1e2)

    # Convert to complex
    real = magnitude * mx.cos(phase)
    imag = magnitude * mx.sin(phase)

    # Create full spectrum (conjugate symmetric)
    B, F, num_frames = real.shape

    # Mirror the spectrum for negative frequencies
    real_mirror = real[:, 1:-1, :][:, ::-1, :]
    imag_mirror = imag[:, 1:-1, :][:, ::-1, :]
    real_full = mx.concatenate([real, real_mirror], axis=1)
    imag_full = mx.concatenate([imag, -imag_mirror], axis=1)

    # Combine into complex
    spectrum = real_full + 1j * imag_full

    # Inverse FFT
    frames = mx.fft.ifft(spectrum, axis=1)
    frames = mx.real(frames)  # Take real part

    # Apply window
    window_expanded = mx.reshape(window, (1, -1, 1))
    frames = frames * window_expanded

    # Pure MLX overlap-add
    output_length = (num_frames - 1) * hop_length + n_fft

    # Create index arrays for scatter-add
    frame_offsets = mx.arange(num_frames) * hop_length
    sample_indices = mx.arange(n_fft)
    # indices[f, s] = frame_offsets[f] + sample_indices[s]
    indices = frame_offsets[:, None] + sample_indices[None, :]  # (num_frames, n_fft)
    indices_flat = indices.flatten()  # (num_frames * n_fft,)

    # Window squared for normalization - compute once (same for all batches)
    window_sq = window**2
    window_sum = mx.zeros(output_length)
    window_updates = mx.tile(window_sq, (num_frames,))
    window_sum = window_sum.at[indices_flat].add(window_updates)
    window_sum = mx.maximum(window_sum, 1e-8)

    # Vectorized overlap-add for all batch items at once
    # frames is (B, n_fft, num_frames), need (B, num_frames, n_fft) for indexing
    frame_data = mx.swapaxes(frames, 1, 2)  # (B, num_frames, n_fft)
    updates = frame_data.reshape(B, -1)  # (B, num_frames * n_fft)

    # Use vmap-style vectorized scatter-add via index expansion
    # Expand indices to (B, num_frames * n_fft) - same indices for all batches
    output = mx.zeros((B, output_length))
    # Process with a single batched at[] operation using advanced indexing
    # batch_indices: [0,0,...,0, 1,1,...,1, ...] for each element
    batch_indices = mx.repeat(mx.arange(B), num_frames * n_fft)
    flat_indices = mx.tile(indices_flat, (B,))
    # Scatter add using 2D indexing
    flat_output = output.flatten()
    # Convert 2D indices to 1D: batch * output_length + sample_idx
    linear_indices = batch_indices * output_length + flat_indices
    flat_output = flat_output.at[linear_indices].add(updates.flatten())
    output = flat_output.reshape(B, output_length)

    # Normalize with precomputed window_sum (broadcasts across batch)
    output = output / window_sum

    # Remove padding
    pad_length = n_fft // 2
    output = output[:, pad_length:-pad_length]

    return output


class HiFTGenerator(nn.Module):
    """
    HiFi-GAN with Neural Source Filter (HiFT-Net) generator.

    Combines mel-spectrogram upsampling with neural source filter
    for high-quality speech synthesis using ISTFT.

    Reference: https://arxiv.org/abs/2309.09493
    """

    def __init__(
        self,
        in_channels: int = 80,
        base_channels: int = 512,
        nb_harmonics: int = 8,
        sampling_rate: int = 22050,
        nsf_alpha: float = 0.1,
        nsf_sigma: float = 0.003,
        nsf_voiced_threshold: float = 10,
        upsample_rates: List[int] = [8, 8],
        upsample_kernel_sizes: List[int] = [16, 16],
        istft_params: Dict[str, int] = {"n_fft": 16, "hop_len": 4},
        resblock_kernel_sizes: List[int] = [3, 7, 11],
        resblock_dilation_sizes: List[List[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
        source_resblock_kernel_sizes: List[int] = [7, 11],
        source_resblock_dilation_sizes: List[List[int]] = [[1, 3, 5], [1, 3, 5]],
        lrelu_slope: float = 0.1,
        audio_limit: float = 0.99,
        f0_predictor: Optional[nn.Module] = None,
        use_interpolation: bool = False,
    ):
        """
        Args:
            in_channels: Number of input mel channels
            base_channels: Base number of channels
            nb_harmonics: Number of harmonics for NSF
            sampling_rate: Audio sampling rate
            nsf_alpha: NSF sine amplitude
            nsf_sigma: NSF noise std
            nsf_voiced_threshold: F0 threshold for voiced/unvoiced
            upsample_rates: Upsampling factors
            upsample_kernel_sizes: Upsampling kernel sizes
            istft_params: ISTFT parameters (n_fft, hop_len)
            resblock_kernel_sizes: ResBlock kernel sizes
            resblock_dilation_sizes: ResBlock dilations
            source_resblock_kernel_sizes: Source ResBlock kernel sizes
            source_resblock_dilation_sizes: Source ResBlock dilations
            lrelu_slope: LeakyReLU negative slope
            audio_limit: Audio clipping limit
            f0_predictor: Optional F0 prediction module
            use_interpolation: Use interpolation-based phase computation in NSF
        """
        super().__init__()

        self.out_channels = 1
        self.nb_harmonics = nb_harmonics
        self.sampling_rate = sampling_rate
        self.istft_params = istft_params
        self.lrelu_slope = lrelu_slope
        self.audio_limit = audio_limit

        self.num_kernels = len(resblock_kernel_sizes)
        self.num_upsamples = len(upsample_rates)

        upsample_scale = math.prod(upsample_rates) * istft_params["hop_len"]

        # Neural Source Filter
        self.m_source = SourceModuleHnNSF(
            sampling_rate=sampling_rate,
            upsample_scale=upsample_scale,
            harmonic_num=nb_harmonics,
            sine_amp=nsf_alpha,
            add_noise_std=nsf_sigma,
            voiced_threshod=nsf_voiced_threshold,
            use_interpolation=use_interpolation,
        )

        # F0 upsampler
        self.f0_upsample_scale = upsample_scale

        # Pre-convolution
        self.conv_pre = nn.Conv1d(in_channels, base_channels, 7, stride=1, padding=3)

        # Upsampling layers
        self.ups = []
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(
                nn.ConvTranspose1d(
                    base_channels // (2**i),
                    base_channels // (2 ** (i + 1)),
                    k,
                    stride=u,
                    padding=(k - u) // 2,
                )
            )

        # Source downsampling and resblocks for frequency domain fusion
        self.source_downs = []
        self.source_resblocks = []
        downsample_rates = [1] + upsample_rates[::-1][:-1]
        # Compute cumulative product manually (replaces np.cumprod)
        downsample_cum_rates = []
        cum_prod = 1
        for rate in downsample_rates:
            cum_prod *= rate
            downsample_cum_rates.append(cum_prod)

        for i, (u, k, d) in enumerate(
            zip(
                downsample_cum_rates[::-1],
                source_resblock_kernel_sizes,
                source_resblock_dilation_sizes,
            )
        ):
            if u == 1:
                self.source_downs.append(
                    nn.Conv1d(
                        istft_params["n_fft"] + 2,
                        base_channels // (2 ** (i + 1)),
                        1,
                        stride=1,
                    )
                )
            else:
                self.source_downs.append(
                    nn.Conv1d(
                        istft_params["n_fft"] + 2,
                        base_channels // (2 ** (i + 1)),
                        u * 2,
                        stride=u,
                        padding=u // 2,
                    )
                )
            self.source_resblocks.append(
                ResBlock(base_channels // (2 ** (i + 1)), k, d)
            )

        # Residual blocks after each upsampling
        self.resblocks = []
        for i in range(len(self.ups)):
            ch = base_channels // (2 ** (i + 1))
            for j, (k, d) in enumerate(
                zip(resblock_kernel_sizes, resblock_dilation_sizes)
            ):
                self.resblocks.append(ResBlock(ch, k, d))

        # Post-convolution (outputs to frequency domain)
        ch = base_channels // (2 ** len(self.ups))
        self.conv_post = nn.Conv1d(
            ch, istft_params["n_fft"] + 2, 7, stride=1, padding=3
        )

        # STFT window (Hann window - periodic, matching scipy fftbins=True)
        self.stft_window = hann_window_periodic(istft_params["n_fft"])

        self.f0_predictor = f0_predictor

    def _f0_upsample(self, f0: mx.array) -> mx.array:
        """Upsample F0 using nearest neighbor interpolation."""
        # f0 shape: (B, 1, T)
        # Use mx.repeat for efficient nearest-neighbor upsampling
        return mx.repeat(f0, self.f0_upsample_scale, axis=2)

    def _stft(self, x: mx.array) -> tuple:
        """Perform STFT on input signal."""
        return stft(
            x,
            self.istft_params["n_fft"],
            self.istft_params["hop_len"],
            self.stft_window,
        )

    def _istft(self, magnitude: mx.array, phase: mx.array) -> mx.array:
        """Perform inverse STFT."""
        return istft(
            magnitude,
            phase,
            self.istft_params["n_fft"],
            self.istft_params["hop_len"],
            self.stft_window,
        )

    def decode(self, x: mx.array, s: mx.array) -> mx.array:
        """
        Decode mel-spectrogram to waveform.

        Args:
            x: Mel-spectrogram (B, C, T)
            s: Source signal (B, 1, T_s)

        Returns:
            Generated waveform (B, T)
        """
        # STFT of source signal
        s_stft_real, s_stft_imag = self._stft(s.squeeze(1))
        s_stft = mx.concatenate([s_stft_real, s_stft_imag], axis=1)

        # Pre-convolution: (B, C, T) -> transpose -> conv -> transpose back
        x = mx.swapaxes(x, 1, 2)  # (B, C, T) -> (B, T, C)
        x = self.conv_pre(x)
        x = mx.swapaxes(x, 1, 2)  # (B, T, C) -> (B, C, T)

        for i in range(self.num_upsamples):
            x = nn.leaky_relu(x, negative_slope=self.lrelu_slope)
            # ConvTranspose1d: (B, C, T) -> transpose -> conv -> transpose back
            x = mx.swapaxes(x, 1, 2)  # (B, C, T) -> (B, T, C)
            x = self.ups[i](x)
            x = mx.swapaxes(x, 1, 2)  # (B, T, C) -> (B, C, T)

            if i == self.num_upsamples - 1:
                # Reflection pad: pad 1 sample on left using reflection
                # For (1, 0) reflection: prepend x[:, :, 1:2] to the beginning
                x = mx.concatenate([x[:, :, 1:2], x], axis=2)

            # Source fusion: (B, C, T) -> transpose -> conv -> transpose back
            si = mx.swapaxes(s_stft, 1, 2)  # (B, C, T) -> (B, T, C)
            si = self.source_downs[i](si)
            si = mx.swapaxes(si, 1, 2)  # (B, T, C) -> (B, C, T)
            si = self.source_resblocks[i](si)
            x = x + si

            # Apply residual blocks and average their outputs
            # Using mx.stack allows MLX's lazy evaluation to optimize the computation graph
            start_idx = i * self.num_kernels
            x = mx.mean(
                mx.stack(
                    [self.resblocks[start_idx + j](x) for j in range(self.num_kernels)],
                    axis=0,
                ),
                axis=0,
            )

        x = nn.leaky_relu(x, negative_slope=self.lrelu_slope)
        # conv_post: (B, C, T) -> transpose -> conv -> transpose back
        x = mx.swapaxes(x, 1, 2)  # (B, C, T) -> (B, T, C)
        x = self.conv_post(x)
        x = mx.swapaxes(x, 1, 2)  # (B, T, C) -> (B, C, T)

        # Split into magnitude and phase
        n_fft_half = self.istft_params["n_fft"] // 2 + 1
        magnitude = mx.exp(x[:, :n_fft_half, :])
        phase = mx.sin(x[:, n_fft_half:, :])  # sin is redundancy, matches original

        # Inverse STFT
        x = self._istft(magnitude, phase)
        x = mx.clip(x, -self.audio_limit, self.audio_limit)

        return x

    def __call__(self, speech_feat: mx.array, cache_source: mx.array = None) -> tuple:
        """
        Generate waveform from mel-spectrogram.

        Args:
            speech_feat: Mel-spectrogram (B, C, T)
            cache_source: Cached source for streaming

        Returns:
            Tuple of (waveform, source)
        """
        if cache_source is None:
            cache_source = mx.zeros((1, 1, 0))

        # Predict F0 from mel
        f0 = self.f0_predictor(speech_feat)

        # Upsample F0
        s = self._f0_upsample(mx.expand_dims(f0, 1))
        s = mx.swapaxes(s, 1, 2)  # (B, T, 1)

        # Generate source from F0
        s, _, _ = self.m_source(s)
        s = mx.swapaxes(s, 1, 2)  # (B, 1, T)

        # Use cache to avoid glitch in streaming
        if cache_source.shape[2] != 0:
            cache_len = cache_source.shape[2]
            s = mx.concatenate([cache_source, s[:, :, cache_len:]], axis=2)

        # Decode mel + source to audio
        generated_speech = self.decode(x=speech_feat, s=s)

        return generated_speech, s

    def inference(self, speech_feat: mx.array, cache_source: mx.array = None) -> tuple:
        """Inference-mode forward pass."""
        return self(speech_feat, cache_source)
