"""wake_word_node.py — Audio wake-word detector ('hey salty').
Issue #320

Subscribes to raw PCM-16 audio on /social/speech/audio_raw, maintains a
sliding window buffer, and on each hop tick computes log-mel spectrogram
features of the most recent window.  If energy is above the gate threshold
AND cosine similarity to the stored template is above match_threshold the
detection fires: Bool(True) is published on /saltybot/wake_word_detected.

Detection is one-shot (only True is published) and guarded by a cooldown
so rapid re-fires are suppressed.  When no template is loaded (template_path
is empty) the node stays passive — energy gating is applied but no match is
attempted.

Audio format expected
─────────────────────
  UInt8MultiArray, same feed as vad_node:
  raw PCM-16 little-endian mono at ``sample_rate`` Hz.

Subscriptions
─────────────
  /social/speech/audio_raw   std_msgs/UInt8MultiArray  — raw PCM-16 chunks

Publications
────────────
  /saltybot/wake_word_detected   std_msgs/Bool  — True on each detection event

Parameters
──────────
  audio_topic       (str,   "/social/speech/audio_raw")
  output_topic      (str,   "/saltybot/wake_word_detected")
  sample_rate       (int,   16000)  sample rate of incoming audio (Hz)
  window_s          (float, 1.5)   detection window duration (s)
  hop_s             (float, 0.1)   detection timer period (s)
  energy_threshold  (float, 0.02)  RMS gate; below this → skip matching
  match_threshold   (float, 0.82)  cosine-similarity gate; above → detect
  cooldown_s        (float, 2.0)   minimum gap between successive detections
  template_path     (str,   "")    path to .npy template file; "" = passive
  n_fft             (int,   512)   FFT size for mel spectrogram
  n_mels            (int,   40)    number of mel filterbank bands
"""

from __future__ import annotations

import math
import struct
import threading
import time
from collections import deque
from typing import Dict, Optional, Tuple

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

try:
    import numpy as np
    _NP = True
except ImportError:                          # pragma: no cover
    _NP = False

INT16_MAX = 32768.0


# ── Pure DSP helpers (no ROS, numpy only) ──────────────────────────────────────

def pcm16_to_float(data: bytes) -> "np.ndarray":
    """Decode PCM-16 LE bytes → float32 ndarray in [-1.0, 1.0]."""
    n = len(data) // 2
    if n == 0:
        return np.zeros(0, dtype=np.float32)
    samples = struct.unpack(f"<{n}h", data[:n * 2])
    return np.array(samples, dtype=np.float32) / INT16_MAX


def mel_filterbank(sr: int, n_fft: int, n_mels: int,
                   fmin: float = 80.0, fmax: Optional[float] = None) -> "np.ndarray":
    """Build a triangular mel filterbank matrix [n_mels, n_fft//2+1]."""
    if fmax is None:
        fmax = sr / 2.0

    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)

    mel_lo = hz_to_mel(fmin)
    mel_hi = hz_to_mel(fmax)
    mel_pts = np.linspace(mel_lo, mel_hi, n_mels + 2)
    hz_pts  = np.array([mel_to_hz(m) for m in mel_pts])
    freqs   = np.fft.rfftfreq(n_fft, d=1.0 / sr)

    fb = np.zeros((n_mels, len(freqs)), dtype=np.float32)
    for m in range(n_mels):
        lo, center, hi = hz_pts[m], hz_pts[m + 1], hz_pts[m + 2]
        for k, f in enumerate(freqs):
            if lo <= f < center and center > lo:
                fb[m, k] = (f - lo) / (center - lo)
            elif center <= f <= hi and hi > center:
                fb[m, k] = (hi - f) / (hi - center)
    return fb


def compute_log_mel(samples: "np.ndarray", sr: int,
                    n_fft: int = 512, n_mels: int = 40,
                    hop: int = 256) -> "np.ndarray":
    """Return log-mel spectrogram [n_mels, T] of *samples* (float32 [-1,1])."""
    n = len(samples)
    window = np.hanning(n_fft).astype(np.float32)
    frames = []
    for start in range(0, max(n - n_fft + 1, 1), hop):
        chunk = samples[start:start + n_fft]
        if len(chunk) < n_fft:
            chunk = np.pad(chunk, (0, n_fft - len(chunk)))
        power = np.abs(np.fft.rfft(chunk * window)) ** 2
        frames.append(power)
    frames_arr = np.array(frames, dtype=np.float32).T   # [bins, T]
    fb  = mel_filterbank(sr, n_fft, n_mels)
    mel = fb @ frames_arr                               # [n_mels, T]
    mel = np.where(mel > 1e-10, mel, 1e-10)
    return np.log(mel)


def cosine_sim(a: "np.ndarray", b: "np.ndarray") -> float:
    """Cosine similarity between two arrays, matched by minimum length."""
    af = a.flatten().astype(np.float64)
    bf = b.flatten().astype(np.float64)
    min_len = min(len(af), len(bf))
    if min_len == 0:
        return 0.0
    af = af[:min_len]
    bf = bf[:min_len]
    denom = float(np.linalg.norm(af)) * float(np.linalg.norm(bf))
    if denom < 1e-12:
        return 0.0
    return float(np.dot(af, bf) / denom)


def rms(samples: "np.ndarray") -> float:
    """RMS amplitude of a float sample array."""
    if len(samples) == 0:
        return 0.0
    return float(np.sqrt(np.mean(samples.astype(np.float64) ** 2)))


# ── Ring buffer ────────────────────────────────────────────────────────────────

class AudioRingBuffer:
    """Lock-free sliding window for raw float audio samples."""

    def __init__(self, max_samples: int) -> None:
        self._buf: deque = deque(maxlen=max_samples)

    def push(self, samples: "np.ndarray") -> None:
        self._buf.extend(samples.tolist())

    def get_window(self, n_samples: int) -> Optional["np.ndarray"]:
        """Return last n_samples as float32 array, or None if buffer too short."""
        if len(self._buf) < n_samples:
            return None
        return np.array(list(self._buf)[-n_samples:], dtype=np.float32)

    def __len__(self) -> int:
        return len(self._buf)


# ── Detector ───────────────────────────────────────────────────────────────────

class WakeWordDetector:
    """Energy-gated cosine-similarity wake-word detector.

    Args:
        template:          Log-mel feature array of the wake word, or None
                           (passive — never fires when None).
        energy_threshold:  Minimum RMS to proceed to feature matching.
        match_threshold:   Minimum cosine similarity to fire.
        sample_rate:       Expected audio sample rate (Hz).
        n_fft:             FFT size for mel computation.
        n_mels:            Number of mel bands.
    """

    def __init__(self,
                 template: Optional["np.ndarray"],
                 energy_threshold: float = 0.02,
                 match_threshold:  float = 0.82,
                 sample_rate: int   = 16000,
                 n_fft: int         = 512,
                 n_mels: int        = 40) -> None:
        self._template    = template
        self._energy_thr  = energy_threshold
        self._match_thr   = match_threshold
        self._sr          = sample_rate
        self._n_fft       = n_fft
        self._n_mels      = n_mels

    # ------------------------------------------------------------------
    def detect(self, samples: "np.ndarray") -> Tuple[bool, float, float]:
        """Run detection on a window of float samples.

        Returns
        -------
        (detected, rms_value, similarity)
        """
        energy = rms(samples)
        if energy < self._energy_thr:
            return False, energy, 0.0
        if self._template is None:
            return False, energy, 0.0

        hop = max(1, self._n_fft // 2)
        feats = compute_log_mel(samples, self._sr, self._n_fft, self._n_mels, hop)
        sim   = cosine_sim(feats, self._template)
        return sim >= self._match_thr, energy, sim

    @property
    def has_template(self) -> bool:
        return self._template is not None


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

class WakeWordNode(Node):
    """ROS2 node: 'hey salty' wake-word detection via energy + template matching."""

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

        self.declare_parameter("audio_topic",      "/social/speech/audio_raw")
        self.declare_parameter("output_topic",     "/saltybot/wake_word_detected")
        self.declare_parameter("sample_rate",      16000)
        self.declare_parameter("window_s",         1.5)
        self.declare_parameter("hop_s",            0.1)
        self.declare_parameter("energy_threshold", 0.02)
        self.declare_parameter("match_threshold",  0.82)
        self.declare_parameter("cooldown_s",       2.0)
        self.declare_parameter("template_path",    "")
        self.declare_parameter("n_fft",            512)
        self.declare_parameter("n_mels",           40)

        audio_topic  = self.get_parameter("audio_topic").value
        output_topic = self.get_parameter("output_topic").value
        self._sr     = int(self.get_parameter("sample_rate").value)
        self._win_s  = float(self.get_parameter("window_s").value)
        hop_s        = float(self.get_parameter("hop_s").value)
        energy_thr   = float(self.get_parameter("energy_threshold").value)
        match_thr    = float(self.get_parameter("match_threshold").value)
        self._cool_s = float(self.get_parameter("cooldown_s").value)
        tmpl_path    = str(self.get_parameter("template_path").value)
        n_fft        = int(self.get_parameter("n_fft").value)
        n_mels       = int(self.get_parameter("n_mels").value)

        # ── Load template ──────────────────────────────────────────────
        template: Optional["np.ndarray"] = None
        if tmpl_path and _NP:
            try:
                template = np.load(tmpl_path)
                self.get_logger().info(
                    f"WakeWord: loaded template {tmpl_path} shape={template.shape}"
                )
            except Exception as exc:
                self.get_logger().warn(
                    f"WakeWord: could not load template '{tmpl_path}': {exc} — passive mode"
                )

        # ── Ring buffer ────────────────────────────────────────────────
        max_samples     = int(self._sr * self._win_s * 4)   # 4× headroom
        self._win_n     = int(self._sr * self._win_s)
        self._buf       = AudioRingBuffer(max_samples)

        # ── Detector ───────────────────────────────────────────────────
        self._detector  = WakeWordDetector(template, energy_thr, match_thr,
                                           self._sr, n_fft, n_mels)

        # ── State ──────────────────────────────────────────────────────
        self._last_det_t: float = 0.0
        self._lock = threading.Lock()

        # ── ROS ────────────────────────────────────────────────────────
        qos = QoSProfile(depth=10)
        self._pub = self.create_publisher(Bool, output_topic, qos)
        self._sub = self.create_subscription(
            UInt8MultiArray, audio_topic, self._on_audio, qos
        )
        self._timer = self.create_timer(hop_s, self._detection_cb)

        tmpl_status = (f"template={tmpl_path}" if template is not None
                       else "passive (no template)")
        self.get_logger().info(
            f"WakeWordNode ready — {tmpl_status}, "
            f"window={self._win_s}s, hop={hop_s}s, "
            f"energy_thr={energy_thr}, match_thr={match_thr}"
        )

    # ── Subscription ───────────────────────────────────────────────────

    def _on_audio(self, msg) -> None:
        if not _NP:
            return
        try:
            raw = bytes(msg.data)
            samples = pcm16_to_float(raw)
        except Exception as exc:
            self.get_logger().warn(f"WakeWord: audio decode error: {exc}")
            return
        with self._lock:
            self._buf.push(samples)

    # ── Detection timer ────────────────────────────────────────────────

    def _detection_cb(self) -> None:
        if not _NP:
            return
        now = time.monotonic()
        with self._lock:
            window = self._buf.get_window(self._win_n)
        if window is None:
            return

        detected, energy, sim = self._detector.detect(window)

        if detected and (now - self._last_det_t) >= self._cool_s:
            self._last_det_t = now
            self.get_logger().info(
                f"WakeWord: 'hey salty' detected "
                f"(rms={energy:.4f}, sim={sim:.3f})"
            )
            out = Bool()
            out.data = True
            self._pub.publish(out)


def main(args=None) -> None:
    rclpy.init(args=args)
    node = WakeWordNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()
