"""
europa.feature
~~~~~~~~~~~~~~
Audio feature extraction.

Features:
    mfcc            - Mel-Frequency Cepstral Coefficients
    mel_spectrogram - Mel-scaled spectrogram
    spectrogram     - Power/magnitude spectrogram (STFT)
    chroma          - Chromagram (pitch class energy)
    spectral_centroid
    spectral_bandwidth
    spectral_rolloff
    spectral_contrast
    zero_crossing_rate
    rms_energy
    tonnetz         - Tonal centroid features
"""

import numpy as np
from ..core.audio import Audio
from ..utils.dsp import stft as _stft, _mel_filterbank


def mfcc(
    audio: Audio,
    n_mfcc: int = 13,
    n_fft: int = 2048,
    hop_length: int = 512,
    n_mels: int = 128,
    fmin: float = 0.0,
    fmax: float = None,
) -> np.ndarray:
    """
    Compute Mel-Frequency Cepstral Coefficients (MFCCs).

    Args:
        audio      : Input Audio object (mono).
        n_mfcc     : Number of MFCCs to return. Default 13.
        n_fft      : FFT window size. Default 2048.
        hop_length : Hop size between frames. Default 512.
        n_mels     : Number of mel bands. Default 128.
        fmin       : Minimum frequency. Default 0.
        fmax       : Maximum frequency. None = sr/2.

    Returns:
        np.ndarray: MFCCs of shape (n_mfcc, n_frames).
    """
    from scipy.fftpack import dct
    mel = mel_spectrogram(audio, n_fft=n_fft, hop_length=hop_length,
                          n_mels=n_mels, fmin=fmin, fmax=fmax)
    log_mel = np.log(mel + 1e-9)
    return dct(log_mel, axis=0, type=2, norm='ortho')[:n_mfcc]


def mel_spectrogram(
    audio: Audio,
    n_fft: int = 2048,
    hop_length: int = 512,
    n_mels: int = 128,
    fmin: float = 0.0,
    fmax: float = None,
    power: float = 2.0,
) -> np.ndarray:
    """
    Compute a Mel-scaled spectrogram.

    Args:
        audio      : Input Audio (mono preferred).
        n_fft      : FFT window size.
        hop_length : Hop length in samples.
        n_mels     : Number of mel frequency bands.
        fmin       : Minimum frequency in Hz.
        fmax       : Maximum frequency. None = sr/2.
        power      : Exponent for the spectrogram (1=magnitude, 2=power).

    Returns:
        np.ndarray: Mel spectrogram of shape (n_mels, n_frames).
    """
    S = spectrogram(audio, n_fft=n_fft, hop_length=hop_length, power=power)
    fmax = fmax or (audio.sr / 2.0)
    mel_fb = _mel_filterbank(audio.sr, n_fft, n_mels, fmin, fmax)
    return mel_fb @ S


def spectrogram(
    audio: Audio,
    n_fft: int = 2048,
    hop_length: int = 512,
    win_length: int = None,
    window: str = "hann",
    power: float = 2.0,
) -> np.ndarray:
    """
    Compute a power or magnitude spectrogram via STFT.

    Args:
        audio      : Input Audio (mono).
        n_fft      : FFT size.
        hop_length : Hop size.
        win_length : Window length. None = n_fft.
        window     : Window type ('hann', 'hamming', etc.).
        power      : 1 = magnitude, 2 = power.

    Returns:
        np.ndarray: Spectrogram of shape (1 + n_fft//2, n_frames).
    """
    data = audio.data if audio.is_mono else np.mean(audio.data, axis=0)
    D = _stft(data, n_fft=n_fft, hop_length=hop_length,
               win_length=win_length, window=window)
    return np.abs(D) ** power


def chroma(
    audio: Audio,
    n_fft: int = 2048,
    hop_length: int = 512,
    n_chroma: int = 12,
    tuning: float = 0.0,
) -> np.ndarray:
    """
    Compute a chromagram from audio.

    Args:
        audio     : Input Audio.
        n_fft     : FFT size.
        hop_length: Hop size.
        n_chroma  : Number of chroma bins (default 12 = one per semitone).
        tuning    : Tuning offset in fractions of a bin.

    Returns:
        np.ndarray: Chromagram of shape (n_chroma, n_frames).
    """
    data = audio.data if audio.is_mono else np.mean(audio.data, axis=0)
    D = _stft(data, n_fft=n_fft, hop_length=hop_length)
    S = np.abs(D) ** 2

    # Build chroma filter bank
    freqs = np.fft.rfftfreq(n_fft, d=1.0 / audio.sr)
    chroma_fb = _chroma_filterbank(freqs, n_chroma, audio.sr, tuning)
    return np.dot(chroma_fb, S)


def spectral_centroid(
    audio: Audio,
    n_fft: int = 2048,
    hop_length: int = 512,
) -> np.ndarray:
    """
    Compute spectral centroid (brightness indicator).

    Returns:
        np.ndarray: Shape (1, n_frames) in Hz.
    """
    data = audio.data if audio.is_mono else np.mean(audio.data, axis=0)
    D = np.abs(_stft(data, n_fft=n_fft, hop_length=hop_length))
    freqs = np.fft.rfftfreq(n_fft, d=1.0 / audio.sr)
    norm = np.sum(D, axis=0, keepdims=True)
    norm = np.where(norm == 0, 1, norm)
    return np.sum(freqs[:, None] * D, axis=0, keepdims=True) / norm


def spectral_bandwidth(
    audio: Audio,
    n_fft: int = 2048,
    hop_length: int = 512,
    p: float = 2.0,
) -> np.ndarray:
    """
    Compute spectral bandwidth.

    Returns:
        np.ndarray: Shape (1, n_frames).
    """
    data = audio.data if audio.is_mono else np.mean(audio.data, axis=0)
    D = np.abs(_stft(data, n_fft=n_fft, hop_length=hop_length))
    freqs = np.fft.rfftfreq(n_fft, d=1.0 / audio.sr)
    centroid = spectral_centroid(audio, n_fft=n_fft, hop_length=hop_length)
    deviation = np.abs(freqs[:, None] - centroid)
    norm = np.sum(D, axis=0, keepdims=True)
    norm = np.where(norm == 0, 1, norm)
    return (np.sum((deviation ** p) * D, axis=0, keepdims=True) / norm) ** (1.0 / p)


def spectral_rolloff(
    audio: Audio,
    n_fft: int = 2048,
    hop_length: int = 512,
    roll_percent: float = 0.85,
) -> np.ndarray:
    """
    Compute spectral rolloff frequency.

    Args:
        roll_percent: Fraction of total spectral energy. Default 0.85.

    Returns:
        np.ndarray: Shape (1, n_frames) in Hz.
    """
    data = audio.data if audio.is_mono else np.mean(audio.data, axis=0)
    D = np.abs(_stft(data, n_fft=n_fft, hop_length=hop_length)) ** 2
    freqs = np.fft.rfftfreq(n_fft, d=1.0 / audio.sr)
    cumsum = np.cumsum(D, axis=0)
    threshold = roll_percent * cumsum[-1, :]
    rolloff_idx = np.argmax(cumsum >= threshold, axis=0)
    return freqs[rolloff_idx][None, :]


def zero_crossing_rate(
    audio: Audio,
    frame_length: int = 2048,
    hop_length: int = 512,
) -> np.ndarray:
    """
    Compute zero-crossing rate per frame.

    Returns:
        np.ndarray: Shape (1, n_frames).
    """
    from ..core.convert import _frame
    data = audio.data if audio.is_mono else np.mean(audio.data, axis=0)
    frames = _frame(data, frame_length, hop_length)
    zcr = np.sum(np.abs(np.diff(np.sign(frames), axis=0)), axis=0) / (2 * frame_length)
    return zcr[None, :]


def rms_energy(
    audio: Audio,
    frame_length: int = 2048,
    hop_length: int = 512,
) -> np.ndarray:
    """
    Compute root mean square energy per frame.

    Returns:
        np.ndarray: Shape (1, n_frames).
    """
    from ..core.convert import _frame
    data = audio.data if audio.is_mono else np.mean(audio.data, axis=0)
    frames = _frame(data, frame_length, hop_length)
    return np.sqrt(np.mean(frames ** 2, axis=0))[None, :]


def tonnetz(audio: Audio, n_fft: int = 4096, hop_length: int = 512) -> np.ndarray:
    """
    Compute tonal centroid features (Tonnetz).

    Based on the 6-dimensional tonal centroid space.

    Returns:
        np.ndarray: Shape (6, n_frames).
    """
    chroma_feat = chroma(audio, n_fft=n_fft, hop_length=hop_length)
    # Tonnetz transformation matrix
    phi = 2 * np.pi * np.arange(12) / 12.0
    T = np.array([
        np.sin(phi),
        np.cos(phi),
        np.sin(2 * phi),
        np.cos(2 * phi),
        np.sin(5 * phi),
        np.cos(5 * phi),
    ])
    chroma_norm = chroma_feat / (np.sum(chroma_feat, axis=0, keepdims=True) + 1e-9)
    return T @ chroma_norm


# ── Internal helpers ──────────────────────────────────────────────────────────

def _chroma_filterbank(freqs, n_chroma, sr, tuning=0.0):
    """Build a chroma filterbank."""
    n_bins = len(freqs)
    fb = np.zeros((n_chroma, n_bins))
    for i in range(1, n_bins):
        if freqs[i] <= 0:
            continue
        hz = freqs[i]
        midi = 12 * np.log2(hz / 440.0) + 69 + tuning
        chroma_bin = int(round(midi)) % n_chroma
        fb[chroma_bin, i] += 1.0
    return fb
