#!/usr/bin/env python3
"""Low-level example that runs with either SDK's ``robot_interface`` module."""

import time

import robot_interface as sdk


# Unitree's official Python wrapper does not export these C++ constants.
LOWLEVEL = 0xff
FR_1 = 1


def control_step(udp, safe, cmd, state, target_q=1.2):
    """Execute one control iteration using only the official Python API."""
    udp.Recv()
    udp.GetRecv(state)

    motor = cmd.motorCmd[FR_1]
    motor.q = target_q
    motor.dq = 0.0
    motor.Kp = 5.0
    motor.Kd = 1.0
    motor.tau = 0.0

    safe.PositionLimit(cmd)
    if safe.PowerProtect(cmd, state, 1) < 0:
        raise RuntimeError("PowerProtect rejected the command")

    udp.SetSend(cmd)
    udp.Send()


def initialize_transport(udp, cmd, state, frames=50, dt=0.01):
    """Send only stop sentinels until the PRO MCU can return initial state."""
    udp.SetSend(cmd)
    received = False
    for _ in range(frames):
        udp.Send()
        time.sleep(dt)
        if udp.Recv() > 0:
            received = True
        udp.GetRecv(state)
    if not received:
        raise TimeoutError("No LowState received during initialization")


def main():
    udp = sdk.UDP(LOWLEVEL, 8080, "192.168.123.10", 8007)
    safe = sdk.Safety(sdk.LeggedType.Go1)
    cmd = sdk.LowCmd()
    state = sdk.LowState()
    udp.InitCmdData(cmd)
    initialize_transport(udp, cmd, state)

    while True:
        control_step(udp, safe, cmd, state)
        time.sleep(0.002)


if __name__ == "__main__":
    main()