Drive from Python

# The SDK counts the 0.5 s a robot in STAND must be upright to arm from its own samples.
time.sleep(0.6)
print(f"robot mode {robot.get_state().mode.name}, armed {robot.armed}")
for action in ("stand", "move", "trajectory"):
    print(robot.preflight(action))  # "ready to move", or "not ready to move:" and why

try:
    # Checks the robot first. Returns once it reports STAND and has been upright for
    # 0.5 s (armed): the firmware accepts MOVE only after that.
    robot.stand()
except (NotReadyError, WaitTimeoutError) as exc:
    print(exc)  # what is wrong and what fixes it
    sys.exit(1)
print(f"robot mode {robot.get_state().mode.name}, armed {robot.armed}")

try:
    # Checks the robot first when it is in STAND. Returns once it reports MOVE.
    robot.balance()
except (NotReadyError, WaitTimeoutError) as exc:
    print(exc)  # what is wrong and what fixes it
    sys.exit(1)
print(f"robot mode {robot.get_state().mode.name}, balancing in place")

while True:
    show(robot, last)
    key = next_key(TICK_S)
    if key is None:
        continue
    if key in QUIT:
        break
    if key in MOVES:
        last = move(robot, key)
    elif key == " ":
        last = balance(robot)
    elif key == "t":
        last = stand(robot)
    elif key == "b":
        last = confirm_and_damp(robot, next_key)

# damp() is never refused, so there is no check here, only the question.
if not YES:
    mode = robot.get_state().mode.name
    question = f"Robot mode {mode}. Is the robot supported? Damp now? [y/N] "
    if input(question).strip().lower() not in ("y", "yes"):
        print("nothing sent")
        sys.exit(0)
try:
    robot.damp()  # returns once the robot reports DAMP
except WaitTimeoutError as exc:  # sent, but DAMP was not reported: the message says what next
    print(exc)
    sys.exit(1)
print(f"robot mode {robot.get_state().mode.name}")

if robot.has("camera"):
    photo = robot.camera.photo()  # the next fresh frame, rgb8
    Path(PHOTO).write_bytes(photo.to_jpeg())
    print(f"saved {PHOTO}, {photo.width}x{photo.height}")
    clip = robot.camera.capture_clip(CLIP_S, audio=robot.has("microphone"))
    frames = clip.save_frames(CLIP_DIR)
    print(f"saved {len(frames)} frames to {CLIP_DIR}/ at {clip.fps:.1f} fps")
    if clip.audio:
        print("saved", clip.save_wav(CLIP_WAV))
else:
    print("this connection carries no camera; use hybrid or livekit")

if robot.has("speaker"):
    # 16-bit mono PCM: a sine wave at TONE_HZ for TONE_S.
    tone = b"".join(
        struct.pack("<h", int(8000 * math.sin(2 * math.pi * TONE_HZ * i / SAMPLE_RATE_HZ)))
        for i in range(int(SAMPLE_RATE_HZ * TONE_S))
    )
    robot.speaker.play_pcm(tone, sample_rate_hz=SAMPLE_RATE_HZ)
    print(f"played {TONE_S:.0f} s at {TONE_HZ:.0f} Hz")

if not robot.has("microphone"):
    print("this connection carries no microphone; use hybrid or livekit")
    sys.exit(1)
# Every chunk in order, as 16-bit PCM. chunks() raises WaitTimeoutError when the
# microphone says nothing for its timeout.
chunks: list[AudioChunk] = []
recorded_s = 0.0
for chunk in robot.microphone.chunks(timeout=5.0):
    chunks.append(chunk)
    recorded_s += chunk.duration_s
    if recorded_s >= SECONDS:
        break
with wave.open(PATH, "wb") as wav:
    wav.setnchannels(chunks[0].channels)
    wav.setsampwidth(2)  # 16-bit samples
    wav.setframerate(chunks[0].sample_rate_hz)
    wav.writeframes(b"".join(c.data for c in chunks))
print(f"saved {PATH}, {recorded_s:.1f} s at {chunks[0].sample_rate_hz} Hz")

if not robot.has("speaker"):
    print("this connection carries no speaker; use hybrid or livekit")
    sys.exit(1)
if os.path.exists(PATH):
    with wave.open(PATH, "rb") as wav:
        if wav.getsampwidth() != 2:
            print(f"{PATH} is not 16-bit PCM; convert it first")
            sys.exit(1)
        rate, channels = wav.getframerate(), wav.getnchannels()
        pcm = wav.readframes(wav.getnframes())
    name = PATH
else:  # 16-bit mono PCM: one second of a 440 Hz sine wave
    rate, channels, name = 16_000, 1, f"a 440 Hz tone (no {PATH} here)"
    pcm = b"".join(
        struct.pack("<h", int(8000 * math.sin(2 * math.pi * 440 * i / rate)))
        for i in range(rate)
    )
# Sent in short pieces, in order: play_pcm() blocks until the room has taken each one,
# so the file plays at its own pace.
step = int(rate * CHUNK_S) * 2 * channels
for start in range(0, len(pcm), step):
    robot.speaker.play_pcm(pcm[start : start + step], sample_rate_hz=rate, channels=channels)
print(f"played {name}, {len(pcm) / (2 * channels * rate):.1f} s at {rate} Hz")