import statistics
import time

from fast_lowcmd import FastLowCmdBuilder, default_state_path
from go1_pro_sdk import Blowfish, LowCmd, MotorCmd, MotorMode
from go1_pro_sdk.codec.lowcmd_builder import build_low_cmd_encrypted, build_low_cmd_plain


def bench(name, fn, n=5000):
    for _ in range(100):
        fn()
    ts = []
    for _ in range(n):
        t = time.perf_counter()
        fn()
        ts.append(time.perf_counter() - t)
    ts.sort()
    print(
        f"{name:28s} median {statistics.median(ts) * 1e6:8.1f} us "
        f"p95 {ts[int(.95 * (n - 1))] * 1e6:8.1f} us "
        f"p99 {ts[int(.99 * (n - 1))] * 1e6:8.1f} us "
        f"max {max(ts) * 1e6:8.1f} us"
    )


def main():
    q = [-0.1, 0.8, -1.5, 0.1, 0.8, -1.5, -0.1, 1.0, -1.5, 0.1, 1.0, -1.5]

    cmd = LowCmd()
    for i in range(12):
        cmd.set_motor(i, MotorCmd(mode=MotorMode.Servo, q=q[i], Kp=28.0, Kd=0.7))

    bf = Blowfish.from_state_file(default_state_path())
    builder = FastLowCmdBuilder(default_state_path())

    assert builder.build_plain_servo12(q, kp=28.0, kd=0.7) == build_low_cmd_plain(cmd)
    assert builder.build_encrypted_servo12(q, kp=28.0, kd=0.7) == build_low_cmd_encrypted(cmd, bf)

    bench("python build_plain", lambda: build_low_cmd_plain(cmd), n=5000)
    bench("python build_encrypted", lambda: build_low_cmd_encrypted(cmd, bf), n=2000)
    bench("native build_plain_lowcmd", lambda: builder.build_plain_lowcmd(cmd), n=10000)
    bench("native build_encrypted_lowcmd", lambda: builder.build_encrypted_lowcmd(cmd), n=10000)
    bench("native build_plain_servo12", lambda: builder.build_plain_servo12(q, kp=28.0, kd=0.7), n=20000)
    bench("native build_encrypted_servo12", lambda: builder.build_encrypted_servo12(q, kp=28.0, kd=0.7), n=10000)
    bench("native damping encrypted", lambda: builder.build_encrypted_damping(), n=10000)


if __name__ == "__main__":
    main()