Trixy3/trixy_core/ml/wakeword/optimization.py

"""
Performance optimization utilities for wakeword detection.

This module provides tools for optimizing the wakeword detection system
for real-time performance, including model optimization, memory management,
and latency reduction techniques.
"""

import gc
import time
import threading
import psutil
import warnings
from typing import Dict, Any, Optional, List, Tuple
from dataclasses import dataclass
from pathlib import Path

import torch
import torch.nn as nn
import numpy as np


@dataclass
class PerformanceMetrics:
    """Performance metrics for monitoring system performance."""
    
    # Latency metrics (milliseconds)
    inference_latency_ms: float = 0.0
    feature_extraction_latency_ms: float = 0.0
    total_processing_latency_ms: float = 0.0
    
    # Throughput metrics
    chunks_per_second: float = 0.0
    audio_realtime_factor: float = 0.0  # >1.0 means faster than realtime
    
    # Memory metrics (MB)
    memory_usage_mb: float = 0.0
    peak_memory_mb: float = 0.0
    gpu_memory_mb: float = 0.0
    
    # CPU metrics
    cpu_usage_percent: float = 0.0
    cpu_load_average: float = 0.0
    
    # Accuracy metrics
    confidence_scores: List[float] = None
    detection_rate: float = 0.0
    false_positive_rate: float = 0.0
    
    def __post_init__(self):
        if self.confidence_scores is None:
            self.confidence_scores = []


class PerformanceMonitor:
    """
    Real-time performance monitoring for wakeword detection.
    
    Tracks various performance metrics and provides alerts for
    performance degradation.
    """
    
    def __init__(self, monitoring_interval: float = 1.0,
                 memory_warning_threshold: float = 512.0,  # MB
                 latency_warning_threshold: float = 100.0):  # ms
        """
        Initialize performance monitor.
        
        Args:
            monitoring_interval: How often to collect metrics (seconds)
            memory_warning_threshold: Memory usage warning threshold (MB)
            latency_warning_threshold: Latency warning threshold (ms)
        """
        self.monitoring_interval = monitoring_interval
        self.memory_warning_threshold = memory_warning_threshold
        self.latency_warning_threshold = latency_warning_threshold
        
        self._is_monitoring = False
        self._monitor_thread = None
        self._metrics_history: List[PerformanceMetrics] = []
        self._lock = threading.Lock()
        
        # Performance counters
        self._chunk_count = 0
        self._total_inference_time = 0.0
        self._total_feature_time = 0.0
        self._start_time = None
        
        # Warnings tracking
        self._warnings_issued = set()
    
    def start_monitoring(self):
        """Start performance monitoring."""
        if self._is_monitoring:
            return
        
        self._is_monitoring = True
        self._start_time = time.time()
        self._monitor_thread = threading.Thread(
            target=self._monitoring_loop,
            name="PerformanceMonitor",
            daemon=True
        )
        self._monitor_thread.start()
    
    def stop_monitoring(self):
        """Stop performance monitoring."""
        self._is_monitoring = False
        if self._monitor_thread:
            self._monitor_thread.join(timeout=2.0)
    
    def record_inference(self, inference_time_ms: float, feature_time_ms: float):
        """Record inference performance metrics."""
        with self._lock:
            self._chunk_count += 1
            self._total_inference_time += inference_time_ms
            self._total_feature_time += feature_time_ms
    
    def _monitoring_loop(self):
        """Main monitoring loop."""
        while self._is_monitoring:
            try:
                metrics = self._collect_metrics()
                
                with self._lock:
                    self._metrics_history.append(metrics)
                    # Keep only recent metrics (last 100 samples)
                    if len(self._metrics_history) > 100:
                        self._metrics_history.pop(0)
                
                # Check for performance issues
                self._check_performance_warnings(metrics)
                
                time.sleep(self.monitoring_interval)
                
            except Exception as e:
                warnings.warn(f"Error in performance monitoring: {e}")
    
    def _collect_metrics(self) -> PerformanceMetrics:
        """Collect current performance metrics."""
        with self._lock:
            # Calculate latencies
            if self._chunk_count > 0:
                avg_inference_latency = self._total_inference_time / self._chunk_count
                avg_feature_latency = self._total_feature_time / self._chunk_count
                total_latency = avg_inference_latency + avg_feature_latency
            else:
                avg_inference_latency = avg_feature_latency = total_latency = 0.0
            
            # Calculate throughput
            elapsed_time = time.time() - self._start_time if self._start_time else 1.0
            chunks_per_second = self._chunk_count / elapsed_time
        
        # System metrics
        process = psutil.Process()
        memory_info = process.memory_info()
        memory_usage_mb = memory_info.rss / (1024 * 1024)
        cpu_percent = process.cpu_percent()
        
        # GPU memory if available
        gpu_memory_mb = 0.0
        if torch.cuda.is_available():
            gpu_memory_mb = torch.cuda.memory_allocated() / (1024 * 1024)
        
        return PerformanceMetrics(
            inference_latency_ms=avg_inference_latency,
            feature_extraction_latency_ms=avg_feature_latency,
            total_processing_latency_ms=total_latency,
            chunks_per_second=chunks_per_second,
            memory_usage_mb=memory_usage_mb,
            gpu_memory_mb=gpu_memory_mb,
            cpu_usage_percent=cpu_percent
        )
    
    def _check_performance_warnings(self, metrics: PerformanceMetrics):
        """Check for performance issues and issue warnings."""
        # Memory warning
        if metrics.memory_usage_mb > self.memory_warning_threshold:
            warning_key = "high_memory"
            if warning_key not in self._warnings_issued:
                warnings.warn(f"High memory usage: {metrics.memory_usage_mb:.1f}MB")
                self._warnings_issued.add(warning_key)
        
        # Latency warning
        if metrics.total_processing_latency_ms > self.latency_warning_threshold:
            warning_key = "high_latency"
            if warning_key not in self._warnings_issued:
                warnings.warn(f"High processing latency: {metrics.total_processing_latency_ms:.1f}ms")
                self._warnings_issued.add(warning_key)
    
    def get_current_metrics(self) -> Optional[PerformanceMetrics]:
        """Get the most recent performance metrics."""
        with self._lock:
            return self._metrics_history[-1] if self._metrics_history else None
    
    def get_metrics_history(self) -> List[PerformanceMetrics]:
        """Get performance metrics history."""
        with self._lock:
            return self._metrics_history.copy()
    
    def get_performance_summary(self) -> Dict[str, Any]:
        """Get performance summary statistics."""
        with self._lock:
            if not self._metrics_history:
                return {}
            
            recent_metrics = self._metrics_history[-10:]  # Last 10 samples
            
            return {
                'avg_inference_latency_ms': np.mean([m.inference_latency_ms for m in recent_metrics]),
                'max_inference_latency_ms': max(m.inference_latency_ms for m in recent_metrics),
                'avg_memory_usage_mb': np.mean([m.memory_usage_mb for m in recent_metrics]),
                'max_memory_usage_mb': max(m.memory_usage_mb for m in recent_metrics),
                'avg_cpu_usage_percent': np.mean([m.cpu_usage_percent for m in recent_metrics]),
                'chunks_processed': self._chunk_count,
                'total_runtime_seconds': time.time() - self._start_time if self._start_time else 0
            }


class ModelOptimizer:
    """
    Model optimization utilities for improving inference performance.
    
    Provides various optimization techniques including quantization,
    pruning, and compilation.
    """
    
    @staticmethod
    def optimize_for_inference(model: nn.Module, example_input: torch.Tensor,
                             optimization_level: str = "moderate") -> nn.Module:
        """
        Optimize PyTorch model for inference.
        
        Args:
            model: PyTorch model to optimize
            example_input: Example input tensor for optimization
            optimization_level: "light", "moderate", or "aggressive"
            
        Returns:
            Optimized model
        """
        optimized_model = model.eval()
        
        if optimization_level in ["moderate", "aggressive"]:
            # Try to use TorchScript
            try:
                optimized_model = torch.jit.trace(optimized_model, example_input)
            except Exception as e:
                warnings.warn(f"TorchScript optimization failed: {e}")
        
        if optimization_level == "aggressive":
            # Try quantization (CPU only)
            if example_input.device.type == "cpu":
                try:
                    optimized_model = torch.quantization.quantize_dynamic(
                        optimized_model, {nn.Linear, nn.Conv2d}, dtype=torch.qint8
                    )
                except Exception as e:
                    warnings.warn(f"Quantization failed: {e}")
        
        return optimized_model
    
    @staticmethod
    def apply_model_compilation(model: nn.Module) -> nn.Module:
        """
        Apply PyTorch 2.0+ compilation if available.
        
        Args:
            model: Model to compile
            
        Returns:
            Compiled model or original if compilation not available
        """
        try:
            if hasattr(torch, 'compile'):
                return torch.compile(model, mode="reduce-overhead")
            else:
                warnings.warn("torch.compile not available in this PyTorch version")
                return model
        except Exception as e:
            warnings.warn(f"Model compilation failed: {e}")
            return model
    
    @staticmethod
    def estimate_model_flops(model: nn.Module, input_shape: Tuple[int, ...]) -> int:
        """
        Estimate FLOPs for model inference.
        
        Args:
            model: PyTorch model
            input_shape: Input tensor shape
            
        Returns:
            Estimated FLOPs
        """
        try:
            # Try to use thop if available
            import thop
            dummy_input = torch.randn(1, *input_shape)
            flops, _ = thop.profile(model, inputs=(dummy_input,), verbose=False)
            return int(flops)
        except ImportError:
            # Fallback to simple estimation
            total_params = sum(p.numel() for p in model.parameters())
            # Rough estimation: assume each parameter is used once
            return total_params
    
    @staticmethod
    def benchmark_model(model: nn.Module, input_shape: Tuple[int, ...],
                       device: torch.device, num_runs: int = 100) -> Dict[str, float]:
        """
        Benchmark model inference performance.
        
        Args:
            model: Model to benchmark
            input_shape: Input tensor shape
            device: Device to run benchmark on
            num_runs: Number of benchmark runs
            
        Returns:
            Performance metrics dictionary
        """
        model = model.to(device).eval()
        dummy_input = torch.randn(1, *input_shape, device=device)
        
        # Warm up
        for _ in range(10):
            with torch.no_grad():
                _ = model(dummy_input)
        
        if device.type == "cuda":
            torch.cuda.synchronize()
        
        # Benchmark
        times = []
        for _ in range(num_runs):
            start_time = time.time()
            
            with torch.no_grad():
                _ = model(dummy_input)
            
            if device.type == "cuda":
                torch.cuda.synchronize()
            
            end_time = time.time()
            times.append((end_time - start_time) * 1000)  # ms
        
        return {
            'mean_inference_time_ms': np.mean(times),
            'std_inference_time_ms': np.std(times),
            'min_inference_time_ms': min(times),
            'max_inference_time_ms': max(times),
            'median_inference_time_ms': np.median(times),
            'percentile_95_ms': np.percentile(times, 95),
            'throughput_fps': 1000.0 / np.mean(times)
        }


class MemoryManager:
    """
    Memory management utilities for efficient resource usage.
    
    Helps manage memory allocation and garbage collection for
    optimal performance.
    """
    
    def __init__(self, gc_frequency: int = 100,
                 memory_limit_mb: Optional[float] = None):
        """
        Initialize memory manager.
        
        Args:
            gc_frequency: How often to run garbage collection (every N chunks)
            memory_limit_mb: Memory limit in MB (None for no limit)
        """
        self.gc_frequency = gc_frequency
        self.memory_limit_mb = memory_limit_mb
        
        self._chunk_counter = 0
        self._memory_pool = {}
        self._peak_memory = 0.0
    
    def process_chunk(self, func, *args, **kwargs):
        """
        Process a chunk with memory management.
        
        Args:
            func: Function to call
            *args, **kwargs: Arguments for the function
            
        Returns:
            Function result
        """
        self._chunk_counter += 1
        
        # Run function
        result = func(*args, **kwargs)
        
        # Check memory usage
        current_memory = self._get_memory_usage()
        self._peak_memory = max(self._peak_memory, current_memory)
        
        # Run garbage collection if needed
        if self._chunk_counter % self.gc_frequency == 0:
            self._run_garbage_collection()
        
        # Check memory limit
        if self.memory_limit_mb and current_memory > self.memory_limit_mb:
            warnings.warn(f"Memory usage ({current_memory:.1f}MB) exceeds limit ({self.memory_limit_mb}MB)")
            self._emergency_cleanup()
        
        return result
    
    def _get_memory_usage(self) -> float:
        """Get current memory usage in MB."""
        process = psutil.Process()
        return process.memory_info().rss / (1024 * 1024)
    
    def _run_garbage_collection(self):
        """Run garbage collection."""
        gc.collect()
        
        # Also clear PyTorch cache if on GPU
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    
    def _emergency_cleanup(self):
        """Emergency memory cleanup."""
        # Force garbage collection
        self._run_garbage_collection()
        
        # Clear memory pool
        self._memory_pool.clear()
        
        # Additional PyTorch cleanup
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
    
    def get_memory_stats(self) -> Dict[str, float]:
        """Get memory usage statistics."""
        current_memory = self._get_memory_usage()
        
        stats = {
            'current_memory_mb': current_memory,
            'peak_memory_mb': self._peak_memory,
            'chunks_processed': self._chunk_counter,
            'gc_runs': self._chunk_counter // self.gc_frequency
        }
        
        if torch.cuda.is_available():
            stats.update({
                'gpu_memory_allocated_mb': torch.cuda.memory_allocated() / (1024 * 1024),
                'gpu_memory_reserved_mb': torch.cuda.memory_reserved() / (1024 * 1024),
                'gpu_memory_cached_mb': torch.cuda.memory_cached() / (1024 * 1024)
            })
        
        return stats


class LatencyOptimizer:
    """
    Latency optimization utilities for reducing processing delays.
    
    Provides techniques for minimizing end-to-end latency in
    wakeword detection.
    """
    
    @staticmethod
    def optimize_audio_buffer_size(target_latency_ms: float,
                                  sample_rate: int,
                                  processing_time_ms: float) -> Tuple[int, float]:
        """
        Calculate optimal audio buffer size for target latency.
        
        Args:
            target_latency_ms: Target total latency
            sample_rate: Audio sample rate
            processing_time_ms: Expected processing time
            
        Returns:
            Tuple of (buffer_samples, actual_latency_ms)
        """
        # Account for processing overhead
        available_latency_ms = target_latency_ms - processing_time_ms
        
        # Calculate buffer size
        buffer_duration_s = max(0.010, available_latency_ms / 1000.0)  # Minimum 10ms
        buffer_samples = int(buffer_duration_s * sample_rate)
        
        # Calculate actual latency
        actual_buffer_latency_ms = (buffer_samples / sample_rate) * 1000
        actual_total_latency_ms = actual_buffer_latency_ms + processing_time_ms
        
        return buffer_samples, actual_total_latency_ms
    
    @staticmethod
    def profile_inference_pipeline(detector, num_samples: int = 100) -> Dict[str, Any]:
        """
        Profile the complete inference pipeline for latency bottlenecks.
        
        Args:
            detector: WakewordDetector instance
            num_samples: Number of samples to profile
            
        Returns:
            Profiling results
        """
        # Create dummy audio data
        sample_rate = detector.config.audio_config.sample_rate
        chunk_duration = detector.config.audio_config.chunk_duration
        chunk_samples = int(chunk_duration * sample_rate)
        
        # Profile stages
        audio_generation_times = []
        feature_extraction_times = []
        inference_times = []
        total_times = []
        
        for _ in range(num_samples):
            # Generate dummy audio
            start_time = time.time()
            dummy_audio = np.random.randn(chunk_samples).astype(np.float32)
            audio_generation_times.append((time.time() - start_time) * 1000)
            
            # Feature extraction
            start_time = time.time()
            features = detector.feature_extractor.extract_features(
                torch.from_numpy(dummy_audio)
            )
            feature_extraction_times.append((time.time() - start_time) * 1000)
            
            # Inference
            start_time = time.time()
            with torch.no_grad():
                _ = detector.model(features)
            inference_times.append((time.time() - start_time) * 1000)
            
            total_times.append(feature_extraction_times[-1] + inference_times[-1])
        
        return {
            'feature_extraction': {
                'mean_ms': np.mean(feature_extraction_times),
                'std_ms': np.std(feature_extraction_times),
                'min_ms': min(feature_extraction_times),
                'max_ms': max(feature_extraction_times)
            },
            'inference': {
                'mean_ms': np.mean(inference_times),
                'std_ms': np.std(inference_times),
                'min_ms': min(inference_times),
                'max_ms': max(inference_times)
            },
            'total_processing': {
                'mean_ms': np.mean(total_times),
                'std_ms': np.std(total_times),
                'min_ms': min(total_times),
                'max_ms': max(total_times)
            },
            'throughput_hz': 1000.0 / np.mean(total_times)
        }


def create_performance_monitor(config: Dict[str, Any]) -> PerformanceMonitor:
    """
    Factory function to create performance monitor.
    
    Args:
        config: Configuration dictionary
        
    Returns:
        PerformanceMonitor instance
    """
    return PerformanceMonitor(
        monitoring_interval=config.get('monitoring_interval', 1.0),
        memory_warning_threshold=config.get('memory_warning_threshold', 512.0),
        latency_warning_threshold=config.get('latency_warning_threshold', 100.0)
    )


def optimize_for_deployment(model: nn.Module, config: Dict[str, Any],
                          example_input: torch.Tensor) -> Tuple[nn.Module, Dict[str, Any]]:
    """
    Comprehensive optimization for deployment.
    
    Args:
        model: Model to optimize
        config: Optimization configuration
        example_input: Example input for optimization
        
    Returns:
        Tuple of (optimized_model, optimization_report)
    """
    optimizer = ModelOptimizer()
    
    # Benchmark original model
    original_metrics = optimizer.benchmark_model(
        model, example_input.shape[1:], example_input.device
    )
    
    # Apply optimizations
    optimized_model = model
    
    if config.get('enable_torchscript', True):
        optimized_model = optimizer.optimize_for_inference(
            optimized_model, example_input, 
            optimization_level=config.get('optimization_level', 'moderate')
        )
    
    if config.get('enable_compilation', True):
        optimized_model = optimizer.apply_model_compilation(optimized_model)
    
    # Benchmark optimized model
    optimized_metrics = optimizer.benchmark_model(
        optimized_model, example_input.shape[1:], example_input.device
    )
    
    # Create optimization report
    report = {
        'original_performance': original_metrics,
        'optimized_performance': optimized_metrics,
        'improvement_factor': original_metrics['mean_inference_time_ms'] / optimized_metrics['mean_inference_time_ms'],
        'optimizations_applied': list(config.keys())
    }
    
    return optimized_model, report