"""
_audio_scene.py — MFCC + spectral nearest-centroid audio scene classifier (no ROS2 deps).

Classifies 1-second audio clips into one of four environment labels:
  'indoor'  — enclosed space (room tone ~440 Hz, low flux)
  'outdoor' — open ambient  (wind/ambient noise, ~1500 Hz band)
  'traffic' — road/engine   (low-frequency fundamental ~80 Hz)
  'park'    — natural space (bird-like high-frequency ~3–5 kHz)

Feature vector (16 dimensions)
--------------------------------
  [0..12]  13 MFCC coefficients (mean across frames, from 26-mel filterbank)
  [13]     Spectral centroid  (Hz, mean across frames)
  [14]     Spectral rolloff   (Hz at 85 % cumulative energy, mean across frames)
  [15]     Zero-crossing rate (crossings / sample, mean across frames)

Classifier
----------
  Nearest-centroid: class whose prototype centroid has the smallest L2 distance
  to the (per-dimension normalised) query feature vector.
  Centroids are computed once at import from seeded synthetic prototype signals,
  so they are always consistent with the feature extractor.

Confidence
----------
  conf = 1 / (1 + dist_to_nearest_centroid)   in normalised feature space.

Public API
----------
  SCENE_LABELS                        tuple[str, ...]  ordered class names
  AudioSceneResult                    NamedTuple       single-clip result
  NearestCentroidClassifier           class
    .predict(features) → (label, conf)
  extract_features(samples, sr)  → np.ndarray  shape (16,)
  classify_audio(samples, sr)    → AudioSceneResult

  _CLASSIFIER                         module-level singleton
"""

from __future__ import annotations

import math
from typing import NamedTuple, Tuple

import numpy as np


# ── Constants ─────────────────────────────────────────────────────────────────

SCENE_LABELS: Tuple[str, ...] = ('indoor', 'outdoor', 'traffic', 'park')

_SR_DEFAULT   = 16_000   # Hz
_N_MFCC       = 13
_N_MELS       = 26
_N_FFT        = 512
_WIN_LENGTH   = 400      # 25 ms @ 16 kHz
_HOP_LENGTH   = 160      # 10 ms @ 16 kHz
_N_FEATURES   = _N_MFCC + 3   # 16 total


# ── Result type ───────────────────────────────────────────────────────────────

class AudioSceneResult(NamedTuple):
    label:      str            # one of SCENE_LABELS
    confidence: float          # 0.0–1.0
    features:   np.ndarray     # shape (16,)  raw feature vector


# ── Low-level DSP helpers ─────────────────────────────────────────────────────

def _frame(x: np.ndarray, frame_length: int, hop_length: int) -> np.ndarray:
    """Slice a 1-D signal into overlapping frames. Returns (n_frames, frame_length)."""
    n_frames = max(1, 1 + (len(x) - frame_length) // hop_length)
    idx = (
        np.arange(frame_length)[np.newaxis, :]
        + np.arange(n_frames)[:, np.newaxis] * hop_length
    )
    # Pad x if the last index exceeds the signal length
    needed = int(idx.max()) + 1
    if needed > len(x):
        x = np.pad(x, (0, needed - len(x)))
    return x[idx]


def _mel_filterbank(
    sr: int,
    n_fft: int,
    n_mels: int,
    f_min: float = 0.0,
    f_max: float | None = None,
) -> np.ndarray:
    """Build (n_mels, n_fft//2+1) triangular mel filterbank matrix."""
    if f_max is None:
        f_max = sr / 2.0

    def _hz2mel(f: float) -> float:
        return 2595.0 * math.log10(1.0 + f / 700.0)

    def _mel2hz(m: float) -> float:
        return 700.0 * (10.0 ** (m / 2595.0) - 1.0)

    mel_pts = np.linspace(_hz2mel(f_min), _hz2mel(f_max), n_mels + 2)
    hz_pts  = np.array([_mel2hz(m) for m in mel_pts])

    freqs   = np.fft.rfftfreq(n_fft, d=1.0 / sr)     # (F,)
    f       = freqs[:, np.newaxis]                     # (F, 1)
    f_left  = hz_pts[:-2]                              # (n_mels,)
    f_ctr   = hz_pts[1:-1]
    f_right = hz_pts[2:]

    left_slope  = (f - f_left)  / np.maximum(f_ctr - f_left,  1e-10)
    right_slope = (f_right - f) / np.maximum(f_right - f_ctr, 1e-10)

    fbank = np.maximum(0.0, np.minimum(left_slope, right_slope)).T  # (n_mels, F)
    return fbank


def _dct2(x: np.ndarray) -> np.ndarray:
    """Type-II DCT of each row.  x: (n_frames, N) → (n_frames, N)."""
    N = x.shape[1]
    n = np.arange(N, dtype=np.float64)
    k = np.arange(N, dtype=np.float64)[:, np.newaxis]
    D = np.cos(math.pi * k * (2.0 * n + 1.0) / (2.0 * N))  # (N, N)
    return x @ D.T                                            # (n_frames, N)


# ── Feature extraction ────────────────────────────────────────────────────────

# Pre-build filterbank once (module-level, shared across all calls at default SR)
_FBANK: np.ndarray = _mel_filterbank(_SR_DEFAULT, _N_FFT, _N_MELS)
_WINDOW: np.ndarray = np.hamming(_WIN_LENGTH)
_FREQS:  np.ndarray = np.fft.rfftfreq(_N_FFT, d=1.0 / _SR_DEFAULT)


def extract_features(
    samples: np.ndarray,
    sr: int = _SR_DEFAULT,
) -> np.ndarray:
    """
    Extract a 16-d feature vector from a mono float audio clip.

    Parameters
    ----------
    samples : 1-D array, float32 or float64, range roughly [-1, 1]
    sr      : sample rate (Hz); must match the audio

    Returns
    -------
    features : shape (16,)  — [MFCC_0..12, centroid_hz, rolloff_hz, zcr]
    """
    x = np.asarray(samples, dtype=np.float64)
    if len(x) == 0:
        return np.zeros(_N_FEATURES, dtype=np.float64)

    # Rebuild filterbank/freqs if non-default sample rate
    if sr != _SR_DEFAULT:
        fbank = _mel_filterbank(sr, _N_FFT, _N_MELS)
        freqs = np.fft.rfftfreq(_N_FFT, d=1.0 / sr)
    else:
        fbank = _FBANK
        freqs = _FREQS

    # Frame + window
    frames = _frame(x, _WIN_LENGTH, _HOP_LENGTH)      # (T, W)
    frames = frames * _WINDOW                           # Hamming

    # Magnitude spectrum
    mag = np.abs(np.fft.rfft(frames, n=_N_FFT, axis=1))  # (T, F)

    # ── MFCC ────────────────────────────────────────────────────────────────
    mel_e  = mag @ fbank.T                             # (T, n_mels)
    log_me = np.log(np.maximum(mel_e, 1e-10))
    mfccs  = _dct2(log_me)[:, :_N_MFCC]               # (T, 13)
    mfcc_mean = mfccs.mean(axis=0)                     # (13,)

    # ── Spectral centroid ───────────────────────────────────────────────────
    power  = mag ** 2                                   # (T, F)
    p_sum  = power.sum(axis=1, keepdims=True)
    p_sum  = np.maximum(p_sum, 1e-20)
    sc_hz  = (power @ freqs) / p_sum.squeeze(1)        # (T,)  weighted mean freq
    centroid_mean = float(sc_hz.mean())

    # ── Spectral rolloff (85 %) ─────────────────────────────────────────────
    cumsum = np.cumsum(power, axis=1)                  # (T, F)
    thresh = 0.85 * p_sum                              # (T, 1)
    rolloff_idx = np.argmax(cumsum >= thresh, axis=1)  # (T,)  first bin ≥ 85 %
    rolloff_hz  = freqs[rolloff_idx]                   # (T,)
    rolloff_mean = float(rolloff_hz.mean())

    # ── Zero-crossing rate ──────────────────────────────────────────────────
    signs  = np.sign(x)
    zcr    = np.mean(np.abs(np.diff(signs)) / 2.0)    # crossings / sample
    zcr_val = float(zcr)

    features = np.concatenate([
        mfcc_mean,
        [centroid_mean, rolloff_mean, zcr_val],
    ]).astype(np.float64)
    return features


# ── Nearest-centroid classifier ───────────────────────────────────────────────

class NearestCentroidClassifier:
    """
    Classify a feature vector by nearest L2 distance in normalised feature space.

    Parameters
    ----------
    centroids : (n_classes, n_features) array — one row per class
    labels    : sequence of class name strings, same order as centroids rows
    """

    def __init__(
        self,
        centroids: np.ndarray,
        labels: Tuple[str, ...],
    ) -> None:
        self.centroids = np.asarray(centroids, dtype=np.float64)
        self.labels    = tuple(labels)
        # Per-dimension normalisation derived from centroid spread
        self._min   = self.centroids.min(axis=0)
        self._range = np.maximum(
            self.centroids.max(axis=0) - self._min, 1e-6
        )

    def predict(self, features: np.ndarray) -> Tuple[str, float]:
        """
        Return (label, confidence) for the given feature vector.

        confidence = 1 / (1 + min_distance)  in normalised feature space.
        """
        norm_feat = (features - self._min) / self._range
        norm_cent = (self.centroids - self._min) / self._range
        dists  = np.linalg.norm(norm_cent - norm_feat, axis=1)  # (n_classes,)
        best   = int(np.argmin(dists))
        conf   = float(1.0 / (1.0 + dists[best]))
        return self.labels[best], conf


# ── Prototype centroid construction ──────────────────────────────────────────

def _make_prototype(
    label: str,
    sr: int = _SR_DEFAULT,
    rng: np.random.RandomState | None = None,
) -> np.ndarray:
    """Generate a 1-second synthetic prototype signal for each scene class."""
    if rng is None:
        rng = np.random.RandomState(42)
    n = sr
    t = np.linspace(0.0, 1.0, n, endpoint=False)

    if label == 'indoor':
        # Room tone: dominant 440 Hz + light broadband noise
        sig = 0.5 * np.sin(2 * math.pi * 440.0 * t) + 0.05 * rng.randn(n)

    elif label == 'outdoor':
        # Wind/ambient: bandpass noise centred ~1500 Hz (800–2500 Hz)
        noise = rng.randn(n)
        # Band-pass via FFT zeroing
        spec = np.fft.rfft(noise)
        freqs = np.fft.rfftfreq(n, d=1.0 / sr)
        spec[(freqs < 800) | (freqs > 2500)] = 0.0
        sig = np.fft.irfft(spec, n=n).real * 0.5

    elif label == 'traffic':
        # Engine / road noise: 80 Hz fundamental + 2nd harmonic + light noise
        sig = (
            0.5 * np.sin(2 * math.pi * 80.0 * t)
            + 0.25 * np.sin(2 * math.pi * 160.0 * t)
            + 0.1 * rng.randn(n)
        )

    elif label == 'park':
        # Birdsong: high-frequency components at 3200 Hz + 4800 Hz
        sig = (
            0.4 * np.sin(2 * math.pi * 3200.0 * t)
            + 0.3 * np.sin(2 * math.pi * 4800.0 * t)
            + 0.05 * rng.randn(n)
        )

    else:
        raise ValueError(f'Unknown scene label: {label!r}')

    return sig.astype(np.float64)


def _build_centroids(sr: int = _SR_DEFAULT) -> np.ndarray:
    """Compute (4, 16) centroid matrix from prototype signals."""
    rng = np.random.RandomState(42)
    rows = []
    for label in SCENE_LABELS:
        proto = _make_prototype(label, sr=sr, rng=rng)
        rows.append(extract_features(proto, sr=sr))
    return np.stack(rows, axis=0)   # (4, 16)


# Module-level singleton — built once at import
_CENTROIDS:   np.ndarray          = _build_centroids()
_CLASSIFIER:  NearestCentroidClassifier = NearestCentroidClassifier(
    _CENTROIDS, SCENE_LABELS
)


# ── Main entry point ──────────────────────────────────────────────────────────

def classify_audio(
    samples: np.ndarray,
    sr: int = _SR_DEFAULT,
) -> AudioSceneResult:
    """
    Classify a mono audio clip into one of four scene categories.

    Parameters
    ----------
    samples : 1-D float array (any length ≥ 1 sample)
    sr      : sample rate in Hz

    Returns
    -------
    AudioSceneResult with .label, .confidence, .features
    """
    feat  = extract_features(samples, sr=sr)
    label, conf = _CLASSIFIER.predict(feat)
    return AudioSceneResult(label=label, confidence=conf, features=feat)
