trikkeSensors/tools/capture_serial.py

#!/usr/bin/env python3
"""Capture Trikke sensor CSV output until interrupted with Ctrl-C."""

import argparse
import glob
import signal
import sys
from datetime import datetime
from pathlib import Path

import serial

EXPECTED_COLUMNS = [
    "sequence",
    "poll_timestamp_us",
    "accel_x_raw",
    "accel_y_raw",
    "accel_z_raw",
    "gyro_x_raw",
    "gyro_y_raw",
    "gyro_z_raw",
    "accel_x_mg",
    "accel_y_mg",
    "accel_z_mg",
    "gyro_x_mdps",
    "gyro_y_mdps",
    "gyro_z_mdps",
    "accel_native_x_raw",
    "accel_native_y_raw",
    "accel_native_z_raw",
    "gyro_native_x_raw",
    "gyro_native_y_raw",
    "gyro_native_z_raw",
    "accel_int_source",
    "gyro_status",
    "loop_overrun_count",
]

ACCEL_STATUS_INDEX = EXPECTED_COLUMNS.index("accel_int_source")
GYRO_STATUS_INDEX = EXPECTED_COLUMNS.index("gyro_status")
LOOP_OVERRUN_INDEX = EXPECTED_COLUMNS.index("loop_overrun_count")

ACCEL_DATA_READY_MASK = 0x80
ACCEL_OVERRUN_MASK = 0x01
GYRO_DATA_READY_MASK = 0x08
GYRO_OVERRUN_MASK = 0x80


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", help="serial port; auto-detected when omitted")
    parser.add_argument(
        "--baud",
        type=int,
        default=115200,
        help="serial API baud rate (ignored by USB Serial/JTAG)",
    )
    parser.add_argument("--output", type=Path)
    return parser.parse_args()


def resolve_port(requested_port: str | None) -> str:
    if requested_port:
        return requested_port

    candidates = sorted(glob.glob("/dev/cu.usbmodem*"))
    if not candidates:
        raise RuntimeError("no /dev/cu.usbmodem* serial device found")
    if len(candidates) > 1:
        joined = ", ".join(candidates)
        raise RuntimeError(f"multiple serial devices found ({joined}); specify --port")
    return candidates[0]


def main() -> int:
    args = parse_args()
    try:
        port = resolve_port(args.port)
    except RuntimeError as exc:
        print(f"Port error: {exc}", file=sys.stderr)
        return 2

    output = args.output or Path("captures") / datetime.now().strftime(
        "motion_%Y%m%d_%H%M%S.csv"
    )
    output.parent.mkdir(parents=True, exist_ok=True)

    stop_requested = False

    def request_stop(_signum: int, _frame: object) -> None:
        nonlocal stop_requested
        stop_requested = True

    signal.signal(signal.SIGINT, request_stop)
    signal.signal(signal.SIGTERM, request_stop)

    sample_count = 0
    missing_sequence_count = 0
    sequence_reset_count = 0
    timing_anomaly_count = 0
    ignored_line_count = 0
    rejected_record_count = 0
    accel_stale_count = 0
    accel_overrun_count = 0
    gyro_stale_count = 0
    gyro_overrun_count = 0
    final_loop_overrun_count = 0
    previous_sequence = None
    previous_timestamp_us = None
    print(f"Recording {port} to {output}; press Ctrl-C to stop", flush=True)

    try:
        with serial.Serial(port, args.baud, timeout=0.25) as sensor, output.open(
            "w", encoding="utf-8", newline=""
        ) as capture:
            capture.write(",".join(EXPECTED_COLUMNS) + "\n")

            while not stop_requested:
                raw_line = sensor.readline()
                if not raw_line:
                    continue

                line = raw_line.decode("utf-8", errors="replace").rstrip("\r\n")
                fields = line.split(",")

                if line.startswith("sequence,"):
                    if fields != EXPECTED_COLUMNS:
                        print("Firmware CSV header does not match capture schema", file=sys.stderr)
                        return 3
                    continue
                if not line or not line[0].isdigit():
                    ignored_line_count += 1
                    continue
                if len(fields) != len(EXPECTED_COLUMNS):
                    rejected_record_count += 1
                    continue

                try:
                    values = [int(field) for field in fields]
                except ValueError:
                    rejected_record_count += 1
                    continue

                sequence = values[0]
                timestamp_us = values[1]
                accel_status = values[ACCEL_STATUS_INDEX]
                gyro_status = values[GYRO_STATUS_INDEX]
                final_loop_overrun_count = values[LOOP_OVERRUN_INDEX]

                if not accel_status & ACCEL_DATA_READY_MASK:
                    accel_stale_count += 1
                if accel_status & ACCEL_OVERRUN_MASK:
                    accel_overrun_count += 1
                if not gyro_status & GYRO_DATA_READY_MASK:
                    gyro_stale_count += 1
                if gyro_status & GYRO_OVERRUN_MASK:
                    gyro_overrun_count += 1

                if previous_sequence is not None:
                    expected_sequence = (previous_sequence + 1) & 0xFFFFFFFF
                    if sequence != expected_sequence:
                        if sequence > expected_sequence:
                            missing_sequence_count += sequence - expected_sequence
                        else:
                            sequence_reset_count += 1
                        print(
                            f"  sequence discontinuity: expected {expected_sequence}, got {sequence}",
                            file=sys.stderr,
                            flush=True,
                        )
                if (
                    previous_timestamp_us is not None
                    and timestamp_us - previous_timestamp_us != 10_000
                ):
                    timing_anomaly_count += 1

                capture.write(line + "\n")
                previous_sequence = sequence
                previous_timestamp_us = timestamp_us
                sample_count += 1
                if sample_count % 500 == 0:
                    capture.flush()
                    print(f"  {sample_count} samples captured", flush=True)
    except serial.SerialException as exc:
        print(f"Serial error: {exc}", file=sys.stderr)
        return 1

    print(
        f"Stopped after {sample_count} samples; missing={missing_sequence_count}, "
        f"resets={sequence_reset_count}, timing_anomalies={timing_anomaly_count}, "
        f"rejected_records={rejected_record_count}, ignored_lines={ignored_line_count}; "
        f"saved {output}",
        flush=True,
    )
    print(
        f"Status totals: accel_stale={accel_stale_count}, "
        f"accel_overrun={accel_overrun_count}, gyro_stale={gyro_stale_count}, "
        f"gyro_overrun={gyro_overrun_count}, "
        f"acquisition_loop_overruns={final_loop_overrun_count}",
        flush=True,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())