"""
europa.utils.dsp
~~~~~~~~~~~~~~~~
Low-level DSP utilities: STFT, mel filterbank, window functions.
"""

import numpy as np
from typing import Optional


def stft(
    y: np.ndarray,
    n_fft: int = 2048,
    hop_length: int = 512,
    win_length: Optional[int] = None,
    window: str = "hann",
) -> np.ndarray:
    """
    Short-Time Fourier Transform.

    Returns:
        np.ndarray: Complex STFT of shape (1 + n_fft//2, n_frames).
    """
    win_length = win_length or n_fft
    win = _get_window(window, win_length)
    if win_length < n_fft:
        win = np.pad(win, (0, n_fft - win_length))

    n_frames = 1 + (len(y) - n_fft) // hop_length
    if n_frames < 1:
        y = np.pad(y, (0, n_fft))
        n_frames = 1

    shape = (n_fft, n_frames)
    strides = (y.strides[0], y.strides[0] * hop_length)
    frames = np.lib.stride_tricks.as_strided(y, shape=shape, strides=strides)
    windowed = frames * win[:, None]
    return np.fft.rfft(windowed, axis=0)


def istft(
    D: np.ndarray,
    hop_length: int = 512,
    win_length: Optional[int] = None,
    window: str = "hann",
    length: Optional[int] = None,
) -> np.ndarray:
    """
    Inverse Short-Time Fourier Transform.

    Returns:
        np.ndarray: Reconstructed time-domain signal.
    """
    n_fft = (D.shape[0] - 1) * 2
    win_length = win_length or n_fft
    win = _get_window(window, win_length)
    if win_length < n_fft:
        win = np.pad(win, (0, n_fft - win_length))

    n_frames = D.shape[1]
    expected_len = n_fft + hop_length * (n_frames - 1)
    y = np.zeros(expected_len, dtype=np.float32)
    win_sum = np.zeros(expected_len, dtype=np.float32)

    for i in range(n_frames):
        start = i * hop_length
        frame = np.fft.irfft(D[:, i], n=n_fft).real
        y[start:start + n_fft] += frame * win
        win_sum[start:start + n_fft] += win ** 2

    win_sum = np.where(win_sum > 1e-8, win_sum, 1.0)
    y /= win_sum
    return y[:length] if length else y


def _get_window(name: str, length: int) -> np.ndarray:
    """Generate a window function."""
    if name == "hann":
        return np.hanning(length).astype(np.float32)
    elif name == "hamming":
        return np.hamming(length).astype(np.float32)
    elif name == "blackman":
        return np.blackman(length).astype(np.float32)
    elif name == "bartlett":
        return np.bartlett(length).astype(np.float32)
    elif name == "rect" or name == "rectangular":
        return np.ones(length, dtype=np.float32)
    else:
        raise ValueError(f"Unknown window type: {name!r}")


def _mel_filterbank(
    sr: int,
    n_fft: int,
    n_mels: int = 128,
    fmin: float = 0.0,
    fmax: float = None,
) -> np.ndarray:
    """
    Construct a Mel filterbank matrix.

    Returns:
        np.ndarray: Shape (n_mels, 1 + n_fft//2).
    """
    fmax = fmax or (sr / 2.0)
    n_bins = 1 + n_fft // 2
    freqs = np.fft.rfftfreq(n_fft, d=1.0 / sr)

    mel_min = _hz_to_mel(fmin)
    mel_max = _hz_to_mel(fmax)
    mel_points = np.linspace(mel_min, mel_max, n_mels + 2)
    hz_points = _mel_to_hz(mel_points)

    fb = np.zeros((n_mels, n_bins))
    for m in range(1, n_mels + 1):
        f_left   = hz_points[m - 1]
        f_center = hz_points[m]
        f_right  = hz_points[m + 1]
        for k, f in enumerate(freqs):
            if f_left <= f <= f_center:
                fb[m - 1, k] = (f - f_left) / (f_center - f_left + 1e-10)
            elif f_center < f <= f_right:
                fb[m - 1, k] = (f_right - f) / (f_right - f_center + 1e-10)
    return fb.astype(np.float32)


def _hz_to_mel(hz: float) -> float:
    return 2595.0 * np.log10(1.0 + hz / 700.0)


def _mel_to_hz(mel) -> np.ndarray:
    return 700.0 * (10.0 ** (mel / 2595.0) - 1.0)
