#!/usr/bin/env python3
"""
R1 零力矩模式 & 阻尼模式 测试脚本

R1 与 G1 共用 unitree_hg 接口，直接复用 G1 的 LocoClient。

⚠️ 安全警告：
- 测试前必须将 R1 悬挂或安排人扶住！
- ZeroTorque() 和 Damp() 会让机器人失去主动支撑，直接摔倒！
- 首次测试建议先运行 CheckMode() 确认通信正常。

使用方式：
    python3 test_r1_torque_damp.py <网卡名>
    例：python3 test_r1_torque_damp.py enx000e0988a670
"""

import sys
import time

from unitree_sdk2py.core.channel import ChannelFactoryInitialize
from unitree_sdk2py.g1.loco.g1_loco_client import LocoClient


def print_menu():
    print("\n" + "=" * 40)
    print("R1 零力矩 / 阻尼模式 测试")
    print("=" * 40)
    print("  [0] ZeroTorque  — 零力矩模式（电机无力矩，机器人会瘫倒）")
    print("  [1] Damp        — 阻尼模式（关节可被动推动）")
    print("  [2] Start       — 恢复站立/正常状态")
    print("  [3] Squat2StandUp — 蹲姿→站立（如果机器人已倒地）")
    print("  [4] Check FSM ID — 查询当前状态码")
    print("  [q] Quit        — 退出")
    print("=" * 40)


def main():
    if len(sys.argv) < 2:
        print(f"用法: python3 {sys.argv[0]} <网卡名>")
        print(f"示例: python3 {sys.argv[0]} enx000e0988a670")
        sys.exit(1)

    network_interface = sys.argv[1]

    print("\n⚠️  安全确认：请确保 R1 已悬挂或有人扶住！")
    print("    零力矩/阻尼模式会导致机器人失去支撑并摔倒。\n")
    input("确认安全后按 Enter 继续...")

    # 初始化 DDS 通道
    print(f"\n[初始化] 网卡: {network_interface}")
    ChannelFactoryInitialize(0, network_interface)

    # 创建 LocoClient（R1 复用 G1 的客户端）
    client = LocoClient()
    client.SetTimeout(10.0)
    ret = client.Init()

    if ret != 0:
        print(f"[错误] LocoClient Init 失败，返回值: {ret}")
        print("可能原因：")
        print("  1. R1 未进入开发者模式（遥控器按 L2 + R2）")
        print("  2. 网络不通（检查网线、IP、网卡名）")
        print("  3. R1 不支持 LocoClient 接口")
        sys.exit(1)

    print("[成功] LocoClient 初始化完成，与 R1 建立通信。\n")

    # 先查询一次当前状态
    try:
        code, fsm_id = client.GetFsmId()
        if code == 0:
            print(f"[当前状态] FSM ID: {fsm_id}")
        else:
            print(f"[警告] 查询状态失败，code: {code}")
    except Exception as e:
        print(f"[警告] GetFsmId 调用异常: {e}")
        print("继续运行，但 R1 可能不支持全部 LocoClient API。\n")

    while True:
        print_menu()
        choice = input("输入选项: ").strip().lower()

        if choice == "q":
            print("\n退出前恢复站立状态...")
            client.Start()
            time.sleep(1)
            print("已退出。")
            break

        elif choice == "0":
            print("\n>>> 正在进入 零力矩模式 (ZeroTorque)...")
            print("    电机将完全断电，机器人会瘫软！")
            confirm = input("    确认执行? [y/N]: ").strip().lower()
            if confirm == "y":
                code = client.ZeroTorque()
                print(f"    调用完成，返回值: {code} (0=成功)")
            else:
                print("    已取消。")

        elif choice == "1":
            print("\n>>> 正在进入 阻尼模式 (Damp)...")
            print("    电机提供阻尼，关节可被动推动。")
            confirm = input("    确认执行? [y/N]: ").strip().lower()
            if confirm == "y":
                code = client.Damp()
                print(f"    调用完成，返回值: {code} (0=成功)")
            else:
                print("    已取消。")

        elif choice == "2":
            print("\n>>> 正在恢复站立 (Start)...")
            code = client.Start()
            print(f"    调用完成，返回值: {code} (0=成功)")
            time.sleep(0.5)

        elif choice == "3":
            print("\n>>> 正在执行 蹲姿→站立 (Squat2StandUp)...")
            print("    注意：此操作仅在机器人处于蹲姿/倒地时有效。")
            code = client.Squat2StandUp()
            print(f"    调用完成，返回值: {code} (0=成功)")
            time.sleep(0.5)

        elif choice == "4":
            try:
                code, fsm_id = client.GetFsmId()
                if code == 0:
                    print(f"\n[当前 FSM ID]: {fsm_id}")
                    # 常见 FSM ID 含义参考
                    fsm_map = {
                        0: "零力矩 (ZeroTorque)",
                        1: "阻尼 (Damp)",
                        500: "站立/正常 (Start)",
                        702: "躺→站 (Lie2StandUp)",
                        706: "蹲→站 (Squat2StandUp)",
                    }
                    desc = fsm_map.get(fsm_id, "未知状态")
                    print(f"[状态含义]: {desc}")
                else:
                    print(f"\n[错误] 查询失败，code: {code}")
            except Exception as e:
                print(f"\n[错误] 调用异常: {e}")

        else:
            print("\n无效选项，请重新输入。")

        time.sleep(0.3)


if __name__ == "__main__":
    main()