serial_controller_reader.py

#!/usr/bin/env python3
"""
Unitree G1 Serial Controller Reader

Attempts to read controller data from the USB receiver.

Note: This may not produce data when the robot is off, as the receiver
might require the robot's control system to poll/initialize it.
"""

import serial
import serial.tools.list_ports
import struct
import time
import sys
import argparse
from typing import Optional

from controller_protocol import ControllerState, ControllerButtons


class SerialControllerReader:
    """
    Reads raw controller data from the USB serial receiver.
    """
    
    # Common baud rates to try
    BAUD_RATES = [921600, 115200, 460800, 256000, 500000, 57600, 38400, 19200, 9600]
    
    # Possible header patterns (based on common embedded protocols)
    HEADER_PATTERNS = [
        b'\xFE\xEF',  # Common Unitree pattern
        b'\xFE\xFE',
        b'\xAA\x55',
        b'\x55\xAA',
    ]
    
    def __init__(self, port: str = '/dev/ttyUSB0', baud: int = 921600):
        self.port = port
        self.baud = baud
        self.serial: Optional[serial.Serial] = None
        self.last_state = ControllerState()
        
    def find_controller_port(self) -> Optional[str]:
        """Auto-detect the controller receiver port"""
        ports = serial.tools.list_ports.comports()
        
        for port in ports:
            # Look for Silicon Labs CP210x
            if 'CP210' in (port.description or ''):
                print(f"Found CP210x at {port.device}: {port.description}")
                return port.device
            # Also check by VID:PID (10c4:ea60 is Silicon Labs)
            if port.vid == 0x10c4 and port.pid == 0xea60:
                print(f"Found Silicon Labs device at {port.device}")
                return port.device
                
        return None
    
    def open(self) -> bool:
        """Open the serial connection"""
        try:
            self.serial = serial.Serial(
                self.port,
                self.baud,
                timeout=0.1,
                bytesize=serial.EIGHTBITS,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE
            )
            self.serial.reset_input_buffer()
            return True
        except serial.SerialException as e:
            print(f"Error opening {self.port}: {e}")
            return False
    
    def close(self):
        """Close the serial connection"""
        if self.serial and self.serial.is_open:
            self.serial.close()
    
    def read_raw(self, size: int = 40) -> Optional[bytes]:
        """Read raw bytes from serial"""
        if not self.serial or not self.serial.is_open:
            return None
        return self.serial.read(size)
    
    def find_packet(self, buffer: bytes) -> Optional[bytes]:
        """
        Try to find a valid 40-byte packet in buffer.
        Looks for known header patterns.
        """
        for pattern in self.HEADER_PATTERNS:
            idx = buffer.find(pattern)
            if idx >= 0 and idx + 40 <= len(buffer):
                return buffer[idx:idx + 40]
        return None
    
    def try_parse_packet(self, data: bytes) -> Optional[ControllerState]:
        """Attempt to parse data as controller state"""
        if len(data) < 40:
            return None
        
        try:
            state = ControllerState.from_bytes(data)
            
            # Validate: joystick values should be in reasonable range
            if all(abs(v) <= 2.0 for v in [state.lx, state.ly, state.rx, state.ry]):
                return state
        except Exception:
            pass
        
        return None
    
    def scan_baud_rates(self) -> Optional[int]:
        """Try different baud rates to find one that works"""
        print("Scanning baud rates...")
        
        for baud in self.BAUD_RATES:
            print(f"  Trying {baud}...", end=' ')
            
            try:
                if self.serial and self.serial.is_open:
                    self.serial.close()
                
                self.serial = serial.Serial(self.port, baud, timeout=1)
                self.serial.reset_input_buffer()
                
                time.sleep(0.5)
                data = self.serial.read(100)
                
                if data:
                    print(f"Got {len(data)} bytes!")
                    return baud
                else:
                    print("No data")
                    
            except Exception as e:
                print(f"Error: {e}")
        
        return None


def monitor_controller(reader: SerialControllerReader, duration: float = 0.0):
    """
    Monitor controller and print state changes.
    
    Args:
        reader: SerialControllerReader instance
        duration: How long to monitor (0 = forever)
    """
    print("\n=== Controller Monitor ===")
    print("Move joysticks or press buttons...")
    print("Press Ctrl+C to exit\n")
    
    start = time.time()
    buffer = b''
    last_print = time.time()
    
    try:
        while True:
            # Check duration
            if duration > 0 and time.time() - start > duration:
                break
            
            # Read data
            data = reader.read_raw(64)
            
            if data:
                buffer += data
                
                # Print raw data periodically
                if time.time() - last_print > 0.5:
                    print(f"[{len(buffer):4d} bytes] {data[:20].hex()}...")
                    last_print = time.time()
                
                # Try to find and parse packet
                packet = reader.find_packet(buffer)
                if packet:
                    state = reader.try_parse_packet(packet)
                    if state:
                        print("\n" + str(state) + "\n")
                        buffer = b''
                
                # Prevent buffer from growing too large
                if len(buffer) > 1000:
                    buffer = buffer[-100:]
            
            time.sleep(0.01)
            
    except KeyboardInterrupt:
        print("\nStopped.")


def send_probe(reader: SerialControllerReader):
    """Send probe bytes to try to wake up the receiver"""
    probes = [
        b'\x00' * 4,
        b'\xFF' * 4,
        b'\xFE\xEF\x00\x00',
        b'\xAA\x55\x00\x00',
        b'AT\r\n',
    ]
    
    print("\nSending probes...")
    for probe in probes:
        print(f"  Sending: {probe.hex()}", end=' ')
        reader.serial.write(probe)
        time.sleep(0.3)
        response = reader.read_raw(100)
        if response:
            print(f"-> Got: {response.hex()}")
        else:
            print("-> No response")


def main():
    parser = argparse.ArgumentParser(description='Unitree Controller Serial Reader')
    parser.add_argument('-p', '--port', default='/dev/ttyUSB0', help='Serial port')
    parser.add_argument('-b', '--baud', type=int, default=921600, help='Baud rate')
    parser.add_argument('-s', '--scan', action='store_true', help='Scan for working baud rate')
    parser.add_argument('-t', '--time', type=float, default=0, help='Monitor duration (0=forever)')
    parser.add_argument('--probe', action='store_true', help='Send probe bytes')
    parser.add_argument('--find', action='store_true', help='Auto-find controller port')
    
    args = parser.parse_args()
    
    reader = SerialControllerReader(args.port, args.baud)
    
    # Auto-find port if requested
    if args.find:
        port = reader.find_controller_port()
        if port:
            reader.port = port
        else:
            print("Could not auto-detect controller port")
            return 1
    
    # Open port
    if not reader.open():
        return 1
    
    print(f"Opened {reader.port} at {reader.baud} baud")
    
    try:
        # Scan baud rates if requested
        if args.scan:
            baud = reader.scan_baud_rates()
            if baud:
                print(f"\nFound working baud rate: {baud}")
                reader.baud = baud
            else:
                print("\nNo baud rate produced data")
        
        # Send probes if requested
        if args.probe:
            send_probe(reader)
        
        # Monitor controller
        monitor_controller(reader, args.time)
        
    finally:
        reader.close()
    
    return 0


if __name__ == "__main__":
    sys.exit(main())