Trixy3/trixy_core/ml/wakeword/audio_buffer.py

"""
Real-time audio buffer management for wakeword detection.

This module provides efficient audio buffering and streaming capabilities
optimized for real-time wakeword detection with minimal latency.
"""

import threading
import time
from collections import deque
from typing import Optional, Callable, List, Tuple, Union
from dataclasses import dataclass
import numpy as np
import torch


@dataclass
class AudioChunk:
    """Represents a chunk of audio data with metadata."""
    
    data: Union[np.ndarray, torch.Tensor]
    timestamp: float
    sample_rate: int
    channels: int = 1
    chunk_id: int = 0
    
    def __post_init__(self):
        """Validate audio chunk parameters."""
        if isinstance(self.data, torch.Tensor):
            self.data = self.data.detach().cpu().numpy()
        
        self.data = np.asarray(self.data, dtype=np.float32)
        
        if self.data.ndim == 1:
            # Mono audio
            self.channels = 1
        elif self.data.ndim == 2:
            # Multi-channel audio
            self.channels = self.data.shape[0] if self.data.shape[0] < self.data.shape[1] else self.data.shape[1]
        else:
            raise ValueError(f"Invalid audio data shape: {self.data.shape}")
    
    @property
    def duration(self) -> float:
        """Get duration of audio chunk in seconds."""
        return len(self.data) / self.sample_rate
    
    @property
    def samples(self) -> int:
        """Get number of samples in the chunk."""
        return len(self.data)
    
    def to_tensor(self, device: Optional[torch.device] = None) -> torch.Tensor:
        """Convert audio data to PyTorch tensor."""
        tensor = torch.from_numpy(self.data).float()
        if device is not None:
            tensor = tensor.to(device)
        return tensor


class AudioBuffer:
    """
    Base class for audio buffering with configurable size and overlap.
    """
    
    def __init__(self, buffer_duration: float, sample_rate: int,
                 overlap_ratio: float = 0.5, max_chunks: int = 1000):
        """
        Initialize audio buffer.
        
        Args:
            buffer_duration: Duration of buffer in seconds
            sample_rate: Audio sample rate
            overlap_ratio: Overlap ratio between chunks (0.0 to 1.0)
            max_chunks: Maximum number of chunks to keep in memory
        """
        self.buffer_duration = buffer_duration
        self.sample_rate = sample_rate
        self.overlap_ratio = overlap_ratio
        self.max_chunks = max_chunks
        
        self.buffer_samples = int(buffer_duration * sample_rate)
        self.overlap_samples = int(self.buffer_samples * overlap_ratio)
        self.stride_samples = self.buffer_samples - self.overlap_samples
        
        self._buffer = np.zeros(self.buffer_samples, dtype=np.float32)
        self._current_position = 0
        self._total_samples_written = 0
        self._chunk_counter = 0
        
        self._lock = threading.Lock()
        
    def add_audio(self, audio_data: Union[np.ndarray, torch.Tensor]) -> List[AudioChunk]:
        """
        Add audio data to buffer and return complete chunks.
        
        Args:
            audio_data: Audio samples to add
            
        Returns:
            List of complete audio chunks ready for processing
        """
        if isinstance(audio_data, torch.Tensor):
            audio_data = audio_data.detach().cpu().numpy()
        
        audio_data = np.asarray(audio_data, dtype=np.float32)
        
        if audio_data.ndim > 1:
            # Convert to mono if multi-channel
            audio_data = np.mean(audio_data, axis=0)
        
        chunks = []
        
        with self._lock:
            data_idx = 0
            
            while data_idx < len(audio_data):
                # Calculate how much we can write to current buffer position
                remaining_buffer = self.buffer_samples - self._current_position
                remaining_data = len(audio_data) - data_idx
                samples_to_write = min(remaining_buffer, remaining_data)
                
                # Write data to buffer
                end_pos = self._current_position + samples_to_write
                self._buffer[self._current_position:end_pos] = audio_data[data_idx:data_idx + samples_to_write]
                
                self._current_position += samples_to_write
                self._total_samples_written += samples_to_write
                data_idx += samples_to_write
                
                # Check if buffer is full
                if self._current_position >= self.buffer_samples:
                    # Create chunk from current buffer
                    chunk_data = self._buffer.copy()
                    chunk = AudioChunk(
                        data=chunk_data,
                        timestamp=time.time(),
                        sample_rate=self.sample_rate,
                        chunk_id=self._chunk_counter
                    )
                    chunks.append(chunk)
                    self._chunk_counter += 1
                    
                    # Shift buffer for overlap
                    if self.overlap_samples > 0:
                        self._buffer[:self.overlap_samples] = self._buffer[-self.overlap_samples:]
                        self._current_position = self.overlap_samples
                    else:
                        self._current_position = 0
        
        return chunks
    
    def get_current_buffer(self) -> Optional[AudioChunk]:
        """
        Get current buffer state (even if not full).
        
        Returns:
            Current audio chunk or None if buffer is empty
        """
        with self._lock:
            if self._current_position == 0:
                return None
            
            # Pad the current buffer to full size
            current_data = np.zeros(self.buffer_samples, dtype=np.float32)
            current_data[:self._current_position] = self._buffer[:self._current_position]
            
            return AudioChunk(
                data=current_data,
                timestamp=time.time(),
                sample_rate=self.sample_rate,
                chunk_id=self._chunk_counter
            )
    
    def clear(self):
        """Clear the buffer."""
        with self._lock:
            self._buffer.fill(0)
            self._current_position = 0
            self._total_samples_written = 0
    
    def get_stats(self) -> dict:
        """Get buffer statistics."""
        with self._lock:
            return {
                'buffer_duration': self.buffer_duration,
                'buffer_samples': self.buffer_samples,
                'current_position': self._current_position,
                'total_samples_written': self._total_samples_written,
                'total_chunks_created': self._chunk_counter,
                'overlap_ratio': self.overlap_ratio,
                'buffer_fill_ratio': self._current_position / self.buffer_samples
            }


class CircularAudioBuffer:
    """
    Circular audio buffer for continuous streaming with automatic chunk extraction.
    
    This buffer continuously accepts audio data and provides overlapping chunks
    for real-time processing.
    """
    
    def __init__(self, buffer_duration: float, chunk_duration: float,
                 sample_rate: int, overlap_ratio: float = 0.5,
                 chunk_callback: Optional[Callable[[AudioChunk], None]] = None):
        """
        Initialize circular audio buffer.
        
        Args:
            buffer_duration: Total buffer duration in seconds
            chunk_duration: Duration of each processing chunk in seconds
            sample_rate: Audio sample rate
            overlap_ratio: Overlap ratio between chunks
            chunk_callback: Optional callback for new chunks
        """
        self.buffer_duration = buffer_duration
        self.chunk_duration = chunk_duration
        self.sample_rate = sample_rate
        self.overlap_ratio = overlap_ratio
        self.chunk_callback = chunk_callback
        
        self.buffer_samples = int(buffer_duration * sample_rate)
        self.chunk_samples = int(chunk_duration * sample_rate)
        self.overlap_samples = int(self.chunk_samples * overlap_ratio)
        self.stride_samples = self.chunk_samples - self.overlap_samples
        
        self._buffer = np.zeros(self.buffer_samples, dtype=np.float32)
        self._write_pos = 0
        self._last_chunk_pos = 0
        self._chunk_counter = 0
        
        self._lock = threading.Lock()
        self._stats = {
            'total_samples_written': 0,
            'total_chunks_extracted': 0,
            'buffer_overruns': 0,
            'last_chunk_time': 0.0
        }
    
    def write_audio(self, audio_data: Union[np.ndarray, torch.Tensor]):
        """
        Write audio data to the circular buffer.
        
        Args:
            audio_data: Audio samples to write
        """
        if isinstance(audio_data, torch.Tensor):
            audio_data = audio_data.detach().cpu().numpy()
        
        audio_data = np.asarray(audio_data, dtype=np.float32)
        
        if audio_data.ndim > 1:
            audio_data = np.mean(audio_data, axis=0)
        
        with self._lock:
            data_len = len(audio_data)
            
            # Check for buffer overrun
            if data_len > self.buffer_samples:
                self._stats['buffer_overruns'] += 1
                # Take only the last part that fits
                audio_data = audio_data[-self.buffer_samples:]
                data_len = len(audio_data)
            
            # Write data to circular buffer
            if self._write_pos + data_len <= self.buffer_samples:
                # Simple case: no wrap-around
                self._buffer[self._write_pos:self._write_pos + data_len] = audio_data
            else:
                # Wrap-around case
                first_part = self.buffer_samples - self._write_pos
                self._buffer[self._write_pos:] = audio_data[:first_part]
                self._buffer[:data_len - first_part] = audio_data[first_part:]
            
            self._write_pos = (self._write_pos + data_len) % self.buffer_samples
            self._stats['total_samples_written'] += data_len
            
            # Check if we can extract new chunks
            self._extract_chunks()
    
    def _extract_chunks(self):
        """Extract ready chunks from the buffer."""
        while True:
            # Calculate available samples since last chunk
            if self._write_pos >= self._last_chunk_pos:
                available_samples = self._write_pos - self._last_chunk_pos
            else:
                available_samples = (self.buffer_samples - self._last_chunk_pos) + self._write_pos
            
            # Check if we have enough samples for a new chunk
            if available_samples < self.stride_samples:
                break
            
            # Extract chunk
            chunk_data = np.zeros(self.chunk_samples, dtype=np.float32)
            
            if self._last_chunk_pos + self.chunk_samples <= self.buffer_samples:
                # No wrap-around
                chunk_data = self._buffer[self._last_chunk_pos:self._last_chunk_pos + self.chunk_samples].copy()
            else:
                # Wrap-around
                first_part = self.buffer_samples - self._last_chunk_pos
                chunk_data[:first_part] = self._buffer[self._last_chunk_pos:]
                chunk_data[first_part:] = self._buffer[:self.chunk_samples - first_part]
            
            # Create chunk object
            chunk = AudioChunk(
                data=chunk_data,
                timestamp=time.time(),
                sample_rate=self.sample_rate,
                chunk_id=self._chunk_counter
            )
            
            self._chunk_counter += 1
            self._stats['total_chunks_extracted'] += 1
            self._stats['last_chunk_time'] = chunk.timestamp
            
            # Update position for next chunk
            self._last_chunk_pos = (self._last_chunk_pos + self.stride_samples) % self.buffer_samples
            
            # Call callback if provided
            if self.chunk_callback:
                try:
                    self.chunk_callback(chunk)
                except Exception as e:
                    # Don't let callback errors break the buffer
                    print(f"Error in chunk callback: {e}")
    
    def get_latest_chunk(self) -> Optional[AudioChunk]:
        """
        Get the latest available chunk without affecting the buffer state.
        
        Returns:
            Latest audio chunk or None if not enough data
        """
        with self._lock:
            if self._write_pos >= self._last_chunk_pos:
                available_samples = self._write_pos - self._last_chunk_pos
            else:
                available_samples = (self.buffer_samples - self._last_chunk_pos) + self._write_pos
            
            if available_samples < self.chunk_samples:
                return None
            
            # Extract chunk without updating position
            chunk_data = np.zeros(self.chunk_samples, dtype=np.float32)
            
            # Get the most recent chunk_samples
            start_pos = (self._write_pos - self.chunk_samples) % self.buffer_samples
            
            if start_pos + self.chunk_samples <= self.buffer_samples:
                chunk_data = self._buffer[start_pos:start_pos + self.chunk_samples].copy()
            else:
                first_part = self.buffer_samples - start_pos
                chunk_data[:first_part] = self._buffer[start_pos:]
                chunk_data[first_part:] = self._buffer[:self.chunk_samples - first_part]
            
            return AudioChunk(
                data=chunk_data,
                timestamp=time.time(),
                sample_rate=self.sample_rate,
                chunk_id=self._chunk_counter
            )
    
    def clear(self):
        """Clear the buffer."""
        with self._lock:
            self._buffer.fill(0)
            self._write_pos = 0
            self._last_chunk_pos = 0
            self._chunk_counter = 0
    
    def get_stats(self) -> dict:
        """Get buffer statistics."""
        with self._lock:
            stats = self._stats.copy()
            stats.update({
                'buffer_duration': self.buffer_duration,
                'chunk_duration': self.chunk_duration,
                'buffer_samples': self.buffer_samples,
                'chunk_samples': self.chunk_samples,
                'overlap_ratio': self.overlap_ratio,
                'write_position': self._write_pos,
                'last_chunk_position': self._last_chunk_pos,
                'buffer_utilization': (abs(self._write_pos - self._last_chunk_pos) / self.buffer_samples)
            })
            return stats


class StreamingAudioBuffer:
    """
    High-performance streaming audio buffer for real-time wakeword detection.
    
    Optimized for minimal latency and efficient memory usage.
    """
    
    def __init__(self, chunk_size: int, sample_rate: int,
                 max_chunks_in_memory: int = 10,
                 chunk_processor: Optional[Callable[[AudioChunk], None]] = None):
        """
        Initialize streaming audio buffer.
        
        Args:
            chunk_size: Size of each audio chunk in samples
            sample_rate: Audio sample rate
            max_chunks_in_memory: Maximum chunks to keep in memory
            chunk_processor: Optional processor for chunks
        """
        self.chunk_size = chunk_size
        self.sample_rate = sample_rate
        self.max_chunks_in_memory = max_chunks_in_memory
        self.chunk_processor = chunk_processor
        
        self._chunks = deque(maxlen=max_chunks_in_memory)
        self._current_chunk = np.zeros(chunk_size, dtype=np.float32)
        self._current_chunk_pos = 0
        self._chunk_counter = 0
        
        self._lock = threading.Lock()
        
        # Performance monitoring
        self._stats = {
            'chunks_processed': 0,
            'samples_processed': 0,
            'processing_time_total': 0.0,
            'avg_processing_time': 0.0,
            'max_processing_time': 0.0,
            'min_processing_time': float('inf'),
            'chunks_dropped': 0
        }
    
    def process_audio_stream(self, audio_data: Union[np.ndarray, torch.Tensor]) -> List[AudioChunk]:
        """
        Process streaming audio data.
        
        Args:
            audio_data: Input audio stream
            
        Returns:
            List of completed chunks
        """
        if isinstance(audio_data, torch.Tensor):
            audio_data = audio_data.detach().cpu().numpy()
        
        audio_data = np.asarray(audio_data, dtype=np.float32)
        
        if audio_data.ndim > 1:
            audio_data = np.mean(audio_data, axis=0)
        
        completed_chunks = []
        
        with self._lock:
            data_idx = 0
            
            while data_idx < len(audio_data):
                # Calculate how much we can write to current chunk
                remaining_chunk = self.chunk_size - self._current_chunk_pos
                remaining_data = len(audio_data) - data_idx
                samples_to_write = min(remaining_chunk, remaining_data)
                
                # Write to current chunk
                end_pos = self._current_chunk_pos + samples_to_write
                self._current_chunk[self._current_chunk_pos:end_pos] = audio_data[data_idx:data_idx + samples_to_write]
                
                self._current_chunk_pos += samples_to_write
                data_idx += samples_to_write
                
                # Check if chunk is complete
                if self._current_chunk_pos >= self.chunk_size:
                    # Create completed chunk
                    chunk = AudioChunk(
                        data=self._current_chunk.copy(),
                        timestamp=time.time(),
                        sample_rate=self.sample_rate,
                        chunk_id=self._chunk_counter
                    )
                    
                    completed_chunks.append(chunk)
                    self._chunk_counter += 1
                    
                    # Store chunk in memory if there's space
                    if len(self._chunks) < self.max_chunks_in_memory:
                        self._chunks.append(chunk)
                    else:
                        self._stats['chunks_dropped'] += 1
                    
                    # Process chunk if processor is available
                    if self.chunk_processor:
                        start_time = time.time()
                        try:
                            self.chunk_processor(chunk)
                            processing_time = (time.time() - start_time) * 1000  # ms
                            self._update_processing_stats(processing_time)
                        except Exception as e:
                            print(f"Error in chunk processor: {e}")
                    
                    # Reset current chunk
                    self._current_chunk.fill(0)
                    self._current_chunk_pos = 0
        
        return completed_chunks
    
    def _update_processing_stats(self, processing_time_ms: float):
        """Update processing performance statistics."""
        self._stats['chunks_processed'] += 1
        self._stats['processing_time_total'] += processing_time_ms
        self._stats['avg_processing_time'] = (
            self._stats['processing_time_total'] / self._stats['chunks_processed']
        )
        self._stats['max_processing_time'] = max(self._stats['max_processing_time'], processing_time_ms)
        self._stats['min_processing_time'] = min(self._stats['min_processing_time'], processing_time_ms)
    
    def get_recent_chunks(self, n: int = 5) -> List[AudioChunk]:
        """Get the most recent n chunks."""
        with self._lock:
            return list(self._chunks)[-n:]
    
    def get_stats(self) -> dict:
        """Get buffer and processing statistics."""
        with self._lock:
            stats = self._stats.copy()
            stats.update({
                'chunk_size': self.chunk_size,
                'sample_rate': self.sample_rate,
                'current_chunk_fill': self._current_chunk_pos / self.chunk_size,
                'chunks_in_memory': len(self._chunks),
                'max_chunks_in_memory': self.max_chunks_in_memory
            })
            return stats
    
    def clear(self):
        """Clear all buffers and reset state."""
        with self._lock:
            self._chunks.clear()
            self._current_chunk.fill(0)
            self._current_chunk_pos = 0
            self._chunk_counter = 0


def create_audio_buffer(buffer_type: str = "circular", **kwargs) -> Union[AudioBuffer, CircularAudioBuffer, StreamingAudioBuffer]:
    """
    Factory function to create audio buffers.
    
    Args:
        buffer_type: Type of buffer ("basic", "circular", "streaming")
        **kwargs: Arguments for buffer initialization
        
    Returns:
        Audio buffer instance
    """
    if buffer_type == "basic":
        return AudioBuffer(**kwargs)
    elif buffer_type == "circular":
        return CircularAudioBuffer(**kwargs)
    elif buffer_type == "streaming":
        return StreamingAudioBuffer(**kwargs)
    else:
        raise ValueError(f"Unknown buffer type: {buffer_type}")


def estimate_buffer_memory_usage(buffer_duration: float, sample_rate: int,
                                max_chunks: int = 10) -> dict:
    """
    Estimate memory usage for audio buffers.
    
    Args:
        buffer_duration: Buffer duration in seconds
        sample_rate: Audio sample rate
        max_chunks: Maximum chunks in memory
        
    Returns:
        Memory usage estimates in bytes
    """
    samples = int(buffer_duration * sample_rate)
    
    # Assuming 32-bit float samples
    buffer_memory = samples * 4  # bytes
    chunk_memory = buffer_memory * max_chunks
    
    return {
        'buffer_memory_bytes': buffer_memory,
        'buffer_memory_mb': buffer_memory / (1024 * 1024),
        'chunk_memory_bytes': chunk_memory,
        'chunk_memory_mb': chunk_memory / (1024 * 1024),
        'total_memory_bytes': buffer_memory + chunk_memory,
        'total_memory_mb': (buffer_memory + chunk_memory) / (1024 * 1024)
    }