"""16 kHz log-mel, exactly as Swar's style encoders were trained on it."""
from __future__ import annotations

import torch
import torchaudio
from torch import nn

SR, N_FFT, WIN, HOP, N_MELS = 16000, 1024, 800, 200, 80
MEAN, STD = -4.0, 4.0


class Mel(nn.Module):
    def __init__(self):
        super().__init__()
        self.mel = torchaudio.transforms.MelSpectrogram(
            sample_rate=SR, n_fft=N_FFT, win_length=WIN, hop_length=HOP, n_mels=N_MELS,
            f_min=0.0, f_max=8000.0, center=True, pad_mode="reflect", power=2.0)

    def forward(self, wav: torch.Tensor) -> torch.Tensor:
        if wav.dtype == torch.int16:
            wav = wav.float() / 32768.0
        if wav.dim() == 1:
            wav = wav[None]
        return (torch.log(1e-5 + self.mel(wav)) - MEAN) / STD

    @staticmethod
    def frames(n_samples: int) -> int:
        return n_samples // HOP + 1