visualize_sensor_data.py

"""
Real-time visualization of piezoelectric sensor data from Arduino
Displays raw sensor readings, extracted features, and slip predictions
"""

import serial
import serial.tools.list_ports
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from collections import deque
import re
import argparse
import numpy as np
from datetime import datetime

class SensorVisualizer:
    def __init__(self, port='COM3', baud_rate=9600, buffer_size=500):
        """
        Initialize sensor visualizer
        
        Args:
            port: Serial port (e.g., 'COM3' on Windows, '/dev/ttyUSB0' on Linux)
            baud_rate: Serial communication baud rate
            buffer_size: Number of data points to keep in buffer
        """
        self.port = port
        self.baud_rate = baud_rate
        self.buffer_size = buffer_size
        
        # Data buffers
        self.time_buffer = deque(maxlen=buffer_size)
        self.sensor_buffer = deque(maxlen=buffer_size)
        self.raw_sensor_buffer = deque(maxlen=buffer_size * 100)  # Store raw samples
        self.features_buffer = {i: deque(maxlen=buffer_size) for i in range(8)}
        self.prediction_buffer = deque(maxlen=buffer_size)
        self.probability_buffer = deque(maxlen=buffer_size)
        self.has_raw_data = False
        
        # Statistics
        self.start_time = None
        self.sample_count = 0
        self.slip_count = 0
        self.no_slip_count = 0
        
        # Serial connection
        self.serial_conn = None
        
        # Setup plots
        self.setup_plots()
        
    def setup_plots(self):
        """Setup matplotlib figure and subplots"""
        self.fig = plt.figure(figsize=(16, 10))
        self.fig.suptitle('Piezoelectric Sensor Data Visualization', fontsize=14, fontweight='bold')
        
        # Create subplots
        gs = self.fig.add_gridspec(4, 2, hspace=0.3, wspace=0.3)
        
        # Plot 1: Raw sensor signal
        self.ax1 = self.fig.add_subplot(gs[0, :])
        self.ax1.set_title('Raw Sensor Signal (Windowed Mean)', fontweight='bold')
        self.ax1.set_xlabel('Time (s)')
        self.ax1.set_ylabel('Amplitude (0-1023)')
        self.ax1.grid(True, alpha=0.3)
        self.line1, = self.ax1.plot([], [], 'b-', linewidth=1.5, label='Window Mean')
        self.line1_raw = None  # Will be created if raw data is available
        self.ax1.legend()
        self.ax1.set_ylim(0, 1023)
        
        # Plot 2: Extracted Features (first 4)
        self.ax2 = self.fig.add_subplot(gs[1, 0])
        self.ax2.set_title('Features 0-3 (Spectral)', fontweight='bold')
        self.ax2.set_xlabel('Time (s)')
        self.ax2.set_ylabel('Normalized Value')
        self.ax2.grid(True, alpha=0.3)
        self.lines2 = [self.ax2.plot([], [], label=f'Feature {i}')[0] for i in range(4)]
        self.ax2.legend(fontsize=8)
        self.ax2.set_ylim(0, 1)
        
        # Plot 3: Extracted Features (last 4 - temporal)
        self.ax3 = self.fig.add_subplot(gs[1, 1])
        self.ax3.set_title('Features 4-7 (Temporal)', fontweight='bold')
        self.ax3.set_xlabel('Time (s)')
        self.ax3.set_ylabel('Normalized Value')
        self.ax3.grid(True, alpha=0.3)
        self.lines3 = [self.ax3.plot([], [], label=f'Feature {i+4}')[0] for i in range(4)]
        self.ax3.legend(fontsize=8)
        self.ax3.set_ylim(0, 1)
        
        # Plot 4: Slip Prediction
        self.ax4 = self.fig.add_subplot(gs[2, :])
        self.ax4.set_title('Slip Detection Prediction', fontweight='bold')
        self.ax4.set_xlabel('Time (s)')
        self.ax4.set_ylabel('Prediction')
        self.ax4.grid(True, alpha=0.3)
        self.line4, = self.ax4.plot([], [], 'r-', linewidth=2, label='Slip (1) / No Slip (0)')
        self.line5, = self.ax4.plot([], [], 'g--', linewidth=1, alpha=0.5, label='Probability')
        self.ax4.legend()
        self.ax4.set_ylim(-0.1, 1.1)
        self.ax4.set_yticks([0, 1])
        self.ax4.set_yticklabels(['No Slip', 'Slip'])
        
        # Plot 5: Statistics text
        self.ax5 = self.fig.add_subplot(gs[3, :])
        self.ax5.axis('off')
        self.stats_text = self.ax5.text(0.1, 0.5, '', fontsize=12, family='monospace',
                                        verticalalignment='center', transform=self.ax5.transAxes)
        
    def connect_serial(self):
        """Connect to Arduino via serial"""
        try:
            self.serial_conn = serial.Serial(self.port, self.baud_rate, timeout=1)
            print(f"Connected to {self.port} at {self.baud_rate} baud")
            return True
        except serial.SerialException as e:
            print(f"Error connecting to serial port: {e}")
            print("\nAvailable ports:")
            for port in serial.tools.list_ports.comports():
                print(f"  - {port.device}: {port.description}")
            return False
    
    def parse_serial_line(self, line):
        """Parse a line from serial output"""
        # Expected format: "RAW: [...] | Features: [0.xxx, 0.xxx, ...] | Prediction: SLIP/NO_SLIP | Probability: 0.xxx"
        try:
            # Extract raw data if available
            raw_data = None
            raw_match = re.search(r'RAW:\s*\[(.*?)\]', line)
            if raw_match:
                raw_str = raw_match.group(1)
                raw_data = [int(x.strip()) for x in raw_str.split(',')]
                self.has_raw_data = True
            
            # Extract features
            features_match = re.search(r'Features:\s*\[(.*?)\]', line)
            if features_match:
                features_str = features_match.group(1)
                features = [float(x.strip()) for x in features_str.split(',')]
            else:
                return None
            
            # Extract prediction
            pred_match = re.search(r'Prediction:\s*(SLIP|NO_SLIP)', line)
            if pred_match:
                prediction = 1 if pred_match.group(1) == 'SLIP' else 0
            else:
                return None
            
            # Extract probability
            prob_match = re.search(r'Probability:\s*([\d.]+)', line)
            if prob_match:
                probability = float(prob_match.group(1))
            else:
                probability = 0.0
            
            return {
                'raw_data': raw_data,
                'features': features,
                'prediction': prediction,
                'probability': probability
            }
        except Exception as e:
            print(f"Error parsing line: {e}")
            return None
    
    def update_data(self, frame):
        """Update data from serial and refresh plots"""
        if self.serial_conn:
            try:
                # Read all available lines
                while self.serial_conn.in_waiting:
                    line = self.serial_conn.readline().decode('utf-8', errors='ignore').strip()
                    
                    # Also check for initialization messages
                    if 'Initialized' in line or 'Ready' in line:
                        print(f"Arduino: {line}")
                        continue
                    
                    if line and 'Features:' in line:
                        data = self.parse_serial_line(line)
                        if data:
                            # Update time
                            if self.start_time is None:
                                self.start_time = datetime.now()
                            current_time = (datetime.now() - self.start_time).total_seconds()
                            
                            # Store data
                            self.time_buffer.append(current_time)
                            self.prediction_buffer.append(data['prediction'])
                            self.probability_buffer.append(data['probability'])
                            
                            # Store features
                            for i, feat_val in enumerate(data['features']):
                                if i < 8:
                                    self.features_buffer[i].append(feat_val)
                            
                            # Store raw sensor data if available
                            if data['raw_data']:
                                for raw_val in data['raw_data']:
                                    self.raw_sensor_buffer.append(raw_val)
                                # Use mean of raw window as sensor value
                                self.sensor_buffer.append(np.mean(data['raw_data']))
                            else:
                                # Fallback: use feature 0 (peak amplitude) as proxy
                                if len(self.features_buffer[0]) > 0:
                                    sensor_proxy = self.features_buffer[0][-1] * 511.5
                                    self.sensor_buffer.append(sensor_proxy)
                            
                            # Update statistics
                            self.sample_count += 1
                            if data['prediction'] == 1:
                                self.slip_count += 1
                            else:
                                self.no_slip_count += 1
                            
                            # Debug: print first few samples and periodic updates
                            if self.sample_count <= 3:
                                print(f"Sample {self.sample_count}: Features={data['features'][:3]}, Prediction={data['prediction']}")
                            elif self.sample_count % 10 == 0:
                                print(f"Received {self.sample_count} samples...")
                            break  # Process one line per update cycle
                        
            except Exception as e:
                if self.sample_count == 0:  # Only print errors if we haven't received any data yet
                    print(f"Error reading serial: {e}")
                import traceback
                if self.sample_count == 0:
                    traceback.print_exc()
        
        # Always update plots (even if no new data, to refresh display)
        self.update_plots()
        
        return []
    
    def update_plots(self):
        """Update all plot data"""
        if len(self.time_buffer) == 0:
            return
        
        # Debug: print when plots are being updated
        if self.sample_count > 0 and self.sample_count % 10 == 0:
            print(f"Updating plots with {len(self.time_buffer)} data points...")
        
        time_array = np.array(self.time_buffer)
        
        # Plot 1: Raw sensor signal
        if len(self.sensor_buffer) > 0:
            sensor_array = np.array(self.sensor_buffer)
            time_sensor = time_array[-len(sensor_array):]
            self.line1.set_data(time_sensor, sensor_array)
            
            # Plot raw samples if available
            if self.has_raw_data and len(self.raw_sensor_buffer) > 0:
                raw_array = np.array(self.raw_sensor_buffer)
                # Create time array for raw samples (assuming they come in windows)
                if len(time_array) > 0:
                    window_duration = 1.0  # Approximate window duration
                    raw_time = np.linspace(max(0, time_array[-1] - 10), time_array[-1], len(raw_array))
                    # Only show recent raw data
                    recent_mask = raw_time >= max(0, time_array[-1] - 10)
                    if self.line1_raw is None:
                        self.line1_raw, = self.ax1.plot([], [], 'r-', linewidth=0.5, alpha=0.3, label='Raw Samples')
                        self.ax1.legend()
                    self.line1_raw.set_data(raw_time[recent_mask], raw_array[recent_mask])
            
            # Update axes limits
            self.ax1.set_xlim(max(0, time_array[-1] - 10), time_array[-1] + 1)
            if len(sensor_array) > 0:
                y_min, y_max = sensor_array.min(), sensor_array.max()
                y_range = max(y_max - y_min, 50)  # Minimum range of 50
                self.ax1.set_ylim(max(0, y_min - y_range * 0.1), min(1023, y_max + y_range * 0.1))
        
        # Plot 2: Features 0-3
        for i in range(4):
            if len(self.features_buffer[i]) > 0:
                feat_array = np.array(self.features_buffer[i])
                time_feat = time_array[-len(feat_array):]
                self.lines2[i].set_data(time_feat, feat_array)
        if len(time_array) > 0:
            self.ax2.set_xlim(max(0, time_array[-1] - 10), time_array[-1] + 1)
            self.ax2.relim()
            self.ax2.autoscale_view()
        
        # Plot 3: Features 4-7
        for i in range(4):
            if len(self.features_buffer[i+4]) > 0:
                feat_array = np.array(self.features_buffer[i+4])
                time_feat = time_array[-len(feat_array):]
                self.lines3[i].set_data(time_feat, feat_array)
        if len(time_array) > 0:
            self.ax3.set_xlim(max(0, time_array[-1] - 10), time_array[-1] + 1)
            self.ax3.relim()
            self.ax3.autoscale_view()
        
        # Plot 4: Predictions
        if len(self.prediction_buffer) > 0:
            pred_array = np.array(self.prediction_buffer)
            prob_array = np.array(self.probability_buffer)
            time_pred = time_array[-len(pred_array):]
            self.line4.set_data(time_pred, pred_array)
            self.line5.set_data(time_pred, prob_array)
            self.ax4.set_xlim(max(0, time_array[-1] - 10), time_array[-1] + 1)
            self.ax4.relim()
            self.ax4.autoscale_view()
        
        # Plot 5: Statistics
        # Format features safely
        feature_strs = []
        for i in range(8):
            if len(self.features_buffer[i]) > 0:
                feature_strs.append(f'{self.features_buffer[i][-1]:.3f}')
            else:
                feature_strs.append('N/A')
        
        # Calculate slip rate safely
        slip_rate = (self.slip_count / self.sample_count * 100) if self.sample_count > 0 else 0.0
        
        # Get current prediction and probability safely
        current_pred = 'SLIP' if (len(self.prediction_buffer) > 0 and self.prediction_buffer[-1] == 1) else 'NO_SLIP'
        current_prob = self.probability_buffer[-1] if len(self.probability_buffer) > 0 else 0.0
        
        stats_str = f"""Statistics:
Samples: {self.sample_count}  |  Runtime: {time_array[-1]:.1f}s
Slip Detections: {self.slip_count}  |  No Slip: {self.no_slip_count}
Slip Rate: {slip_rate:.1f}%
Current Features: {feature_strs}
Current Prediction: {current_pred}  |  Probability: {current_prob:.3f}
"""
        self.stats_text.set_text(stats_str)
        
        # Force redraw - use multiple methods to ensure update
        try:
            self.fig.canvas.draw_idle()
            self.fig.canvas.flush_events()
            # Small pause to allow GUI to update
            plt.pause(0.001)
        except:
            pass
    
    def run(self):
        """Start the visualization"""
        if not self.connect_serial():
            return
        
        print("Starting visualization...")
        print("Close the plot window to stop.")
        print("Waiting for data from Arduino...")
        
        # Start animation - update every 50ms for smoother updates
        ani = animation.FuncAnimation(
            self.fig, 
            self.update_data, 
            interval=50, 
            blit=False, 
            cache_frame_data=False,
            repeat=True
        )
        
        # Store animation reference to prevent garbage collection
        self.ani = ani
        
        try:
            plt.ion()  # Turn on interactive mode
            plt.show(block=True)
        except KeyboardInterrupt:
            print("\nStopping visualization...")
        finally:
            if self.serial_conn:
                self.serial_conn.close()
                print("Serial connection closed.")


def main():
    parser = argparse.ArgumentParser(description='Visualize piezoelectric sensor data from Arduino')
    parser.add_argument('--port', type=str, default='COM3',
                       help='Serial port (e.g., COM3 on Windows, /dev/ttyUSB0 on Linux)')
    parser.add_argument('--baud', type=int, default=9600,
                       help='Baud rate (default: 9600)')
    parser.add_argument('--buffer', type=int, default=500,
                       help='Buffer size for data points (default: 500)')
    
    args = parser.parse_args()
    
    visualizer = SensorVisualizer(port=args.port, baud_rate=args.baud, buffer_size=args.buffer)
    visualizer.run()


if __name__ == '__main__':
    main()