"""vad_node.py — Energy + zero-crossing rate voice activity detection.
Issue #242

Subscribes to raw PCM-16 audio on /social/speech/audio_raw, computes
per-chunk RMS energy (dBFS) and zero-crossing rate (ZCR), applies
onset/offset hysteresis, and publishes:

  /social/speech/is_speaking  (std_msgs/Bool)    — VAD decision
  /social/speech/energy       (std_msgs/Float32) — linear RMS [0..1]

The combined decision rule:
  active = energy_db >= rms_threshold_db  AND  zcr in [zcr_min, zcr_max]

ZCR bands for 16 kHz audio (typical):
  Silence / low-freq rumble : ZCR < 0.01
  Voiced speech             : ZCR  0.01–0.20
  Unvoiced / sibilants      : ZCR  0.20–0.40
  High-freq noise           : ZCR > 0.40

Parameters:
  sample_rate       (int,   16000) — expected sample rate of incoming audio
  rms_threshold_db  (float, -35.0) — energy gate (dBFS); below = silence
  zcr_min           (float,  0.01) — ZCR lower bound; below = rumble/DC
  zcr_max           (float,  0.40) — ZCR upper bound; above = noise
  onset_frames      (int,    2)    — consecutive active frames to set is_speaking
  offset_frames     (int,    8)    — consecutive silent frames to clear is_speaking
  audio_topic       (str, "/social/speech/audio_raw") — input PCM-16 topic
"""

from __future__ import annotations

import math
import struct
from typing import Optional

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

INT16_MAX = 32768.0


# ── Pure signal helpers (no ROS) ──────────────────────────────────────────────

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


def rms_linear(samples: list) -> float:
    """RMS amplitude in [0.0, 1.0]. Returns 0.0 for empty / silent input."""
    if not samples:
        return 0.0
    mean_sq = sum(s * s for s in samples) / len(samples)
    return math.sqrt(mean_sq) if mean_sq > 0.0 else 0.0


def rms_db(samples: list) -> float:
    """RMS energy in dBFS. Returns -96.0 for silence."""
    rms = rms_linear(samples)
    return 20.0 * math.log10(max(rms, 1e-10))


def zero_crossing_rate(samples: list) -> float:
    """Zero-crossing rate: fraction of consecutive pairs with opposite signs.

    Returns a value in [0.0, 1.0].  Silence returns 0.0.
    """
    n = len(samples)
    if n < 2:
        return 0.0
    crossings = sum(
        1 for i in range(1, n) if samples[i - 1] * samples[i] < 0
    )
    return crossings / (n - 1)


# ── Hysteresis state machine ──────────────────────────────────────────────────

class VadStateMachine:
    """Onset/offset hysteresis on a per-frame boolean signal."""

    def __init__(self, onset_frames: int = 2, offset_frames: int = 8) -> None:
        self.onset_frames  = onset_frames
        self.offset_frames = offset_frames
        self._above = 0
        self._below = 0
        self._active = False

    def update(self, raw_active: bool) -> bool:
        if raw_active:
            self._above += 1
            self._below = 0
            if self._above >= self.onset_frames:
                self._active = True
        else:
            self._below += 1
            self._above = 0
            if self._below >= self.offset_frames:
                self._active = False
        return self._active

    def reset(self) -> None:
        self._above = 0
        self._below = 0
        self._active = False

    @property
    def is_active(self) -> bool:
        return self._active


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

class VadNode(Node):
    """Energy + ZCR voice activity detector — subscribes to raw audio."""

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

        self.declare_parameter("sample_rate",      16000)
        self.declare_parameter("rms_threshold_db", -35.0)
        self.declare_parameter("zcr_min",          0.01)
        self.declare_parameter("zcr_max",          0.40)
        self.declare_parameter("onset_frames",     2)
        self.declare_parameter("offset_frames",    8)
        self.declare_parameter("audio_topic",      "/social/speech/audio_raw")

        self._sample_rate  = self.get_parameter("sample_rate").value
        self._rms_thresh   = self.get_parameter("rms_threshold_db").value
        self._zcr_min      = self.get_parameter("zcr_min").value
        self._zcr_max      = self.get_parameter("zcr_max").value
        audio_topic        = self.get_parameter("audio_topic").value

        self._sm = VadStateMachine(
            onset_frames  = self.get_parameter("onset_frames").value,
            offset_frames = self.get_parameter("offset_frames").value,
        )

        qos = QoSProfile(depth=10)
        self._speaking_pub = self.create_publisher(Bool,    "/social/speech/is_speaking", qos)
        self._energy_pub   = self.create_publisher(Float32, "/social/speech/energy",      qos)

        self._audio_sub = self.create_subscription(
            UInt8MultiArray, audio_topic, self._on_audio, qos
        )

        self.get_logger().info(
            f"VadNode ready (rms_thresh={self._rms_thresh} dBFS, "
            f"zcr=[{self._zcr_min},{self._zcr_max}], "
            f"topic={audio_topic})"
        )

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

        energy_lin = rms_linear(samples)
        energy_db  = rms_db(samples)
        zcr        = zero_crossing_rate(samples)

        raw_active = (
            energy_db  >= self._rms_thresh
            and self._zcr_min <= zcr <= self._zcr_max
        )
        is_speaking = self._sm.update(raw_active)

        bool_msg = Bool()
        bool_msg.data = is_speaking
        self._speaking_pub.publish(bool_msg)

        energy_msg = Float32()
        energy_msg.data = float(energy_lin)
        self._energy_pub.publish(energy_msg)


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