"""Extraction de mel-spectrogrammes **écrite à la main** (Phase 7).

Aucune bibliothèque d'audio pré-entraînée : STFT via `torch.fft`, banc de filtres
mel triangulaires construit à la main, puis log. Utilisé par l'encodeur audio.

Justification : docs/architecture.md §6.
"""

from __future__ import annotations

import math

import torch


def hann_window(n: int, device=None, dtype=torch.float32) -> torch.Tensor:
    """Fenêtre de Hann périodique (comme librosa/torch)."""
    t = torch.arange(n, device=device, dtype=dtype)
    return 0.5 - 0.5 * torch.cos(2 * math.pi * t / n)


def frame_signal(x: torch.Tensor, frame_length: int, hop_length: int,
                 center: bool = True) -> torch.Tensor:
    """Découpe ``(B, N)`` en trames ``(B, n_frames, frame_length)``.

    ``center=True`` ajoute un padding de réflexion de ``frame_length//2`` de chaque
    côté (comme torch.stft), ce qui aligne les trames sur les échantillons.
    """
    if x.dim() == 1:
        x = x.unsqueeze(0)
    if center:
        pad = frame_length // 2
        x = torch.nn.functional.pad(x.unsqueeze(1), (pad, pad), mode="reflect").squeeze(1)
    n_frames = 1 + (x.size(1) - frame_length) // hop_length
    idx = (torch.arange(n_frames, device=x.device) * hop_length).unsqueeze(1)
    idx = idx + torch.arange(frame_length, device=x.device).unsqueeze(0)
    return x[:, idx]  # (B, n_frames, frame_length)


def stft_magnitude(x: torch.Tensor, n_fft: int, hop_length: int,
                   win_length: int | None = None) -> torch.Tensor:
    """Magnitude de la STFT ``(B, n_frames, n_fft//2 + 1)``."""
    win_length = win_length or n_fft
    window = hann_window(win_length, device=x.device, dtype=x.dtype)
    frames = frame_signal(x, win_length, hop_length) * window
    spec = torch.fft.rfft(frames, n=n_fft, dim=-1)
    return spec.abs()


def hz_to_mel(f: torch.Tensor | float) -> torch.Tensor:
    """Échelle mel (formule de O'Shaughnessy)."""
    f = torch.as_tensor(f, dtype=torch.float64)
    return 2595.0 * torch.log10(1.0 + f / 700.0)


def mel_to_hz(m: torch.Tensor | float) -> torch.Tensor:
    m = torch.as_tensor(m, dtype=torch.float64)
    return 700.0 * (10.0 ** (m / 2595.0) - 1.0)


def mel_filterbank(sr: int, n_fft: int, n_mels: int, fmin: float = 0.0,
                   fmax: float | None = None) -> torch.Tensor:
    """Banc de filtres mel triangulaires ``(n_mels, n_fft//2 + 1)``."""
    fmax = fmax if fmax is not None else sr / 2
    n_freqs = n_fft // 2 + 1
    mel_pts = torch.linspace(float(hz_to_mel(fmin)), float(hz_to_mel(fmax)), n_mels + 2)
    hz_pts = mel_to_hz(mel_pts).to(torch.float32)             # (n_mels+2,)
    fft_freqs = torch.linspace(0, sr / 2, n_freqs)            # (n_freqs,)

    fb = torch.zeros(n_mels, n_freqs)
    for i in range(n_mels):
        left, center, right = hz_pts[i], hz_pts[i + 1], hz_pts[i + 2]
        rising = (fft_freqs - left) / max(float(center - left), 1e-8)
        falling = (right - fft_freqs) / max(float(right - center), 1e-8)
        fb[i] = torch.clamp(torch.minimum(rising, falling), min=0.0)
    return fb


class LogMelSpectrogram(torch.nn.Module):
    """Transforme une forme d'onde ``(B, N)`` en log-mel ``(B, 1, n_mels, n_frames)``."""

    def __init__(self, sr: int = 16000, n_fft: int = 256, hop_length: int = 64,
                 n_mels: int = 64, fmin: float = 0.0, fmax: float | None = None):
        super().__init__()
        self.sr, self.n_fft, self.hop_length = sr, n_fft, hop_length
        self.register_buffer("fb", mel_filterbank(sr, n_fft, n_mels, fmin, fmax))

    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        mag = stft_magnitude(waveform, self.n_fft, self.hop_length)  # (B, T, F)
        mel = torch.einsum("btf,mf->btm", mag, self.fb)              # (B, T, n_mels)
        log_mel = torch.log(mel.clamp(min=1e-6))
        return log_mel.transpose(1, 2).unsqueeze(1)                  # (B, 1, n_mels, T)