Trixy3/trixy_core/ml/voice_recognition/optimization.py

"""
Performance optimization for Voice Recognition system.

This module provides comprehensive performance optimization capabilities:
- Real-time performance monitoring
- Model optimization for deployment
- Memory management
- Latency optimization
- Hardware-specific optimizations
"""

import gc
import logging
import time
import threading
from typing import Dict, List, Optional, Any, Tuple, Union
from dataclasses import dataclass, field
from collections import deque
import psutil
import numpy as np
import torch
import torch.nn as nn
from torch.profiler import profile, record_function, ProfilerActivity

# Configure logging
logger = logging.getLogger(__name__)


@dataclass
class VRPerformanceMetrics:
    """Performance metrics for voice recognition system."""
    # Timing metrics
    avg_inference_time_ms: float = 0.0
    max_inference_time_ms: float = 0.0
    min_inference_time_ms: float = float('inf')
    std_inference_time_ms: float = 0.0
    
    # Throughput metrics
    frames_per_second: float = 0.0
    audio_realtime_factor: float = 0.0  # Processing speed vs real-time
    
    # Memory metrics
    memory_usage_mb: float = 0.0
    peak_memory_mb: float = 0.0
    gpu_memory_mb: float = 0.0
    
    # Model metrics
    model_size_mb: float = 0.0
    flops_per_inference: Optional[int] = None
    
    # Quality metrics
    avg_confidence: float = 0.0
    successful_identifications: int = 0
    total_attempts: int = 0
    
    # System metrics
    cpu_usage_percent: float = 0.0
    temperature_celsius: Optional[float] = None
    
    def accuracy_rate(self) -> float:
        """Calculate identification accuracy rate."""
        if self.total_attempts == 0:
            return 0.0
        return self.successful_identifications / self.total_attempts
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert metrics to dictionary."""
        return {
            'timing': {
                'avg_inference_time_ms': self.avg_inference_time_ms,
                'max_inference_time_ms': self.max_inference_time_ms,
                'min_inference_time_ms': self.min_inference_time_ms,
                'std_inference_time_ms': self.std_inference_time_ms
            },
            'throughput': {
                'frames_per_second': self.frames_per_second,
                'audio_realtime_factor': self.audio_realtime_factor
            },
            'memory': {
                'memory_usage_mb': self.memory_usage_mb,
                'peak_memory_mb': self.peak_memory_mb,
                'gpu_memory_mb': self.gpu_memory_mb
            },
            'model': {
                'model_size_mb': self.model_size_mb,
                'flops_per_inference': self.flops_per_inference
            },
            'quality': {
                'avg_confidence': self.avg_confidence,
                'accuracy_rate': self.accuracy_rate(),
                'successful_identifications': self.successful_identifications,
                'total_attempts': self.total_attempts
            },
            'system': {
                'cpu_usage_percent': self.cpu_usage_percent,
                'temperature_celsius': self.temperature_celsius
            }
        }


class VRPerformanceMonitor:
    """Real-time performance monitor for voice recognition system."""
    
    def __init__(self, 
                 window_size: int = 100,
                 monitoring_interval: float = 1.0,
                 enable_detailed_profiling: bool = False):
        """
        Initialize performance monitor.
        
        Args:
            window_size: Number of samples to keep for rolling statistics
            monitoring_interval: Monitoring update interval in seconds
            enable_detailed_profiling: Enable detailed PyTorch profiling
        """
        self.window_size = window_size
        self.monitoring_interval = monitoring_interval
        self.enable_detailed_profiling = enable_detailed_profiling
        
        # Performance data storage
        self.inference_times = deque(maxlen=window_size)
        self.confidence_scores = deque(maxlen=window_size)
        self.memory_usage = deque(maxlen=window_size)
        
        # Monitoring state
        self._is_monitoring = False
        self._monitor_thread: Optional[threading.Thread] = None
        self._start_time = time.time()
        self._frame_count = 0
        self._successful_count = 0
        
        # System monitoring
        self._process = psutil.Process()
        self._peak_memory = 0.0
        
        # Profiling
        self._profiler: Optional[torch.profiler.profile] = None
        self._profiling_active = False
    
    def start_monitoring(self):
        """Start performance monitoring."""
        if self._is_monitoring:
            logger.warning("Performance monitoring already active")
            return
        
        self._is_monitoring = True
        self._start_time = time.time()
        self._frame_count = 0
        self._successful_count = 0
        
        # Start monitoring thread
        self._monitor_thread = threading.Thread(
            target=self._monitoring_loop,
            daemon=True
        )
        self._monitor_thread.start()
        
        logger.info("Performance monitoring started")
    
    def stop_monitoring(self):
        """Stop performance monitoring."""
        self._is_monitoring = False
        
        if self._monitor_thread:
            self._monitor_thread.join(timeout=2.0)
        
        if self._profiling_active:
            self.stop_profiling()
        
        logger.info("Performance monitoring stopped")
    
    def record_inference(self, 
                        inference_time_ms: float,
                        confidence: float,
                        success: bool = True):
        """
        Record inference performance data.
        
        Args:
            inference_time_ms: Inference time in milliseconds
            confidence: Confidence score
            success: Whether inference was successful
        """
        self.inference_times.append(inference_time_ms)
        self.confidence_scores.append(confidence)
        self._frame_count += 1
        
        if success:
            self._successful_count += 1
        
        # Update peak memory
        current_memory = self._get_memory_usage()
        self._peak_memory = max(self._peak_memory, current_memory)
    
    def get_metrics(self) -> VRPerformanceMetrics:
        """Get current performance metrics."""
        metrics = VRPerformanceMetrics()
        
        # Timing metrics
        if self.inference_times:
            times = list(self.inference_times)
            metrics.avg_inference_time_ms = np.mean(times)
            metrics.max_inference_time_ms = np.max(times)
            metrics.min_inference_time_ms = np.min(times)
            metrics.std_inference_time_ms = np.std(times)
        
        # Throughput metrics
        elapsed_time = time.time() - self._start_time
        if elapsed_time > 0:
            metrics.frames_per_second = self._frame_count / elapsed_time
            
            # Estimate real-time factor (assuming 16kHz audio with 30ms frames)
            audio_duration_s = self._frame_count * 0.03  # 30ms frames
            metrics.audio_realtime_factor = audio_duration_s / elapsed_time
        
        # Memory metrics
        metrics.memory_usage_mb = self._get_memory_usage()
        metrics.peak_memory_mb = self._peak_memory
        metrics.gpu_memory_mb = self._get_gpu_memory_usage()
        
        # Quality metrics
        if self.confidence_scores:
            metrics.avg_confidence = np.mean(list(self.confidence_scores))
        metrics.successful_identifications = self._successful_count
        metrics.total_attempts = self._frame_count
        
        # System metrics
        metrics.cpu_usage_percent = self._process.cpu_percent()
        
        # Temperature (if available)
        try:
            temps = psutil.sensors_temperatures()
            if temps:
                cpu_temps = temps.get('coretemp', [])
                if cpu_temps:
                    metrics.temperature_celsius = cpu_temps[0].current
        except:
            pass
        
        return metrics
    
    def start_profiling(self, 
                       activities: List[ProfilerActivity] = None,
                       record_shapes: bool = True):
        """Start detailed PyTorch profiling."""
        if not self.enable_detailed_profiling:
            logger.warning("Detailed profiling not enabled")
            return
        
        if self._profiling_active:
            logger.warning("Profiling already active")
            return
        
        if activities is None:
            activities = [ProfilerActivity.CPU]
            if torch.cuda.is_available():
                activities.append(ProfilerActivity.CUDA)
        
        self._profiler = profile(
            activities=activities,
            record_shapes=record_shapes,
            with_stack=True
        )
        self._profiler.__enter__()
        self._profiling_active = True
        
        logger.info("PyTorch profiling started")
    
    def stop_profiling(self) -> Optional[str]:
        """Stop profiling and return results."""
        if not self._profiling_active or not self._profiler:
            return None
        
        self._profiler.__exit__(None, None, None)
        self._profiling_active = False
        
        # Get profiling results
        results = self._profiler.key_averages().table(sort_by="cpu_time_total", row_limit=20)
        
        logger.info("PyTorch profiling stopped")
        return results
    
    def _monitoring_loop(self):
        """Background monitoring loop."""
        while self._is_monitoring:
            try:
                # Update memory usage
                current_memory = self._get_memory_usage()
                self.memory_usage.append(current_memory)
                
                # Log periodic stats
                if len(self.inference_times) > 10:
                    metrics = self.get_metrics()
                    logger.debug(f"Performance: {metrics.avg_inference_time_ms:.1f}ms avg, "
                               f"{metrics.frames_per_second:.1f} FPS, "
                               f"{metrics.memory_usage_mb:.1f}MB memory")
                
                time.sleep(self.monitoring_interval)
                
            except Exception as e:
                logger.error(f"Error in monitoring loop: {e}")
                time.sleep(self.monitoring_interval)
    
    def _get_memory_usage(self) -> float:
        """Get current memory usage in MB."""
        try:
            memory_info = self._process.memory_info()
            return memory_info.rss / (1024 * 1024)  # Convert to MB
        except:
            return 0.0
    
    def _get_gpu_memory_usage(self) -> float:
        """Get GPU memory usage in MB."""
        if not torch.cuda.is_available():
            return 0.0
        
        try:
            memory_allocated = torch.cuda.memory_allocated() / (1024 * 1024)
            return memory_allocated
        except:
            return 0.0


class VRModelOptimizer:
    """Model optimization utilities for voice recognition."""
    
    @staticmethod
    def optimize_for_inference(model: nn.Module,
                             enable_torch_compile: bool = True,
                             use_half_precision: bool = False,
                             enable_jit: bool = False) -> nn.Module:
        """
        Optimize model for inference.
        
        Args:
            model: PyTorch model to optimize
            enable_torch_compile: Use torch.compile (PyTorch 2.0+)
            use_half_precision: Convert to half precision
            enable_jit: Use TorchScript JIT compilation
            
        Returns:
            Optimized model
        """
        logger.info("Optimizing model for inference...")
        
        # Set to evaluation mode
        model.eval()
        
        # Disable gradients
        for param in model.parameters():
            param.requires_grad = False
        
        # Half precision
        if use_half_precision:
            model = model.half()
            logger.info("Converted model to half precision")
        
        # TorchScript JIT
        if enable_jit:
            try:
                # Create example input for tracing
                example_input = VRModelOptimizer._create_example_input(model)
                model = torch.jit.trace(model, example_input)
                logger.info("Applied TorchScript JIT compilation")
            except Exception as e:
                logger.warning(f"JIT compilation failed: {e}")
        
        # Torch compile (PyTorch 2.0+)
        if enable_torch_compile and hasattr(torch, 'compile'):
            try:
                model = torch.compile(model, mode='reduce-overhead')
                logger.info("Applied torch.compile optimization")
            except Exception as e:
                logger.warning(f"torch.compile failed: {e}")
        
        return model
    
    @staticmethod
    def quantize_model(model: nn.Module,
                      quantization_backend: str = 'fbgemm',
                      calibration_data: Optional[torch.Tensor] = None) -> nn.Module:
        """
        Quantize model for improved performance.
        
        Args:
            model: Model to quantize
            quantization_backend: Quantization backend ('fbgemm', 'qnnpack')
            calibration_data: Calibration data for quantization
            
        Returns:
            Quantized model
        """
        logger.info(f"Quantizing model with {quantization_backend} backend...")
        
        try:
            torch.backends.quantized.engine = quantization_backend
            
            # Prepare model for quantization
            model.eval()
            model_prepared = torch.quantization.prepare(model)
            
            # Calibrate if data provided
            if calibration_data is not None:
                with torch.no_grad():
                    model_prepared(calibration_data)
            
            # Convert to quantized model
            quantized_model = torch.quantization.convert(model_prepared)
            
            logger.info("Model quantization completed successfully")
            return quantized_model
            
        except Exception as e:
            logger.error(f"Model quantization failed: {e}")
            return model
    
    @staticmethod
    def estimate_model_size(model: nn.Module) -> float:
        """
        Estimate model size in MB.
        
        Args:
            model: PyTorch model
            
        Returns:
            Model size in MB
        """
        param_size = 0
        buffer_size = 0
        
        for param in model.parameters():
            param_size += param.nelement() * param.element_size()
        
        for buffer in model.buffers():
            buffer_size += buffer.nelement() * buffer.element_size()
        
        total_size = param_size + buffer_size
        return total_size / (1024 * 1024)  # Convert to MB
    
    @staticmethod
    def _create_example_input(model: nn.Module) -> torch.Tensor:
        """Create example input for model tracing."""
        # This is a simplified example - in practice, you'd need to
        # determine the correct input shape for your specific model
        return torch.randn(1, 40, 100)  # Batch, features, time


class VRMemoryManager:
    """Memory management utilities for voice recognition."""
    
    def __init__(self, max_memory_mb: int = 512):
        """
        Initialize memory manager.
        
        Args:
            max_memory_mb: Maximum memory usage limit in MB
        """
        self.max_memory_mb = max_memory_mb
        self._cleanup_threshold = max_memory_mb * 0.8  # Trigger cleanup at 80%
        self._last_cleanup = time.time()
        self._cleanup_interval = 30.0  # Cleanup interval in seconds
    
    def check_memory_usage(self) -> bool:
        """
        Check if memory usage is within limits.
        
        Returns:
            True if memory usage is acceptable
        """
        current_usage = self._get_current_memory_usage()
        
        if current_usage > self._cleanup_threshold:
            self._trigger_cleanup()
            return False
        
        return current_usage <= self.max_memory_mb
    
    def force_cleanup(self):
        """Force memory cleanup."""
        self._trigger_cleanup()
    
    def _get_current_memory_usage(self) -> float:
        """Get current process memory usage in MB."""
        try:
            process = psutil.Process()
            memory_info = process.memory_info()
            return memory_info.rss / (1024 * 1024)
        except:
            return 0.0
    
    def _trigger_cleanup(self):
        """Trigger memory cleanup operations."""
        current_time = time.time()
        
        # Rate limit cleanup operations
        if current_time - self._last_cleanup < self._cleanup_interval:
            return
        
        logger.info("Triggering memory cleanup...")
        
        # Python garbage collection
        gc.collect()
        
        # PyTorch cache cleanup
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
        
        self._last_cleanup = current_time
        
        # Log memory usage after cleanup
        usage_after = self._get_current_memory_usage()
        logger.info(f"Memory usage after cleanup: {usage_after:.1f}MB")


class VRLatencyOptimizer:
    """Latency optimization utilities for voice recognition."""
    
    def __init__(self, target_latency_ms: float = 100.0):
        """
        Initialize latency optimizer.
        
        Args:
            target_latency_ms: Target latency in milliseconds
        """
        self.target_latency_ms = target_latency_ms
        self._latency_history = deque(maxlen=50)
        self._optimization_suggestions = []
    
    def record_latency(self, latency_ms: float):
        """Record inference latency."""
        self._latency_history.append(latency_ms)
        
        # Check if optimization needed
        if len(self._latency_history) >= 10:
            avg_latency = np.mean(list(self._latency_history))
            if avg_latency > self.target_latency_ms * 1.2:  # 20% tolerance
                self._generate_optimization_suggestions(avg_latency)
    
    def get_optimization_suggestions(self) -> List[str]:
        """Get optimization suggestions based on performance."""
        return self._optimization_suggestions.copy()
    
    def _generate_optimization_suggestions(self, avg_latency: float):
        """Generate optimization suggestions."""
        suggestions = []
        
        if avg_latency > self.target_latency_ms * 2:
            suggestions.append("Consider using model quantization")
            suggestions.append("Enable half precision inference")
            suggestions.append("Reduce audio processing length")
        elif avg_latency > self.target_latency_ms * 1.5:
            suggestions.append("Enable torch.compile optimization")
            suggestions.append("Consider using a smaller model architecture")
        else:
            suggestions.append("Fine-tune batch size for optimal performance")
        
        # Update suggestions if they've changed
        if suggestions != self._optimization_suggestions:
            self._optimization_suggestions = suggestions
            logger.info(f"Generated optimization suggestions: {suggestions}")


def create_vr_performance_monitor(window_size: int = 100,
                                monitoring_interval: float = 1.0,
                                enable_profiling: bool = False) -> VRPerformanceMonitor:
    """
    Factory function to create a VRPerformanceMonitor.
    
    Args:
        window_size: Number of samples for rolling statistics
        monitoring_interval: Monitoring update interval
        enable_profiling: Enable detailed profiling
        
    Returns:
        Configured VRPerformanceMonitor instance
    """
    return VRPerformanceMonitor(
        window_size=window_size,
        monitoring_interval=monitoring_interval,
        enable_detailed_profiling=enable_profiling
    )


def optimize_vr_for_deployment(model: nn.Module,
                             target_device: str = "cpu",
                             target_latency_ms: float = 100.0,
                             enable_quantization: bool = False,
                             calibration_data: Optional[torch.Tensor] = None) -> Dict[str, Any]:
    """
    Comprehensive optimization for voice recognition deployment.
    
    Args:
        model: Model to optimize
        target_device: Target deployment device
        target_latency_ms: Target latency requirement
        enable_quantization: Whether to apply quantization
        calibration_data: Data for quantization calibration
        
    Returns:
        Dictionary containing optimized model and metrics
    """
    logger.info("Starting comprehensive voice recognition optimization...")
    
    results = {
        'original_model': model,
        'optimized_model': None,
        'optimization_applied': [],
        'performance_improvement': {},
        'model_size_reduction': 0.0
    }
    
    # Measure original performance
    original_size = VRModelOptimizer.estimate_model_size(model)
    
    # Apply optimizations
    optimized_model = model
    
    # Basic inference optimization
    optimized_model = VRModelOptimizer.optimize_for_inference(
        optimized_model,
        enable_torch_compile=True,
        use_half_precision=(target_device != "cpu"),
        enable_jit=False
    )
    results['optimization_applied'].append('inference_optimization')
    
    # Quantization if requested
    if enable_quantization:
        optimized_model = VRModelOptimizer.quantize_model(
            optimized_model,
            calibration_data=calibration_data
        )
        results['optimization_applied'].append('quantization')
    
    # Measure optimized performance
    optimized_size = VRModelOptimizer.estimate_model_size(optimized_model)
    results['model_size_reduction'] = (original_size - optimized_size) / original_size * 100
    
    results['optimized_model'] = optimized_model
    results['performance_improvement'] = {
        'original_size_mb': original_size,
        'optimized_size_mb': optimized_size,
        'size_reduction_percent': results['model_size_reduction']
    }
    
    logger.info(f"Optimization completed. Model size reduced by {results['model_size_reduction']:.1f}%")
    
    return results