"""NeMo-compatible log-mel front-end (numpy only, no torch/librosa at runtime).

Mirrors ``FilterbankFeatures`` (eval mode) from NeMo Speech 3.0 with the
Sortformer streaming configuration:

    sample_rate=16000, n_fft=512, win_length=400 (hann, periodic=False),
    hop_length=160, n_mels=128, preemph=0.97, mag_power=2.0,
    log(x + 2**-24), normalize="NA", pad_to=16

``log_mel_spectrogram`` returns ``(num_frames, 128)`` float32 where
``num_frames = floor(num_samples / hop)`` rounded up to a multiple of 16.
"""

import numpy as np

SAMPLE_RATE = 16000
N_FFT = 512
WIN_LENGTH = 400
HOP_LENGTH = 160
N_MELS = 128
PREEMPH = 0.97
LOG_ZERO_GUARD = 2.0 ** -24
PAD_TO = 16

_MEL_CACHE = {}


def _hz_to_mel(freq):
    # Slaney mel scale (librosa default, htk=False)
    f_sp = 200.0 / 3.0
    log_step = np.log(6.4) / 27.0
    min_log_hz = 1000.0
    min_log_mel = min_log_hz / f_sp
    freq = np.asarray(freq, dtype=np.float64)
    mel_linear = freq / f_sp
    mel_log = min_log_mel + np.log(np.maximum(freq, min_log_hz) / min_log_hz) / log_step
    return np.where(freq >= min_log_hz, mel_log, mel_linear)


def _mel_to_hz(mel):
    f_sp = 200.0 / 3.0
    log_step = np.log(6.4) / 27.0
    min_log_hz = 1000.0
    min_log_mel = min_log_hz / f_sp
    mel = np.asarray(mel, dtype=np.float64)
    hz_linear = f_sp * mel
    hz_log = min_log_hz * np.exp(log_step * (mel - min_log_mel))
    return np.where(mel >= min_log_mel, hz_log, hz_linear)


def mel_filterbank(sample_rate: int = SAMPLE_RATE, n_fft: int = N_FFT, n_mels: int = N_MELS):
    """librosa-compatible Slaney-normalized Slaney-scale mel filterbank."""
    key = (sample_rate, n_fft, n_mels)
    if key in _MEL_CACHE:
        return _MEL_CACHE[key]

    fmin, fmax = 0.0, sample_rate / 2.0
    mels = np.linspace(_hz_to_mel(fmin), _hz_to_mel(fmax), n_mels + 2)
    hz = _mel_to_hz(mels)
    freqs = np.linspace(0.0, sample_rate / 2.0, 1 + n_fft // 2)

    fdiff = np.diff(hz)
    ramps = np.subtract.outer(hz, freqs)
    lower = -ramps[np.arange(n_mels), :] / fdiff[np.arange(n_mels)][:, None]
    upper = ramps[np.arange(2, n_mels + 2), :] / fdiff[np.arange(1, n_mels + 1)][:, None]
    weights = np.maximum(0.0, np.minimum(lower, upper))

    enorm = 2.0 / (hz[2 : n_mels + 2] - hz[:n_mels])
    weights *= enorm[:, None]
    result = weights.astype(np.float64)
    _MEL_CACHE[key] = result
    return result


def _hann_window(length: int):
    return 0.5 - 0.5 * np.cos(2.0 * np.pi * np.arange(length) / (length - 1))


def log_mel_spectrogram(
    waveform: np.ndarray,
    sample_rate: int = SAMPLE_RATE,
    *,
    preemph: float = PREEMPH,
    pad_to: int = PAD_TO,
    mel_filter: np.ndarray = None,
) -> np.ndarray:
    """Compute the NeMo Sortformer streaming log-mel features for a mono waveform."""
    if sample_rate != SAMPLE_RATE:
        raise ValueError(f"expected {SAMPLE_RATE} Hz input, got {sample_rate}")
    waveform = np.asarray(waveform, dtype=np.float32).reshape(-1)

    if preemph is not None and waveform.size > 0:
        emphasized = np.empty_like(waveform)
        emphasized[0] = waveform[0]
        emphasized[1:] = waveform[1:] - preemph * waveform[:-1]
        waveform = emphasized

    padded = np.pad(waveform, (N_FFT // 2, N_FFT // 2), mode="constant")
    num_samples = waveform.shape[0]
    num_frames = num_samples // HOP_LENGTH
    if num_frames == 0:
        return np.zeros((0, N_MELS), dtype=np.float32)

    window = _hann_window(WIN_LENGTH).astype(np.float32)
    left_pad = (N_FFT - WIN_LENGTH) // 2
    frame_starts = np.arange(num_frames) * HOP_LENGTH
    frames = np.lib.stride_tricks.as_strided(
        padded,
        shape=(num_frames, N_FFT),
        strides=(padded.strides[0] * HOP_LENGTH, padded.strides[0]),
        writeable=False,
    ).copy()
    frames[:, :left_pad] = 0.0
    frames[:, left_pad : left_pad + WIN_LENGTH] *= window
    frames[:, left_pad + WIN_LENGTH :] = 0.0

    spectrum = np.fft.rfft(frames, n=N_FFT, axis=1)
    magnitude = np.abs(spectrum).astype(np.float32) ** 2.0

    if mel_filter is None:
        mel_filter = mel_filterbank()
    mel = magnitude @ mel_filter.T.astype(np.float32)
    log_mel = np.log(mel + LOG_ZERO_GUARD, dtype=np.float32)

    if pad_to and log_mel.shape[0] % pad_to:
        pad = pad_to - log_mel.shape[0] % pad_to
        log_mel = np.pad(log_mel, ((0, pad), (0, 0)), mode="constant")
    return log_mel.astype(np.float32)