"""Mel spectrogram utilities for the UI and BiCrossMamba-ST input."""
from __future__ import annotations

from typing import List

import numpy as np
import torch
import torchaudio


_MEL_TRANSFORM_CACHE: dict = {}


def _get_mel(sr: int, n_mels: int, n_fft: int, hop_length: int) -> torchaudio.transforms.MelSpectrogram:
    key = (sr, n_mels, n_fft, hop_length)
    if key not in _MEL_TRANSFORM_CACHE:
        _MEL_TRANSFORM_CACHE[key] = torchaudio.transforms.MelSpectrogram(
            sample_rate=sr,
            n_fft=n_fft,
            hop_length=hop_length,
            n_mels=n_mels,
            f_min=0.0,
            f_max=sr // 2,
            power=2.0,
        )
    return _MEL_TRANSFORM_CACHE[key]


def mel_spectrogram(
    waveform: torch.Tensor,
    sample_rate: int = 16000,
    n_mels: int = 64,
    n_fft: int = 1024,
    hop_length: int = 256,
) -> torch.Tensor:
    """Compute a log-mel spectrogram. Returns shape [n_mels, T]."""
    mel = _get_mel(sample_rate, n_mels, n_fft, hop_length)
    spec = mel(waveform)              # [1, n_mels, T]
    log_spec = torch.log(spec + 1e-6)
    return log_spec.squeeze(0)


def mel_for_ui(
    waveform: torch.Tensor,
    sample_rate: int = 16000,
    n_mels: int = 64,
    n_fft: int = 1024,
    hop_length: int = 256,
    max_time_steps: int = 256,
) -> List[List[float]]:
    """Return a normalised [0,1] list-of-lists suitable for canvas rendering."""
    log_spec = mel_spectrogram(waveform, sample_rate, n_mels, n_fft, hop_length)
    arr = log_spec.detach().cpu().numpy()

    # Down-sample time axis if needed
    T = arr.shape[1]
    if T > max_time_steps:
        bucket = T // max_time_steps
        trimmed = arr[:, : bucket * max_time_steps]
        arr = trimmed.reshape(arr.shape[0], max_time_steps, bucket).mean(axis=2)

    # Per-clip min-max normalize
    lo, hi = float(arr.min()), float(arr.max())
    rng = max(hi - lo, 1e-6)
    arr = (arr - lo) / rng

    return arr.astype(np.float32).tolist()