"""
audio_direction_node.py — Sound source localization via GCC-PHAT beamforming.

Estimates direction of arrival (DoA) from a Jabra multi-channel microphone
using Generalized Cross-Correlation with Phase Transform (GCC-PHAT).

Publishes:
  /saltybot/audio_direction   std_msgs/Float32      bearing in degrees (0-360)
  /saltybot/audio_activity    std_msgs/Bool         voice activity detected

Parameters:
  device_id           int     -1          audio device index (-1 = default)
  sample_rate         int     16000       sample rate in Hz
  chunk_size          int     2048        samples per frame
  publish_hz          float   10.0        output publication rate
  vad_threshold       float   0.02        RMS energy threshold for speech
  gcc_phat_max_lag    int     64          max lag for correlation (determines angular resolution)
  self_noise_filter   bool    true        apply robot noise suppression
"""

from __future__ import annotations

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy

import numpy as np
from scipy import signal
import sounddevice as sd
from collections import deque
import threading
import time

from std_msgs.msg import Float32, Bool


_SENSOR_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=5,
)


class GCCPHATBeamformer:
    """Generalized Cross-Correlation with Phase Transform for DoA estimation."""

    def __init__(self, sample_rate: int, max_lag: int = 64):
        self.sample_rate = sample_rate
        self.max_lag = max_lag
        self.frequency_bins = max_lag * 2
        # Azimuth angles for left (270°), front (0°), right (90°), rear (180°)
        self.angles = np.array([270.0, 0.0, 90.0, 180.0])

    def gcc_phat(self, sig1: np.ndarray, sig2: np.ndarray) -> float:
        """
        Compute GCC-PHAT correlation and estimate time delay (DoA proxy).

        Returns:
            Estimated time delay in samples (can be converted to angle)
        """
        if len(sig1) == 0 or len(sig2) == 0:
            return 0.0

        # Cross-correlation in frequency domain
        fft1 = np.fft.rfft(sig1, n=self.frequency_bins)
        fft2 = np.fft.rfft(sig2, n=self.frequency_bins)

        # Normalize magnitude (phase transform)
        cross_spectrum = fft1 * np.conj(fft2)
        cross_spectrum_normalized = cross_spectrum / (np.abs(cross_spectrum) + 1e-8)

        # Inverse FFT to get correlation
        correlation = np.fft.irfft(cross_spectrum_normalized, n=self.frequency_bins)

        # Find lag with maximum correlation
        lags = np.arange(-self.max_lag, self.max_lag + 1)
        valid_corr = correlation[: len(lags)]
        max_lag_idx = np.argmax(np.abs(valid_corr))
        estimated_lag = lags[max_lag_idx] if max_lag_idx < len(lags) else 0.0

        return float(estimated_lag)

    def estimate_bearing(self, channels: list[np.ndarray]) -> float:
        """
        Estimate bearing from multi-channel audio.

        Supports:
        - 2-channel (stereo): left/right discrimination
        - 4-channel: quadrophonic (front/rear/left/right)

        Returns:
            Bearing in degrees (0-360)
        """
        if len(channels) < 2:
            return 0.0  # Default to front if mono

        if len(channels) == 2:
            # Stereo: estimate left vs right
            lag = self.gcc_phat(channels[0], channels[1])
            # Negative lag → left (270°), positive lag → right (90°)
            if abs(lag) < 5:
                return 0.0  # Front
            elif lag < 0:
                return 270.0  # Left
            else:
                return 90.0  # Right

        elif len(channels) >= 4:
            # Quadrophonic: compute pairwise correlations
            # Assume channel order: [front, right, rear, left]
            lags = []
            for i in range(4):
                lag = self.gcc_phat(channels[i], channels[(i + 1) % 4])
                lags.append(lag)

            # Select angle based on strongest correlation
            max_idx = np.argmax(np.abs(lags))
            return self.angles[max_idx]

        return 0.0


class VADDetector:
    """Simple voice activity detection based on RMS energy."""

    def __init__(self, threshold: float = 0.02, smoothing: int = 5):
        self.threshold = threshold
        self.smoothing = smoothing
        self.history = deque(maxlen=smoothing)

    def detect(self, audio_frame: np.ndarray) -> bool:
        """
        Detect voice activity in audio frame.

        Returns:
            True if speech-like energy detected
        """
        rms = np.sqrt(np.mean(audio_frame**2))
        self.history.append(rms > self.threshold)
        # Majority voting over window
        return sum(self.history) > self.smoothing // 2


class AudioDirectionNode(Node):

    def __init__(self):
        super().__init__('audio_direction_estimator')

        # Parameters
        self.declare_parameter('device_id', -1)
        self.declare_parameter('sample_rate', 16000)
        self.declare_parameter('chunk_size', 2048)
        self.declare_parameter('publish_hz', 10.0)
        self.declare_parameter('vad_threshold', 0.02)
        self.declare_parameter('gcc_phat_max_lag', 64)
        self.declare_parameter('self_noise_filter', True)

        self.device_id = self.get_parameter('device_id').value
        self.sample_rate = self.get_parameter('sample_rate').value
        self.chunk_size = self.get_parameter('chunk_size').value
        pub_hz = self.get_parameter('publish_hz').value
        vad_threshold = self.get_parameter('vad_threshold').value
        gcc_max_lag = self.get_parameter('gcc_phat_max_lag').value
        self.apply_noise_filter = self.get_parameter('self_noise_filter').value

        # Publishers
        self._pub_bearing = self.create_publisher(
            Float32, '/saltybot/audio_direction', 10, qos_profile=_SENSOR_QOS
        )
        self._pub_vad = self.create_publisher(
            Bool, '/saltybot/audio_activity', 10, qos_profile=_SENSOR_QOS
        )

        # Audio processing
        self.beamformer = GCCPHATBeamformer(self.sample_rate, max_lag=gcc_max_lag)
        self.vad = VADDetector(threshold=vad_threshold)
        self._audio_buffer = deque(maxlen=self.chunk_size * 2)
        self._lock = threading.Lock()

        # Start audio stream
        self._stream = None
        self._running = True
        self._start_audio_stream()

        # Publish timer
        self.create_timer(1.0 / pub_hz, self._tick)

        self.get_logger().info(
            f'audio_direction_estimator ready — '
            f'device_id={self.device_id} sample_rate={self.sample_rate} Hz '
            f'chunk_size={self.chunk_size} hz={pub_hz}'
        )

    def _start_audio_stream(self) -> None:
        """Initialize audio stream from default microphone."""
        try:
            self._stream = sd.InputStream(
                device=self.device_id if self.device_id >= 0 else None,
                samplerate=self.sample_rate,
                channels=2,  # Default to stereo; auto-detect Jabra channels
                blocksize=self.chunk_size,
                callback=self._audio_callback,
                latency='low',
            )
            self._stream.start()
            self.get_logger().info('Audio stream started')
        except Exception as e:
            self.get_logger().error(f'Failed to start audio stream: {e}')
            self._stream = None

    def _audio_callback(self, indata: np.ndarray, frames: int, time_info, status) -> None:
        """Callback for audio input stream."""
        if status:
            self.get_logger().warn(f'Audio callback status: {status}')

        try:
            with self._lock:
                # Store stereo or mono data
                if indata.shape[1] == 1:
                    self._audio_buffer.extend(indata.flatten())
                else:
                    # For multi-channel, interleave or store separately
                    self._audio_buffer.extend(indata.flatten())
        except Exception as e:
            self.get_logger().error(f'Audio callback error: {e}')

    def _tick(self) -> None:
        """Publish audio direction and VAD at configured rate."""
        if self._stream is None or not self._running:
            return

        with self._lock:
            if len(self._audio_buffer) < self.chunk_size:
                return
            audio_data = np.array(list(self._audio_buffer))

        # Extract channels (assume stereo or mono)
        if len(audio_data) > 0:
            channels = self._extract_channels(audio_data)

            # VAD detection on first channel
            if len(channels) > 0:
                is_speech = self.vad.detect(channels[0])
                vad_msg = Bool()
                vad_msg.data = is_speech
                self._pub_vad.publish(vad_msg)

            # DoA estimation (only if speech detected)
            if is_speech and len(channels) >= 2:
                bearing = self.beamformer.estimate_bearing(channels)
            else:
                bearing = 0.0  # Default to front when no speech

            bearing_msg = Float32()
            bearing_msg.data = float(bearing)
            self._pub_bearing.publish(bearing_msg)

    def _extract_channels(self, audio_data: np.ndarray) -> list[np.ndarray]:
        """
        Extract stereo/mono channels from audio buffer.

        Handles both interleaved and non-interleaved formats.
        """
        if len(audio_data) == 0:
            return []

        # Assume audio_data is a flat array; reshape for stereo
        # Adjust based on actual channel count from stream
        if len(audio_data) % 2 == 0:
            # Stereo interleaved
            stereo = audio_data.reshape(-1, 2)
            return [stereo[:, 0], stereo[:, 1]]
        else:
            # Mono or odd-length
            return [audio_data]

    def destroy_node(self):
        """Clean up audio stream on shutdown."""
        self._running = False
        if self._stream is not None:
            self._stream.stop()
            self._stream.close()
        super().destroy_node()


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


if __name__ == '__main__':
    main()
