"""Spectral shape features from a short-time magnitude spectrum.

Five low-level descriptors that summarise the *shape* of one frame's spectrum
in a handful of numbers: where its energy sits (centroid), how spread out it is
(spread / bandwidth), where most of it lies below (rolloff), how fast it changes
frame to frame (flux), and whether it is tonal or noise-like (flatness). They
are the bread-and-butter front end of music information retrieval and of cheap
voice-activity and timbre classifiers, computed straight from the same FFT bins
the spectrogram already produced.

Every function takes a one-sided magnitude spectrum ``mag`` (the magnitude of an
``rfft``, i.e. ``np.abs(np.fft.rfft(x))``, or ``np.abs`` of the STFT in
``topics/stft``) and the matching frequency axis ``freqs``
(``np.fft.rfftfreq(n, 1/fs)``). Note the inputs are magnitudes, not the complex
``rfft`` output itself. ``mag`` may be a single frame, shape ``(nfreq,)``, or a whole
spectrogram, shape ``(nfreq, nframes)``; the reductions run along the frequency
axis (axis 0) so a spectrogram yields one feature value per frame.

Weighting convention: centroid, spread and rolloff weight the bins by
**magnitude** ``|X[k]|`` (the common MIR convention, e.g. librosa's default and
Tzanetakis & Cook 2002). Weighting by power ``|X[k]|**2`` instead is an equally
valid variant that leans harder on the strongest bins; pass ``mag**2`` if you
want it. Flatness is defined on the **power** spectrum, where it is the classic
Wiener entropy (geometric mean over arithmetic mean).

References: G. Tzanetakis and P. Cook, "Musical genre classification of audio
signals," IEEE Trans. Speech Audio Process. 10(5):293-302, 2002 (centroid,
rolloff, flux). G. Peeters, "A large set of audio features for sound
description," CUIDADO project tech. report, IRCAM, 2004 (spread, flatness).
"""

import numpy as np

_EPS = 1e-12


def _freq_col(freqs, mag):
    """Reshape ``freqs`` so it broadcasts against ``mag`` along axis 0."""
    freqs = np.asarray(freqs, dtype=float)
    shape = [freqs.size] + [1] * (np.ndim(mag) - 1)
    return freqs.reshape(shape)


def spectral_centroid(freqs, mag):
    """Magnitude-weighted mean frequency of each frame (Hz).

    The "centre of mass" of the spectrum, and the usual proxy for perceived
    brightness: a bright, buzzy sound has a high centroid, a dull or boomy one a
    low centroid. Frames with no energy return 0.
    """
    mag = np.asarray(mag, dtype=float)
    f = _freq_col(freqs, mag)
    total = mag.sum(axis=0)
    num = (f * mag).sum(axis=0)
    return np.where(total > _EPS, num / np.where(total > _EPS, total, 1.0), 0.0)


def spectral_spread(freqs, mag):
    """Magnitude-weighted standard deviation about the centroid (Hz).

    The second central moment, also called spectral bandwidth: how far the
    energy spreads around the centroid. A pure tone has spread ~0; broadband
    noise has a large spread. Frames with no energy return 0.
    """
    mag = np.asarray(mag, dtype=float)
    f = _freq_col(freqs, mag)
    c = spectral_centroid(freqs, mag)
    total = mag.sum(axis=0)
    var = (((f - c) ** 2) * mag).sum(axis=0)
    var = np.where(total > _EPS, var / np.where(total > _EPS, total, 1.0), 0.0)
    return np.sqrt(var)


def spectral_rolloff(freqs, mag, roll_percent=0.85):
    """Frequency below which ``roll_percent`` of the magnitude lies (Hz).

    The rolloff is the lowest frequency whose running sum of magnitude reaches
    the given fraction of the frame total. It marks the upper edge of where the
    energy is concentrated, separating (say) a low-pitched voiced sound from a
    hissy fricative. ``roll_percent`` defaults to 0.85 (Tzanetakis & Cook 2002).
    Frames with no energy return 0.
    """
    if not 0.0 < roll_percent <= 1.0:
        raise ValueError("roll_percent must be in (0, 1]")
    mag = np.asarray(mag, dtype=float)
    freqs = np.asarray(freqs, dtype=float)
    cumulative = np.cumsum(mag, axis=0)
    total = cumulative[-1]
    threshold = roll_percent * total                     # share of the frame total
    reached = cumulative >= threshold
    idx = np.argmax(reached, axis=0)                     # first bin that reaches it
    # a silent frame has no rolloff; return 0 rather than freqs[0], which is only
    # 0 when the axis starts at DC, not for a sub-band slice
    return np.where(total > _EPS, freqs[idx], 0.0)


def spectral_flux(mag, normalize=True):
    """How much the spectrum changes from one frame to the next.

    The L2 distance between successive (optionally energy-normalised) magnitude
    spectra. It peaks at onsets and transients and stays low while a sound is
    steady. Needs a spectrogram, shape ``(nfreq, nframes)``; returns one value
    per transition, shape ``(nframes - 1,)``.

    With ``normalize=True`` each frame is scaled to unit L2 norm first, so flux
    measures a change of *shape* rather than of loudness. The half-wave
    rectified variant (keep only bins that grew) is the classic onset detector;
    this returns the full symmetric distance.
    """
    mag = np.asarray(mag, dtype=float)
    if mag.ndim != 2:
        raise ValueError("flux needs a (nfreq, nframes) spectrogram")
    if normalize:
        norms = np.sqrt((mag ** 2).sum(axis=0, keepdims=True))
        mag = np.divide(mag, norms, out=np.zeros_like(mag), where=norms > _EPS)
    diff = np.diff(mag, axis=1)
    return np.sqrt((diff ** 2).sum(axis=0))


def spectral_flatness(mag):
    """Wiener entropy: geometric mean over arithmetic mean of the power spectrum.

    A number in (0, 1]. Close to 1 means a flat, noise-like spectrum (every bin
    similar); close to 0 means a peaky, tonal one (energy in a few bins). Often
    quoted in dB as ``10*log10(flatness)``. Computed on the power spectrum
    ``|X|**2`` with a small floor so silent bins do not send the log to minus
    infinity.
    """
    mag = np.asarray(mag, dtype=float)
    power = np.maximum(mag ** 2, _EPS)
    gmean = np.exp(np.mean(np.log(power), axis=0))
    amean = np.mean(power, axis=0)
    return gmean / amean


def frame_features(freqs, mag, roll_percent=0.85):
    """Bundle the per-frame features (everything except the inter-frame flux).

    Returns a dict with ``centroid``, ``spread``, ``rolloff`` and ``flatness``,
    each a scalar for a single frame or a length-``nframes`` array for a
    spectrogram. Flux is excluded because it is defined between frames; compute
    it from the full spectrogram with :func:`spectral_flux`.
    """
    return {
        "centroid": spectral_centroid(freqs, mag),
        "spread": spectral_spread(freqs, mag),
        "rolloff": spectral_rolloff(freqs, mag, roll_percent),
        "flatness": spectral_flatness(mag),
    }