trikkeSensors/tools/capture_binary.py

#!/usr/bin/env python3
"""Capture validated TRK1 frames and render their samples to CSV."""

import argparse
import csv
import glob
import signal
import sys
from contextlib import ExitStack
from datetime import datetime
from pathlib import Path

import serial

from trikke_protocol import (
    CSV_COLUMNS,
    PACKET_TYPE_METADATA,
    Frame,
    IntegrityTracker,
    Metadata,
    StreamParser,
    sample_to_csv_row,
)


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)
    parser.add_argument("--output", type=Path, help="validated binary .trk output")
    parser.add_argument("--csv", type=Path, help="decoded CSV output")
    parser.add_argument(
        "--wire",
        type=Path,
        help="optional byte-for-byte serial capture, including startup text",
    )
    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:
        raise RuntimeError(
            f"multiple serial devices found ({', '.join(candidates)}); 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

    stem = datetime.now().strftime("binary_%Y%m%d_%H%M%S")
    output = args.output or Path("captures") / f"{stem}.trk"
    csv_output = args.csv or output.with_suffix(".csv")
    output.parent.mkdir(parents=True, exist_ok=True)
    csv_output.parent.mkdir(parents=True, exist_ok=True)
    if args.wire is not None:
        args.wire.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)

    parser = StreamParser()
    metadata: Metadata | None = None
    pending_frames: list[Frame] = []
    sample_count = 0
    metadata_count = 0
    integrity = IntegrityTracker()
    serial_error: serial.SerialException | None = None

    def render_frame(frame: Frame, writer: csv.writer) -> None:
        nonlocal sample_count
        if metadata is None:
            pending_frames.append(frame)
            return
        for sample in frame.samples:
            writer.writerow(sample_to_csv_row(sample, metadata, frame.loop_overrun_count))
            sample_count += 1
            if sample_count % 500 == 0:
                print(f"  {sample_count} samples captured", flush=True)

    print(f"Recording {port} to {output} and {csv_output}; press Ctrl-C to stop")
    try:
        with ExitStack() as stack:
            sensor = stack.enter_context(
                serial.Serial(port, args.baud, timeout=0.25)
            )
            raw_capture = stack.enter_context(output.open("wb"))
            decoded = stack.enter_context(
                csv_output.open("w", encoding="utf-8", newline="")
            )
            wire_capture = (
                stack.enter_context(args.wire.open("wb"))
                if args.wire is not None
                else None
            )
            writer = csv.writer(decoded)
            writer.writerow(CSV_COLUMNS)
            while not stop_requested:
                chunk = sensor.read(4096)
                if not chunk:
                    continue
                if wire_capture is not None:
                    wire_capture.write(chunk)
                for frame in parser.feed(chunk):
                    raw_capture.write(frame.raw)
                    integrity.observe(frame)
                    if frame.packet_type == PACKET_TYPE_METADATA:
                        metadata = frame.metadata
                        metadata_count += 1
                        for pending in pending_frames:
                            render_frame(pending, writer)
                        pending_frames.clear()
                    else:
                        render_frame(frame, writer)
                raw_capture.flush()
                decoded.flush()
                if wire_capture is not None:
                    wire_capture.flush()
    except serial.SerialException as exc:
        print(f"Serial error: {exc}", file=sys.stderr)
        serial_error = exc

    print(
        f"Stopped after {sample_count} samples and {metadata_count} metadata frames; "
        f"packet_gaps={integrity.packet_gap_count}, "
        f"packet_resets={integrity.packet_reset_count}, "
        f"sample_gaps={integrity.sample_gap_count}, "
        f"sample_resets={integrity.sample_reset_count}, "
        f"timing_anomalies={integrity.timing_anomaly_count}, "
        "timestamp_saturation_frames="
        f"{integrity.timestamp_saturation_frame_count}, "
        f"startup_crc_rejects={parser.startup_crc_errors}, "
        f"stream_crc_errors={parser.crc_errors}, "
        f"header_errors={parser.header_errors}, "
        f"skipped_nonframe_bytes={parser.skipped_bytes}, "
        f"trailing_partial_bytes={parser.buffered_bytes}"
    )
    print(
        f"Status totals: accel_stale={integrity.accel_stale_count}, "
        f"accel_overrun={integrity.accel_overrun_count}, "
        f"gyro_stale={integrity.gyro_stale_count}, "
        f"gyro_overrun={integrity.gyro_overrun_count}, "
        f"dropped={integrity.final_dropped_sample_count}, "
        f"acquisition_loop_overruns={integrity.final_loop_overrun_count}"
    )
    saved = f"Saved {output} and {csv_output}"
    if args.wire is not None:
        saved += f"; raw wire saved to {args.wire}"
    print(saved)
    if serial_error is not None:
        return 1
    return 0 if metadata is not None else 4


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