controller_reader.py

#!/usr/bin/env python3
"""
Cross-platform Xbox Series S Controller Input Reader
Reads controller inputs and prints them to console in real-time
"""

import sys
import struct
import os
import time
from datetime import datetime
from dataclasses import dataclass, asdict
from typing import Dict
import json

try:
    import pygame
except ImportError:
    print("Error: 'pygame' library not found.")
    print("Install with: pip install pygame")
    sys.exit(1)


@dataclass
class ControllerState:
    """
    Controller state that can be easily serialized for MAVLink transmission.
    All values are normalized where applicable.
    """
    left_stick_x: float = 0.0
    left_stick_y: float = 0.0
    right_stick_x: float = 0.0
    right_stick_y: float = 0.0
    left_trigger: float = 0.0
    right_trigger: float = 0.0
    dpad_up: int = 0
    dpad_down: int = 0
    dpad_left: int = 0
    dpad_right: int = 0
    button_a: int = 0
    button_b: int = 0
    button_x: int = 0
    button_y: int = 0
    left_bumper: int = 0
    right_bumper: int = 0
    button_view: int = 0
    button_menu: int = 0
    button_share: int = 0
    button_home: int = 0
    left_stick_click: int = 0
    right_stick_click: int = 0
    
    def to_dict(self) -> Dict:
        return asdict(self)
    
    def to_json(self) -> str:
        return json.dumps(self.to_dict())
    
    def to_bytes(self) -> bytes:
        """Pack controller state into 40 bytes for MAVLink transmission"""
        packed = struct.pack(
            '<6f16B',  # 6 floats (24 bytes) + 16 buttons (16 bytes) = 40 bytes
            self.left_stick_x,
            self.left_stick_y,
            self.right_stick_x,
            self.right_stick_y,
            self.left_trigger,
            self.right_trigger,
            self.dpad_up,
            self.dpad_down,
            self.dpad_left,
            self.dpad_right,
            self.button_a,
            self.button_b,
            self.button_x,
            self.button_y,
            self.left_bumper,
            self.right_bumper,
            self.button_view,
            self.button_menu,
            self.button_share,
            self.button_home,
            self.left_stick_click,
            self.right_stick_click
        )
        return packed
    
    @classmethod
    def from_bytes(cls, data: bytes) -> 'ControllerState':
        """Unpack controller state from bytes"""
        unpacked = struct.unpack('<6f16B', data)
        return cls(
            left_stick_x=unpacked[0],
            left_stick_y=unpacked[1],
            right_stick_x=unpacked[2],
            right_stick_y=unpacked[3],
            left_trigger=unpacked[4],
            right_trigger=unpacked[5],
            dpad_up=unpacked[6],
            dpad_down=unpacked[7],
            dpad_left=unpacked[8],
            dpad_right=unpacked[9],
            button_a=unpacked[10],
            button_b=unpacked[11],
            button_x=unpacked[12],
            button_y=unpacked[13],
            left_bumper=unpacked[14],
            right_bumper=unpacked[15],
            button_view=unpacked[16],
            button_menu=unpacked[17],
            button_share=unpacked[18],
            button_home=unpacked[19],
            left_stick_click=unpacked[20],
            right_stick_click=unpacked[21]
        )


class DataLogger:
    """Logs controller data to files with millisecond timestamps"""
    
    def __init__(self, log_dir="controller_logs"):
        self.log_dir = log_dir
        os.makedirs(log_dir, exist_ok=True)
        
        # Find next sequence number
        sequence_num = self._get_next_sequence_number()
        
        # Create filenames with sequence number and timestamp
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.json_log_path = os.path.join(log_dir, f"log_{sequence_num:04d}_{timestamp}.jsonl")
        self.mavlink_log_path = os.path.join(log_dir, f"log_{sequence_num:04d}_{timestamp}.bin")
        
        # Open files
        self.json_file = open(self.json_log_path, 'w')
        self.mavlink_file = open(self.mavlink_log_path, 'wb')
        
        # Track start time for relative timestamps
        self.start_time = time.time()
        
        print(f"Logging session #{sequence_num}:")
        print(f"  JSON: {self.json_log_path}")
        print(f"  MAVLink: {self.mavlink_log_path}")
        print()
    
    def _get_next_sequence_number(self):
        """Find the next sequence number by checking existing log files"""
        if not os.path.exists(self.log_dir):
            return 1
        
        existing_files = os.listdir(self.log_dir)
        sequence_numbers = []
        
        for filename in existing_files:
            if filename.startswith("log_") and "_" in filename:
                try:
                    # Extract sequence number from filename like "log_0001_..."
                    seq_str = filename.split("_")[1]
                    if seq_str.isdigit():
                        sequence_numbers.append(int(seq_str))
                except (IndexError, ValueError):
                    continue
        
        if sequence_numbers:
            return max(sequence_numbers) + 1
        return 1
    
    def log(self, state: ControllerState):
        """Log controller state to both files"""
        # Get timestamp in milliseconds since start
        current_time = time.time()
        timestamp_ms = int((current_time - self.start_time) * 1000)
        
        # Create JSON log entry with timestamp
        log_entry = {
            "timestamp_ms": timestamp_ms,
            "timestamp_iso": datetime.now().isoformat(),
            "state": state.to_dict()
        }
        
        # Write JSON line
        self.json_file.write(json.dumps(log_entry) + '\n')
        self.json_file.flush()  # Ensure data is written immediately
        
        # Write MAVLink binary data with timestamp header
        # Format: timestamp (8 bytes, uint64) + state data (40 bytes)
        mavlink_packet = struct.pack('<Q', timestamp_ms) + state.to_bytes()
        self.mavlink_file.write(mavlink_packet)
        self.mavlink_file.flush()  # Ensure data is written immediately
    
    def close(self):
        """Close log files"""
        self.json_file.close()
        self.mavlink_file.close()
        print(f"\nLogs saved:")
        print(f"  JSON: {self.json_log_path}")
        print(f"  MAVLink: {self.mavlink_log_path}")


class XboxController:
    """Xbox Controller Input Handler"""
    
    def __init__(self):
        pygame.init()
        pygame.joystick.init()
        
        joystick_count = pygame.joystick.get_count()
        if joystick_count == 0:
            raise RuntimeError("No controller detected!")
        
        self.joystick = pygame.joystick.Joystick(0)
        self.joystick.init()
        
        print(f"Controller: {self.joystick.get_name()}")
        print(f"Axes: {self.joystick.get_numaxes()}")
        print(f"Buttons: {self.joystick.get_numbuttons()}")
        print(f"Hats: {self.joystick.get_numhats()}")
        print()
        
        self.state = ControllerState()
        self.deadzone = 0.1
    
    def apply_deadzone(self, value: float) -> float:
        if abs(value) < self.deadzone:
            return 0.0
        return value
    
    def normalize_trigger(self, value: float) -> float:
        """Normalize trigger from -1.0~1.0 to 0.0~1.0"""
        return max(0.0, min(1.0, (value + 1.0) / 2.0))
    
    def update(self):
        """Poll controller and update state"""
        pygame.event.pump()
        
        # Read analog sticks
        if self.joystick.get_numaxes() >= 4:
            self.state.left_stick_x = self.apply_deadzone(self.joystick.get_axis(0))
            self.state.left_stick_y = self.apply_deadzone(self.joystick.get_axis(1))
            self.state.right_stick_x = self.apply_deadzone(self.joystick.get_axis(2))
            self.state.right_stick_y = self.apply_deadzone(self.joystick.get_axis(3))
        
        # Read triggers
        if self.joystick.get_numaxes() >= 6:
            self.state.left_trigger = self.normalize_trigger(self.joystick.get_axis(4))
            self.state.right_trigger = self.normalize_trigger(self.joystick.get_axis(5))
        
        # Read buttons - Xbox Series S on macOS specific mapping
        if self.joystick.get_numbuttons() >= 16:
            # Face buttons (standard)
            self.state.button_a = 1 if self.joystick.get_button(0) else 0
            self.state.button_b = 1 if self.joystick.get_button(1) else 0
            self.state.button_x = 1 if self.joystick.get_button(2) else 0
            self.state.button_y = 1 if self.joystick.get_button(3) else 0
            
            # Bumpers
            self.state.left_bumper = 1 if self.joystick.get_button(9) else 0
            self.state.right_bumper = 1 if self.joystick.get_button(10) else 0
            
            # Special buttons - Xbox Series S specific
            self.state.button_view = 1 if self.joystick.get_button(4) else 0
            self.state.button_menu = 1 if self.joystick.get_button(6) else 0
            self.state.button_share = 1 if self.joystick.get_button(15) else 0
            self.state.button_home = 1 if self.joystick.get_button(5) else 0
            
            # Stick clicks
            self.state.left_stick_click = 1 if self.joystick.get_button(7) else 0
            self.state.right_stick_click = 1 if self.joystick.get_button(8) else 0
            
            # D-pad as buttons (not hat on this controller)
            self.state.dpad_up = 1 if self.joystick.get_button(11) else 0
            self.state.dpad_down = 1 if self.joystick.get_button(12) else 0
            self.state.dpad_left = 1 if self.joystick.get_button(13) else 0
            self.state.dpad_right = 1 if self.joystick.get_button(14) else 0
    
    def get_state(self) -> ControllerState:
        return self.state
    
    def get_raw_inputs(self):
        """Get all raw inputs for debugging"""
        pygame.event.pump()
        
        raw = {
            'axes': [],
            'buttons': [],
            'hats': []
        }
        
        for i in range(self.joystick.get_numaxes()):
            raw['axes'].append(self.joystick.get_axis(i))
        
        for i in range(self.joystick.get_numbuttons()):
            raw['buttons'].append(self.joystick.get_button(i))
        
        for i in range(self.joystick.get_numhats()):
            raw['hats'].append(self.joystick.get_hat(i))
        
        return raw


def clear_screen():
    """Clear terminal screen"""
    os.system('clear' if os.name != 'nt' else 'cls')


def print_state(state: ControllerState, verbose: bool = False, debug_mode: bool = False, raw_data=None, logger=None):
    """Print controller state in real-time (overwrites previous output)"""
    
    # Clear screen completely
    print("\033[2J\033[H", end="")
    
    print("=" * 70)
    print("XBOX CONTROLLER - REAL-TIME INPUT")
    if logger:
        print(f"[LOGGING ENABLED - Time: {int((time.time() - logger.start_time) * 1000)}ms]")
    print("=" * 70)
    
    # Analog Sticks
    left_click_text = "[CLICKED]" if state.left_stick_click else "         "
    right_click_text = "[CLICKED]" if state.right_stick_click else "         "
    
    print(f"\nLEFT STICK   X: {state.left_stick_x:>6.3f}  Y: {state.left_stick_y:>6.3f}  {left_click_text}")
    print(f"RIGHT STICK  X: {state.right_stick_x:>6.3f}  Y: {state.right_stick_y:>6.3f}  {right_click_text}")
    
    # Triggers
    print(f"\nLEFT TRIGGER:  [{('#' * int(state.left_trigger * 20)):20s}] {state.left_trigger:.3f}")
    print(f"RIGHT TRIGGER: [{('#' * int(state.right_trigger * 20)):20s}] {state.right_trigger:.3f}")
    
    # D-pad
    dpad_display = []
    if state.dpad_up: dpad_display.append("↑")
    if state.dpad_down: dpad_display.append("↓")
    if state.dpad_left: dpad_display.append("←")
    if state.dpad_right: dpad_display.append("→")
    print(f"\nD-PAD: {' '.join(dpad_display) if dpad_display else '---'}    ")
    
    # Face Buttons
    buttons = []
    if state.button_a: buttons.append("A")
    if state.button_b: buttons.append("B")
    if state.button_x: buttons.append("X")
    if state.button_y: buttons.append("Y")
    print(f"FACE BUTTONS: {' '.join(buttons) if buttons else '---'}    ")
    
    # Shoulder Buttons
    shoulders = []
    if state.left_bumper: shoulders.append("LB")
    if state.right_bumper: shoulders.append("RB")
    print(f"BUMPERS: {' '.join(shoulders) if shoulders else '---'}    ")
    
    # Special Buttons
    special = []
    if state.button_view: special.append("VIEW")
    if state.button_menu: special.append("MENU")
    if state.button_share: special.append("SHARE")
    if state.button_home: special.append("HOME")
    print(f"SPECIAL: {' '.join(special) if special else '---'}    ")
    
    if debug_mode and raw_data:
        print("\n" + "-" * 70)
        print("DEBUG - RAW INPUTS")
        print("-" * 70)
        print(f"Axes: {[f'{i}:{v:.2f}' for i, v in enumerate(raw_data['axes'])]}")
        print(f"Buttons pressed: {[f'{i}:{v}' for i, v in enumerate(raw_data['buttons']) if v]}")
        print(f"All buttons: {raw_data['buttons']}")
        print(f"Hats: {raw_data['hats']}")
    
    if verbose:
        print("\n" + "-" * 70)
        print("SERIALIZATION DATA")
        print("-" * 70)
        byte_data = state.to_bytes()
        print(f"Bytes (hex): {byte_data.hex()}")
        print(f"Size: {len(byte_data)} bytes")
    
    print("\n" + "=" * 70)
    print("Press Ctrl+C to exit")
    if logger:
        print("Use --log or -l to enable logging")
    else:
        print("Logging: DISABLED (use --log to enable)")
    print("Use --debug for raw mappings | --verbose for serialization data")
    print("=" * 70)


def main():
    """Main loop"""
    verbose = '--verbose' in sys.argv or '-v' in sys.argv
    debug = '--debug' in sys.argv or '-d' in sys.argv
    enable_logging = '--log' in sys.argv or '-l' in sys.argv
    
    print("Xbox Controller Input Reader")
    print("=" * 70)
    print("Initializing...")
    if enable_logging:
        print("Logging enabled!")
    print()
    
    logger = None
    
    try:
        controller = XboxController()
        
        # Initialize logger if requested
        if enable_logging:
            logger = DataLogger()
        
        # Clear screen and start
        clear_screen()
        
        clock = pygame.time.Clock()
        
        while True:
            controller.update()
            state = controller.get_state()
            
            # Log data if logging is enabled
            if logger:
                logger.log(state)
            
            raw_data = None
            if debug:
                raw_data = controller.get_raw_inputs()
            
            print_state(state, verbose=verbose, debug_mode=debug, raw_data=raw_data, logger=logger)
            
            # 30 updates per second
            clock.tick(30)
    
    except KeyboardInterrupt:
        print("\n\nExiting...")
        if logger:
            logger.close()
        pygame.quit()
        sys.exit(0)
    except RuntimeError as e:
        print(f"\nError: {e}")
        print("\nTroubleshooting:")
        print("1. Make sure Xbox controller is connected (Bluetooth or USB)")
        print("2. On macOS: System Settings > Privacy & Security > Input Monitoring")
        if logger:
            logger.close()
        pygame.quit()
        sys.exit(1)
    except Exception as e:
        print(f"\nUnexpected error: {e}")
        import traceback
        traceback.print_exc()
        if logger:
            logger.close()
        pygame.quit()
        sys.exit(1)


if __name__ == "__main__":
    main()