visualize_raw_sensor.py

"""
Visualize raw piezoelectric sensor data from Arduino
Simple real-time plotting of raw sensor values
"""

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

class RawSensorVisualizer:
    def __init__(self, port='COM3', baud_rate=9600, buffer_size=1000):
        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)
        
        # Statistics
        self.start_time = None
        self.sample_count = 0
        self.serial_conn = None
        
        # Setup plot
        self.setup_plot()
    
    def setup_plot(self):
        """Setup matplotlib figure"""
        self.fig, (self.ax1, self.ax2) = plt.subplots(2, 1, figsize=(12, 8))
        self.fig.suptitle('Raw Piezoelectric Sensor Output', fontsize=14, fontweight='bold')
        
        # Plot 1: Time series
        self.ax1.set_title('Sensor Reading Over Time', fontweight='bold')
        self.ax1.set_xlabel('Time (s)')
        self.ax1.set_ylabel('Raw Value (0-1023)')
        self.ax1.grid(True, alpha=0.3)
        self.line1, = self.ax1.plot([], [], 'b-', linewidth=1.5, label='Raw Sensor')
        self.ax1.legend()
        self.ax1.set_ylim(0, 1023)
        
        # Plot 2: Recent samples (zoomed)
        self.ax2.set_title('Recent Samples (Last 5 seconds)', fontweight='bold')
        self.ax2.set_xlabel('Time (s)')
        self.ax2.set_ylabel('Raw Value (0-1023)')
        self.ax2.grid(True, alpha=0.3)
        self.line2, = self.ax2.plot([], [], 'r-', linewidth=1.5)
        self.ax2.set_ylim(0, 1023)
        
        # Statistics text
        self.stats_text = self.fig.text(0.02, 0.02, '', fontsize=10, family='monospace',
                                        verticalalignment='bottom')
    
    def connect(self):
        """Connect to Arduino"""
        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: {e}")
            print("\nAvailable ports:")
            for p in serial.tools.list_ports.comports():
                print(f"  - {p.device}: {p.description}")
            return False
    
    def parse_line(self, line):
        """Parse serial line - supports both continuous and batch modes"""
        try:
            # Check for batch format: BATCH:[...]
            if 'BATCH:[' in line:
                batch_match = line.split('BATCH:[')[1].split(']')[0]
                values = [int(x.strip()) for x in batch_match.split(',')]
                return values
            
            # Check for continuous format: timestamp,value
            elif ',' in line:
                parts = line.strip().split(',')
                if len(parts) == 2:
                    timestamp = int(parts[0])
                    value = int(parts[1])
                    return [(timestamp, value)]
            
            # Single value format
            else:
                value = int(line.strip())
                return [value]
        except:
            return None
    
    def update(self, frame):
        """Update function called by animation"""
        from datetime import datetime
        
        if self.serial_conn and self.serial_conn.in_waiting:
            line = self.serial_conn.readline().decode('utf-8', errors='ignore').strip()
            
            if line:
                data = self.parse_line(line)
                if data:
                    # Handle batch data
                    if isinstance(data, list) and len(data) > 0:
                        if isinstance(data[0], tuple):
                            # Continuous mode: (timestamp, value)
                            for timestamp, value in data:
                                self.add_sample(timestamp / 1000.0, value)
                        elif isinstance(data[0], int) and len(data) == 1:
                            # Single value
                            if self.start_time is None:
                                self.start_time = datetime.now()
                            current_time = (datetime.now() - self.start_time).total_seconds()
                            self.add_sample(current_time, data[0])
                        else:
                            # Batch mode: list of values
                            if self.start_time is None:
                                self.start_time = datetime.now()
                            base_time = (datetime.now() - self.start_time).total_seconds()
                            sample_rate = 100  # Adjust based on your Arduino sample rate
                            for i, value in enumerate(data):
                                sample_time = base_time - (len(data) - i) / sample_rate
                                self.add_sample(sample_time, value)
        
        # Update plots
        self.update_plots()
        
        return [self.line1, self.line2]
    
    def add_sample(self, time_val, sensor_val):
        """Add a sample to buffers"""
        self.time_buffer.append(time_val)
        self.sensor_buffer.append(sensor_val)
        self.sample_count += 1
    
    def update_plots(self):
        """Update plot data"""
        if len(self.time_buffer) == 0:
            return
        
        time_array = np.array(self.time_buffer)
        sensor_array = np.array(self.sensor_buffer)
        
        # Plot 1: Full time series
        self.line1.set_data(time_array, sensor_array)
        if len(time_array) > 0:
            self.ax1.set_xlim(max(0, time_array[0]), 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)
                self.ax1.set_ylim(max(0, y_min - y_range * 0.1), min(1023, y_max + y_range * 0.1))
        
        # Plot 2: Recent samples (last 5 seconds)
        if len(time_array) > 0:
            recent_mask = time_array >= max(0, time_array[-1] - 5)
            recent_time = time_array[recent_mask]
            recent_sensor = sensor_array[recent_mask]
            
            self.line2.set_data(recent_time, recent_sensor)
            if len(recent_time) > 0:
                self.ax2.set_xlim(max(0, recent_time[0]), recent_time[-1] + 0.5)
                if len(recent_sensor) > 0:
                    y_min, y_max = recent_sensor.min(), recent_sensor.max()
                    y_range = max(y_max - y_min, 50)
                    self.ax2.set_ylim(max(0, y_min - y_range * 0.1), min(1023, y_max + y_range * 0.1))
        
        # Update statistics
        if len(sensor_array) > 0:
            stats = f"""Statistics:
Samples: {self.sample_count}  |  Runtime: {time_array[-1]:.1f}s
Current Value: {sensor_array[-1]}  |  Min: {sensor_array.min()}  |  Max: {sensor_array.max()}
Mean: {sensor_array.mean():.1f}  |  Std Dev: {sensor_array.std():.1f}
Sample Rate: {len(sensor_array) / time_array[-1]:.1f} Hz (if time_array[-1] > 0 else 0)
"""
            self.stats_text.set_text(stats)
        
        # Force redraw
        self.fig.canvas.draw()
    
    def run(self):
        """Start visualization"""
        if not self.connect():
            return
        
        print("Starting visualization...")
        print("Close the plot window to stop.")
        
        ani = animation.FuncAnimation(
            self.fig,
            self.update,
            interval=50,
            blit=False,
            cache_frame_data=False
        )
        
        try:
            plt.show(block=True)
        except KeyboardInterrupt:
            print("\nStopping...")
        finally:
            if self.serial_conn:
                self.serial_conn.close()
                print("Serial connection closed.")


def main():
    parser = argparse.ArgumentParser(description='Visualize raw piezoelectric sensor data')
    parser.add_argument('--port', type=str, default='COM3', help='Serial port')
    parser.add_argument('--baud', type=int, default=9600, help='Baud rate')
    
    args = parser.parse_args()
    
    visualizer = RawSensorVisualizer(port=args.port, baud_rate=args.baud)
    visualizer.run()


if __name__ == '__main__':
    main()