src/realtime_processor.py

"""
Baby Cry AI - Real-time Audio Processing
Step 5: Implement real-time audio capture and processing
"""

import pyaudio
import numpy as np
import threading
import queue
import time
from collections import deque
import warnings
warnings.filterwarnings('ignore')

from models.baseline_model import BaselineModel
from audio_processor import AudioProcessor

class RealTimeCryDetector:
    def __init__(self, model_path="models/baseline_model.pkl"):
        self.model = BaselineModel(model_path)
        self.audio_processor = AudioProcessor()
        
        # Audio parameters
        self.chunk_size = 1024
        self.sample_rate = 22050
        self.channels = 1
        self.format = pyaudio.paFloat32
        
        # Audio buffer for analysis
        self.audio_buffer = deque(maxlen=int(self.sample_rate * 3))  # 3 seconds buffer
        
        # Control flags
        self.is_recording = False
        self.is_processing = False
        
        # Audio stream
        self.audio = None
        self.stream = None
        
        # Processing queue
        self.audio_queue = queue.Queue()
        
        # Callbacks
        self.on_prediction = None
        self.on_error = None
        
        # Load model
        if not self.model.load_model():
            print("⚠️ Warning: No trained model found. Real-time analysis will not work.")
    
    def audio_callback(self, in_data, frame_count, time_info, status):
        """Callback function for audio input"""
        if self.is_recording:
            # Convert bytes to float32 numpy array
            audio_data = np.frombuffer(in_data, dtype=np.float32)
            
            # Add to buffer
            self.audio_buffer.extend(audio_data)
            
            # Add to processing queue
            self.audio_queue.put(audio_data)
        
        return (None, pyaudio.paContinue)
    
    def start_recording(self):
        """Start real-time audio recording"""
        if self.is_recording:
            print("⚠️ Already recording")
            return False
        
        try:
            # Initialize PyAudio
            self.audio = pyaudio.PyAudio()
            
            # Create audio stream
            self.stream = self.audio.open(
                format=self.format,
                channels=self.channels,
                rate=self.sample_rate,
                input=True,
                frames_per_buffer=self.chunk_size,
                stream_callback=self.audio_callback
            )
            
            # Start recording
            self.is_recording = True
            self.stream.start_stream()
            
            # Start processing thread
            self.is_processing = True
            self.processing_thread = threading.Thread(target=self._process_audio_continuously, daemon=True)
            self.processing_thread.start()
            
            print("🎤 Real-time recording started")
            return True
            
        except Exception as e:
            print(f"❌ Error starting recording: {e}")
            if self.on_error:
                self.on_error(str(e))
            return False
    
    def stop_recording(self):
        """Stop real-time audio recording"""
        if not self.is_recording:
            print("⚠️ Not currently recording")
            return
        
        # Stop recording
        self.is_recording = False
        self.is_processing = False
        
        # Stop and close stream
        if self.stream:
            self.stream.stop_stream()
            self.stream.close()
        
        # Terminate PyAudio
        if self.audio:
            self.audio.terminate()
        
        # Clear buffer
        self.audio_buffer.clear()
        
        print("⏹️ Real-time recording stopped")
    
    def _process_audio_continuously(self):
        """Continuously process audio chunks"""
        last_analysis_time = 0
        analysis_interval = 2.0  # Analyze every 2 seconds
        
        while self.is_processing:
            try:
                current_time = time.time()
                
                # Check if it's time for analysis
                if current_time - last_analysis_time >= analysis_interval:
                    if len(self.audio_buffer) >= self.sample_rate * 2:  # At least 2 seconds of audio
                        # Get recent audio data
                        recent_audio = np.array(list(self.audio_buffer)[-int(self.sample_rate * 3):])
                        
                        # Analyze
                        prediction, confidence = self._analyze_audio_chunk(recent_audio)
                        
                        if prediction and confidence > 0.6:  # Only report high-confidence predictions
                            if self.on_prediction:
                                self.on_prediction(prediction, confidence)
                        
                        last_analysis_time = current_time
                
                # Small sleep to prevent excessive CPU usage
                time.sleep(0.1)
                
            except Exception as e:
                print(f"❌ Error in audio processing: {e}")
                if self.on_error:
                    self.on_error(str(e))
                break
    
    def _analyze_audio_chunk(self, audio_data):
        """Analyze a chunk of audio data"""
        try:
            if not self.model.is_trained:
                return None, None
            
            # Extract features
            features = self.audio_processor.extract_features_from_array(audio_data, self.sample_rate)
            
            if features is None:
                return None, None
            
            # Make prediction
            prediction, confidence = self.model.predict(features)
            
            return prediction, confidence
            
        except Exception as e:
            print(f"❌ Error analyzing audio chunk: {e}")
            return None, None
    
    def get_current_audio_level(self):
        """Get current audio level (for visualization)"""
        if not self.audio_buffer:
            return 0.0
        
        # Get RMS of recent audio
        recent_audio = np.array(list(self.audio_buffer)[-1024:])  # Last 1024 samples
        rms = np.sqrt(np.mean(recent_audio**2))
        
        return min(rms * 100, 100.0)  # Scale to 0-100
    
    def set_prediction_callback(self, callback):
        """Set callback function for predictions"""
        self.on_prediction = callback
    
    def set_error_callback(self, callback):
        """Set callback function for errors"""
        self.on_error = callback
    
    def is_ready(self):
        """Check if the detector is ready for real-time analysis"""
        return self.model.is_trained
    
    def get_status(self):
        """Get current status"""
        return {
            'recording': self.is_recording,
            'processing': self.is_processing,
            'model_ready': self.model.is_trained,
            'audio_level': self.get_current_audio_level(),
            'buffer_size': len(self.audio_buffer)
        }

# Example usage and testing
if __name__ == "__main__":
    print("🎤 Real-time Cry Detector Test")
    print("=" * 40)
    
    # Initialize detector
    detector = RealTimeCryDetector()
    
    if not detector.is_ready():
        print("❌ Model not ready. Please train the model first.")
        exit(1)
    
    # Set up callbacks
    def on_prediction(prediction, confidence):
        print(f"🔍 Prediction: {prediction} (Confidence: {confidence:.2f})")
    
    def on_error(error):
        print(f"❌ Error: {error}")
    
    detector.set_prediction_callback(on_prediction)
    detector.set_error_callback(on_error)
    
    try:
        # Start recording
        if detector.start_recording():
            print("🎤 Recording started. Speak into your microphone...")
            print("Press Ctrl+C to stop")
            
            # Keep running
            while True:
                status = detector.get_status()
                print(f"\r📊 Status: Recording={status['recording']}, Level={status['audio_level']:.1f}%", end="")
                time.sleep(0.5)
        
    except KeyboardInterrupt:
        print("\n⏹️ Stopping recording...")
        detector.stop_recording()
        print("✅ Recording stopped")
    
    except Exception as e:
        print(f"❌ Error: {e}")
        detector.stop_recording()