"""Mel-Frequency Cepstral Coefficients: the classic speech feature front end.

An MFCC vector is a short, decorrelated summary of one frame's spectral
envelope, warped to match how human hearing resolves frequency. It is the
feature that powered a generation of speech recognisers and still anchors many
small-footprint keyword spotters, because the whole pipeline is cheap and the
output is a handful of numbers per frame.

The pipeline is a fixed sequence of small steps, each of which this module
implements from first principles rather than wrapping ``librosa`` so every stage
is visible and testable:

1. pre-emphasis: a one-tap high-pass that lifts the weak high end of voiced
   speech (:func:`pre_emphasis`);
2. framing and windowing: cut the signal into short overlapping frames and taper
   each (:func:`frame_signal`);
3. power spectrum: the periodogram of each frame (:func:`power_spectrum`);
4. mel filterbank: sum the power into a few perceptually-spaced bands
   (:func:`mel_filterbank`), then take the log (:func:`log_mel_spectrogram`);
5. DCT: a type-II discrete cosine transform that decorrelates the log-band
   energies and keeps the low-order coefficients (:func:`dct2`, :func:`mfcc`);
6. deltas: first and second time differences add coarse dynamics
   (:func:`delta`).

Conventions: input is a real 1-D signal at sample rate ``fs``. Frame and hop
lengths are in samples. The mel warp uses the common O'Shaughnessy 1987 closed
form ``mel = 2595 log10(1 + f/700)``; the perceptual scale it approximates is
Stevens, Volkmann & Newman 1937. The cepstral pipeline and its use for word
recognition are Davis & Mermelstein 1980.

References: S. Davis and P. Mermelstein, "Comparison of parametric
representations for monosyllabic word recognition in continuously spoken
sentences," IEEE Trans. ASSP 28(4):357-366, 1980. S. S. Stevens, J. Volkmann
and E. B. Newman, "A scale for the measurement of the psychological magnitude
pitch," J. Acoust. Soc. Am. 8(3):185-190, 1937. D. O'Shaughnessy, Speech
Communication: Human and Machine, Addison-Wesley, 1987 (the 2595/700 mel form).
"""

import numpy as np
from scipy.signal import get_window

_EPS = 1e-12


# --- frequency warp ---------------------------------------------------------

def hz_to_mel(f):
    """Hertz to mel, O'Shaughnessy 1987 form: ``2595 log10(1 + f/700)``."""
    f = np.asarray(f, dtype=float)
    return 2595.0 * np.log10(1.0 + f / 700.0)


def mel_to_hz(m):
    """Mel back to hertz, the exact inverse of :func:`hz_to_mel`."""
    m = np.asarray(m, dtype=float)
    return 700.0 * (10.0 ** (m / 2595.0) - 1.0)


# --- stage 1: pre-emphasis --------------------------------------------------

def pre_emphasis(x, coeff=0.97):
    """First-order high-pass pre-emphasis ``y[n] = x[n] - coeff*x[n-1]``.

    Voiced speech rolls off at roughly -6 dB per octave above a few hundred Hz,
    so the high formants arrive weak. This one-tap filter tilts the spectrum
    back up by about +6 dB per octave, balancing the energy across bands before
    the filterbank sees it. ``y[0]`` keeps ``x[0]`` (no sample exists to its
    left). ``coeff`` is conventionally 0.95-0.97.
    """
    x = np.asarray(x, dtype=float)
    y = np.empty_like(x)
    y[0] = x[0]
    y[1:] = x[1:] - coeff * x[:-1]
    return y


# --- stage 2: framing and windowing ----------------------------------------

def frame_signal(x, frame_len, hop, window="hamming"):
    """Cut ``x`` into overlapping windowed frames.

    Returns an array of shape ``(n_frames, frame_len)`` holding every frame that
    fits fully inside the signal, each multiplied by ``window``. This is the same
    framing the [STFT](../stft/stft.py) uses; speech MFCCs traditionally taper
    with a Hamming window over ~25 ms frames at a ~10 ms hop.
    """
    x = np.asarray(x, dtype=float)
    if frame_len <= 0 or hop <= 0:
        raise ValueError("frame_len and hop must be positive")
    win = _resolve_window(window, frame_len)
    if len(x) < frame_len:
        return np.empty((0, frame_len))
    starts = np.arange(0, len(x) - frame_len + 1, hop, dtype=int)
    return np.stack([x[s:s + frame_len] * win for s in starts])


# --- stage 3: power spectrum -------------------------------------------------

def power_spectrum(frames, n_fft):
    """Periodogram of each frame: ``|rfft|**2 / n_fft``.

    Takes ``(n_frames, frame_len)`` and returns ``(n_frames, n_fft//2+1)``. The
    ``1/n_fft`` scaling makes the result a power estimate, but every later stage
    is invariant to a constant scale, so the choice is cosmetic.
    """
    frames = np.asarray(frames, dtype=float)
    spec = np.fft.rfft(frames, n=n_fft, axis=-1)
    return (np.abs(spec) ** 2) / n_fft


# --- stage 4: mel filterbank + log ------------------------------------------

def mel_filterbank(n_filters, n_fft, fs, fmin=0.0, fmax=None):
    """Triangular mel filterbank as a ``(n_filters, n_fft//2+1)`` matrix.

    Builds ``n_filters`` triangular weighting functions whose centre frequencies
    are equally spaced on the mel scale between ``fmin`` and ``fmax`` (default
    Nyquist). Each triangle rises from its left neighbour's centre to its own and
    falls to its right neighbour's, peaking at 1 in the continuous mel domain. The
    discretized row maximum is close to 1 but depends on how well the centre
    frequency lands on an FFT bin; on a coarse grid the lowest, narrowest bands can
    peak lower (around 0.85 with a 512-point FFT at 16 kHz). Multiplying a power
    spectrum by this matrix sums the energy into perceptually-spaced bands. Returns
    the matrix; apply it with ``fb @ power`` (power as a column) or ``power @ fb.T``.
    """
    if fmax is None:
        fmax = fs / 2.0
    n_bins = n_fft // 2 + 1

    # equally-spaced centres on the mel scale, with the two extra edge points
    # that bound the first and last triangles
    mel_pts = np.linspace(hz_to_mel(fmin), hz_to_mel(fmax), n_filters + 2)
    hz_pts = mel_to_hz(mel_pts)
    bin_freqs = np.fft.rfftfreq(n_fft, 1.0 / fs)        # Hz of each FFT bin

    fb = np.zeros((n_filters, n_bins))
    for m in range(1, n_filters + 1):
        f_left, f_centre, f_right = hz_pts[m - 1], hz_pts[m], hz_pts[m + 1]
        # rising edge, then falling edge; clipped to [0, 1] and to the band
        rising = (bin_freqs - f_left) / max(f_centre - f_left, _EPS)
        falling = (f_right - bin_freqs) / max(f_right - f_centre, _EPS)
        fb[m - 1] = np.clip(np.minimum(rising, falling), 0.0, None)
    return fb


def log_mel_spectrogram(x, fs, n_filters=26, frame_len=400, hop=160,
                        n_fft=512, window="hamming", coeff=0.97, fmin=0.0,
                        fmax=None):
    """Pre-emphasis, framing, power spectrum, mel filterbank, log.

    The full front half of the MFCC pipeline: returns ``(n_frames, n_filters)``
    log mel-band energies. These are still strongly correlated band to band; the
    DCT in :func:`mfcc` is what decorrelates them.
    """
    emphasised = pre_emphasis(x, coeff)
    frames = frame_signal(emphasised, frame_len, hop, window)
    power = power_spectrum(frames, n_fft)               # (n_frames, n_bins)
    fb = mel_filterbank(n_filters, n_fft, fs, fmin, fmax)
    mel_energy = power @ fb.T                           # (n_frames, n_filters)
    return np.log(np.maximum(mel_energy, _EPS))


# --- stage 5: the DCT --------------------------------------------------------

def dct2(x, axis=-1):
    """Orthonormal type-II DCT along ``axis``.

    The same transform as ``scipy.fft.dct(x, type=2, norm='ortho')``, written
    out so the cepstral step is self-contained:

    ``X[k] = s(k) * sum_n x[n] cos(pi*(2n+1)*k / (2N))``,

    with ``s(0) = sqrt(1/N)`` and ``s(k) = sqrt(2/N)`` for ``k > 0``. The DCT-II
    is real, orthonormal, and concentrates the energy of smooth, correlated
    inputs (a log mel envelope) into its first few coefficients, which is exactly
    why MFCCs keep only the low-order terms.
    """
    x = np.asarray(x, dtype=float)
    x = np.moveaxis(x, axis, -1)
    N = x.shape[-1]
    n = np.arange(N)
    k = n.reshape(-1, 1)
    basis = np.cos(np.pi * (2 * n + 1) * k / (2 * N))   # (N, N): rows are k
    scale = np.full(N, np.sqrt(2.0 / N))
    scale[0] = np.sqrt(1.0 / N)
    X = (x @ basis.T) * scale
    return np.moveaxis(X, -1, axis)


def mfcc(x, fs, n_mfcc=13, n_filters=26, frame_len=400, hop=160, n_fft=512,
         window="hamming", coeff=0.97, fmin=0.0, fmax=None):
    """Mel-frequency cepstral coefficients of a signal.

    Runs the whole pipeline and keeps the first ``n_mfcc`` DCT coefficients of
    each frame's log mel spectrum. Returns ``(n_frames, n_mfcc)``. Coefficient 0
    is the overall log energy of the frame; the higher coefficients describe the
    shape of the spectral envelope and are what carry phonetic identity.
    """
    log_mel = log_mel_spectrogram(x, fs, n_filters, frame_len, hop, n_fft,
                                  window, coeff, fmin, fmax)
    cepstra = dct2(log_mel, axis=-1)
    return cepstra[:, :n_mfcc]


# --- stage 6: deltas ---------------------------------------------------------

def delta(feat, width=2):
    """Regression-based delta (time-derivative) coefficients.

    For each coefficient track this fits a slope over a ``+/- width`` window:

    ``d[t] = sum_{n=1..width} n*(c[t+n] - c[t-n]) / (2 * sum_{n=1..width} n^2)``.

    Edges are handled by replicating the first and last frames, so the output has
    the same shape as ``feat`` (``(n_frames, n_mfcc)``). Apply it twice for the
    delta-delta (acceleration) coefficients. Stacking ``[mfcc, delta, delta2]``
    is the standard 39-dimensional speech feature.
    """
    feat = np.asarray(feat, dtype=float)
    if width < 1:
        raise ValueError("width must be >= 1")
    n_frames = feat.shape[0]
    padded = np.pad(feat, ((width, width), (0, 0)), mode="edge")
    denom = 2.0 * sum(n * n for n in range(1, width + 1))
    out = np.zeros_like(feat)
    for n in range(1, width + 1):
        # padded[t+width] is feat[t]; +/-n around it in padded coordinates
        out += n * (padded[width + n: width + n + n_frames]
                    - padded[width - n: width - n + n_frames])
    return out / denom


# --- helpers -----------------------------------------------------------------

def _resolve_window(window, frame_len):
    """Return a length-frame_len window array from a name or an array."""
    if isinstance(window, str):
        return get_window(window, frame_len, fftbins=True)
    win = np.asarray(window, dtype=float)
    if win.shape != (frame_len,):
        raise ValueError("explicit window must have length frame_len")
    return win