"""Log-mel feature frontend.

40 log-mel bands, 16 kHz, 25 ms window / 10 ms hop. Deliberately matches
what is cheap to compute on the ESP32-S3 later (esp-dsp / TFLite-Micro
audio frontend), so the on-device features can be made to line up with
training.
"""

import torch
import torchaudio

SAMPLE_RATE = 16000
N_MELS = 40
WIN_LENGTH = 400   # 25 ms
HOP_LENGTH = 160   # 10 ms
N_FFT = 512


class LogMel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.mel = torchaudio.transforms.MelSpectrogram(
            sample_rate=SAMPLE_RATE,
            n_fft=N_FFT,
            win_length=WIN_LENGTH,
            hop_length=HOP_LENGTH,
            n_mels=N_MELS,
            center=True,
            power=2.0,
        )

    EMA_ALPHA = 0.02  # ~0.5 s time constant at 10 ms hop

    def forward(self, waveform):
        """waveform (B, samples) -> features (B, T, N_MELS).

        Normalization is a causal per-channel EMA mean subtraction so the
        exact same computation can run frame-by-frame on the device.
        """
        mel = self.mel(waveform)                    # (B, n_mels, T)
        logmel = torch.log(mel + 1e-6)
        a = self.EMA_ALPHA
        ema = torchaudio.functional.lfilter(
            logmel,
            a_coeffs=torch.tensor([1.0, -(1.0 - a)], device=logmel.device),
            b_coeffs=torch.tensor([a, 0.0], device=logmel.device),
            clamp=False,
        )
        logmel = logmel - ema
        return logmel.transpose(1, 2)               # (B, T, n_mels)


def num_frames(num_samples):
    """Frame count produced for a waveform length (center=True)."""
    return num_samples // HOP_LENGTH + 1