"""ambient_sound_node.py -- Ambient sound classifier via mel-spectrogram features.
Issue #252

Accumulates 1 s of PCM-16 audio from /social/speech/audio_raw, extracts a
compact mel-spectrogram feature vector, then classifies the scene into one of:

  silence | speech | music | crowd | outdoor | alarm

Publishes the label as std_msgs/String on /saltybot/ambient_sound at 1 Hz.

Signal processing is pure Python + numpy (no torch / onnx dependency).

Feature vector (per 1-s window):
  energy_db    -- overall RMS in dBFS
  zcr          -- mean zero-crossing rate across frames
  mel_centroid -- centre-of-mass of the mel band energies  [0..1]
  mel_flatness -- geometric/arithmetic mean of mel energies [0..1]
                  (1 = white noise, 0 = single sinusoid)
  low_ratio    -- fraction of mel energy in lower third of bands
  high_ratio   -- fraction of mel energy in upper third of bands

Classification cascade (priority-ordered):
  silence  : energy_db < silence_db
  alarm    : energy_db >= alarm_db_min  AND zcr >= alarm_zcr_min
             AND high_ratio >= alarm_high_ratio_min
  speech   : zcr in [speech_zcr_min, speech_zcr_max]
             AND mel_flatness < speech_flatness_max
  music    : zcr < music_zcr_max AND mel_flatness < music_flatness_max
  crowd    : zcr >= crowd_zcr_min AND mel_flatness >= crowd_flatness_min
  outdoor  : catch-all

Parameters:
  sample_rate           (int,   16000)
  window_s              (float, 1.0)   -- accumulation window before classify
  n_fft                 (int,   512)   -- FFT size
  n_mels                (int,   32)    -- mel filterbank bands
  audio_topic           (str,   "/social/speech/audio_raw")
  silence_db            (float, -40.0)
  alarm_db_min          (float, -25.0)
  alarm_zcr_min         (float, 0.12)
  alarm_high_ratio_min  (float, 0.35)
  speech_zcr_min        (float, 0.02)
  speech_zcr_max        (float, 0.25)
  speech_flatness_max   (float, 0.35)
  music_zcr_max         (float, 0.08)
  music_flatness_max    (float, 0.25)
  crowd_zcr_min         (float, 0.10)
  crowd_flatness_min    (float, 0.35)
"""

from __future__ import annotations

import math
import struct
import threading
from typing import Dict, List, Optional

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile
from std_msgs.msg import String, UInt8MultiArray

# numpy used only in DSP helpers — the Jetson always has it
try:
    import numpy as np
    _NUMPY = True
except ImportError:
    _NUMPY = False

INT16_MAX = 32768.0
LABELS = ("silence", "speech", "music", "crowd", "outdoor", "alarm")


# ── PCM helpers ───────────────────────────────────────────────────────────────

def pcm16_bytes_to_float32(data: bytes) -> List[float]:
    """PCM-16 LE bytes → float32 list in [-1.0, 1.0]."""
    n = len(data) // 2
    if n == 0:
        return []
    return [s / INT16_MAX for s in struct.unpack(f"<{n}h", data[: n * 2])]


# ── Mel DSP (numpy path) ──────────────────────────────────────────────────────

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


def mel_to_hz(mel: float) -> float:
    return 700.0 * (10.0 ** (mel / 2595.0) - 1.0)


def build_mel_filterbank(sr: int, n_fft: int, n_mels: int,
                         fmin: float = 0.0, fmax: Optional[float] = None):
    """Return (n_mels, n_fft//2+1) numpy filterbank matrix."""
    import numpy as np
    if fmax is None:
        fmax = sr / 2.0
    n_freqs = n_fft // 2 + 1
    mel_min = hz_to_mel(fmin)
    mel_max = hz_to_mel(fmax)
    mel_pts = np.linspace(mel_min, mel_max, n_mels + 2)
    hz_pts  = np.array([mel_to_hz(m) for m in mel_pts])
    bin_pts = np.floor((n_fft + 1) * hz_pts / sr).astype(int)
    fb = np.zeros((n_mels, n_freqs))
    for m in range(n_mels):
        lo, ctr, hi = bin_pts[m], bin_pts[m + 1], bin_pts[m + 2]
        for k in range(lo, min(ctr, n_freqs)):
            if ctr != lo:
                fb[m, k] = (k - lo) / (ctr - lo)
        for k in range(ctr, min(hi, n_freqs)):
            if hi != ctr:
                fb[m, k] = (hi - k) / (hi - ctr)
    return fb


def compute_mel_spectrogram(samples: List[float], sr: int,
                            n_fft: int = 512, n_mels: int = 32,
                            hop_length: int = 256):
    """Return (n_mels, n_frames) log-mel spectrogram (numpy array)."""
    import numpy as np
    x = np.array(samples, dtype=np.float32)
    fb = build_mel_filterbank(sr, n_fft, n_mels)
    window = np.hanning(n_fft)
    frames = []
    for start in range(0, len(x) - n_fft + 1, hop_length):
        frame = x[start : start + n_fft] * window
        spec  = np.abs(np.fft.rfft(frame)) ** 2
        mel   = fb @ spec
        frames.append(mel)
    if not frames:
        return np.zeros((n_mels, 1), dtype=np.float32)
    return np.column_stack(frames).astype(np.float32)


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

def extract_features(samples: List[float], sr: int,
                     n_fft: int = 512, n_mels: int = 32) -> Dict[str, float]:
    """Extract scalar features from a raw audio window."""
    import numpy as np

    n = len(samples)
    if n == 0:
        return {k: 0.0 for k in
                ("energy_db", "zcr", "mel_centroid", "mel_flatness",
                 "low_ratio", "high_ratio")}

    # Energy
    rms = math.sqrt(sum(s * s for s in samples) / n) if n else 0.0
    energy_db = 20.0 * math.log10(max(rms, 1e-10))

    # ZCR across 30 ms frames
    chunk = max(1, int(sr * 0.030))
    zcr_vals = []
    for i in range(0, n - chunk + 1, chunk):
        seg = samples[i : i + chunk]
        crossings = sum(1 for j in range(1, len(seg))
                        if seg[j - 1] * seg[j] < 0)
        zcr_vals.append(crossings / max(len(seg) - 1, 1))
    zcr = sum(zcr_vals) / len(zcr_vals) if zcr_vals else 0.0

    # Mel spectrogram features
    mel_spec = compute_mel_spectrogram(samples, sr, n_fft, n_mels)
    mel_mean = mel_spec.mean(axis=1)  # (n_mels,) mean energy per band

    total = float(mel_mean.sum()) if mel_mean.sum() > 0 else 1e-10
    indices = np.arange(n_mels, dtype=np.float32)
    mel_centroid = float((indices * mel_mean).sum()) / (n_mels * total / total) / n_mels

    # Spectral flatness: geometric mean / arithmetic mean
    eps = 1e-10
    mel_pos = np.clip(mel_mean, eps, None)
    geo_mean = float(np.exp(np.log(mel_pos).mean()))
    arith_mean = float(mel_pos.mean())
    mel_flatness = min(geo_mean / max(arith_mean, eps), 1.0)

    # Band ratios
    third = max(1, n_mels // 3)
    low_energy  = float(mel_mean[:third].sum())
    high_energy = float(mel_mean[-third:].sum())
    low_ratio   = low_energy  / max(total, eps)
    high_ratio  = high_energy / max(total, eps)

    return {
        "energy_db":    energy_db,
        "zcr":          zcr,
        "mel_centroid": mel_centroid,
        "mel_flatness": mel_flatness,
        "low_ratio":    low_ratio,
        "high_ratio":   high_ratio,
    }


# ── Classifier ────────────────────────────────────────────────────────────────

def classify(features: Dict[str, float],
             silence_db: float = -40.0,
             alarm_db_min: float = -25.0,
             alarm_zcr_min: float = 0.12,
             alarm_high_ratio_min: float = 0.35,
             speech_zcr_min: float = 0.02,
             speech_zcr_max: float = 0.25,
             speech_flatness_max: float = 0.35,
             music_zcr_max: float = 0.08,
             music_flatness_max: float = 0.25,
             crowd_zcr_min: float = 0.10,
             crowd_flatness_min: float = 0.35) -> str:
    """Priority-ordered rule cascade. Returns a label from LABELS."""
    e   = features["energy_db"]
    zcr = features["zcr"]
    fl  = features["mel_flatness"]
    hi  = features["high_ratio"]

    if e < silence_db:
        return "silence"
    if (e >= alarm_db_min
            and zcr >= alarm_zcr_min
            and hi >= alarm_high_ratio_min):
        return "alarm"
    if zcr < music_zcr_max and fl < music_flatness_max:
        return "music"
    if (speech_zcr_min <= zcr <= speech_zcr_max
            and fl < speech_flatness_max):
        return "speech"
    if zcr >= crowd_zcr_min and fl >= crowd_flatness_min:
        return "crowd"
    return "outdoor"


# ── Audio accumulation buffer ─────────────────────────────────────────────────

class AudioBuffer:
    """Thread-safe ring buffer; yields a window of samples when full."""

    def __init__(self, window_samples: int) -> None:
        self._target = window_samples
        self._buf: List[float] = []
        self._lock = threading.Lock()

    def push(self, samples: List[float]) -> Optional[List[float]]:
        """Append samples. Returns a complete window (and resets) when full."""
        with self._lock:
            self._buf.extend(samples)
            if len(self._buf) >= self._target:
                window = self._buf[: self._target]
                self._buf = self._buf[self._target :]
                return window
        return None

    def clear(self) -> None:
        with self._lock:
            self._buf.clear()


# ── ROS2 node ─────────────────────────────────────────────────────────────────

class AmbientSoundNode(Node):
    """Classifies ambient sound from raw audio and publishes label at 1 Hz."""

    def __init__(self) -> None:
        super().__init__("ambient_sound_node")

        self.declare_parameter("sample_rate",          16000)
        self.declare_parameter("window_s",             1.0)
        self.declare_parameter("n_fft",                512)
        self.declare_parameter("n_mels",               32)
        self.declare_parameter("audio_topic",          "/social/speech/audio_raw")
        # Classifier thresholds
        self.declare_parameter("silence_db",           -40.0)
        self.declare_parameter("alarm_db_min",         -25.0)
        self.declare_parameter("alarm_zcr_min",        0.12)
        self.declare_parameter("alarm_high_ratio_min", 0.35)
        self.declare_parameter("speech_zcr_min",       0.02)
        self.declare_parameter("speech_zcr_max",       0.25)
        self.declare_parameter("speech_flatness_max",  0.35)
        self.declare_parameter("music_zcr_max",        0.08)
        self.declare_parameter("music_flatness_max",   0.25)
        self.declare_parameter("crowd_zcr_min",        0.10)
        self.declare_parameter("crowd_flatness_min",   0.35)

        self._sr      = self.get_parameter("sample_rate").value
        self._n_fft   = self.get_parameter("n_fft").value
        self._n_mels  = self.get_parameter("n_mels").value
        window_s      = self.get_parameter("window_s").value
        audio_topic   = self.get_parameter("audio_topic").value

        self._thresholds = {
            k: self.get_parameter(k).value for k in (
                "silence_db", "alarm_db_min", "alarm_zcr_min",
                "alarm_high_ratio_min", "speech_zcr_min", "speech_zcr_max",
                "speech_flatness_max", "music_zcr_max", "music_flatness_max",
                "crowd_zcr_min", "crowd_flatness_min",
            )
        }

        self._buffer = AudioBuffer(int(self._sr * window_s))
        self._last_label = "silence"

        qos = QoSProfile(depth=10)
        self._pub = self.create_publisher(String, "/saltybot/ambient_sound", qos)
        self._audio_sub = self.create_subscription(
            UInt8MultiArray, audio_topic, self._on_audio, qos
        )

        if not _NUMPY:
            self.get_logger().warn(
                "numpy not available — mel features disabled, classifying by energy only"
            )

        self.get_logger().info(
            f"AmbientSoundNode ready "
            f"(sr={self._sr}, window={window_s}s, n_mels={self._n_mels})"
        )

    def _on_audio(self, msg: UInt8MultiArray) -> None:
        samples = pcm16_bytes_to_float32(bytes(msg.data))
        if not samples:
            return
        window = self._buffer.push(samples)
        if window is not None:
            self._classify_and_publish(window)

    def _classify_and_publish(self, samples: List[float]) -> None:
        try:
            if _NUMPY:
                feats = extract_features(samples, self._sr, self._n_fft, self._n_mels)
            else:
                # Numpy-free fallback: energy-only
                rms = math.sqrt(sum(s * s for s in samples) / len(samples))
                e_db = 20.0 * math.log10(max(rms, 1e-10))
                feats = {
                    "energy_db": e_db, "zcr": 0.05,
                    "mel_centroid": 0.5, "mel_flatness": 0.2,
                    "low_ratio": 0.4, "high_ratio": 0.2,
                }
            label = classify(feats, **self._thresholds)
        except Exception as exc:
            self.get_logger().error(f"Classification error: {exc}")
            label = self._last_label

        if label != self._last_label:
            self.get_logger().info(
                f"Ambient sound: {self._last_label} -> {label}"
            )
        self._last_label = label

        msg = String()
        msg.data = label
        self._pub.publish(msg)


def main(args: Optional[list] = None) -> None:
    rclpy.init(args=args)
    node = AmbientSoundNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()
