#!/usr/bin/env python3
"""
R1 状态快速诊断脚本
运行 15 秒，通过 DDS 数据判断机器人当前姿态和电机状态。
"""
import sys
import time

sys.path.insert(0, "/home/qs_huawei/unitree_sdk2_python")
from unitree_sdk2py.core.channel import ChannelFactoryInitialize, ChannelSubscriber
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowState_

ChannelFactoryInitialize(0, "enx000e0988a670")

samples = []
start_time = time.time()
frame_count = 0

def handler(msg: LowState_):
    global samples, frame_count
    frame_count += 1
    
    # 只采样部分数据避免内存爆炸
    if len(samples) < 300:  # 最多存 300 帧（约 0.6s @ 500Hz）
        samples.append({
            "rpy": list(msg.imu_state.rpy),
            "motor_modes": [m.mode for m in msg.motor_state],
            "motor_tau": [m.tau_est for m in msg.motor_state[:12]],  # 前12个电机力矩
            "motor_q": [m.q for m in msg.motor_state[:6]],  # 前6个电机位置
            "timestamp": time.time() - start_time
        })

sub = ChannelSubscriber("rt/lowstate", LowState_)
sub.Init(handler, 10)

print("=" * 55)
print("R1 状态诊断监听中... 持续 15 秒")
print("=" * 55)

time.sleep(15)

print("\n" + "=" * 55)
print("诊断结果")
print("=" * 55)

if not samples:
    print("❌ 未收到任何 DDS 数据！")
    print("   可能原因：R1 未进入开发者模式，或网络不通")
    sys.exit(1)

print(f"总接收帧数: {frame_count}")
print(f"采样帧数: {len(samples)}")

# 1. 姿态分析
avg_rpy = [sum(s["rpy"][i] for s in samples) / len(samples) for i in range(3)]
pitch_deg = avg_rpy[1] * 57.3  # rad to deg
roll_deg = avg_rpy[0] * 57.3

print(f"\n📐 平均姿态:")
print(f"   Roll  = {roll_deg:+.2f}°")
print(f"   Pitch = {pitch_deg:+.2f}°")
print(f"   Yaw   = {avg_rpy[2] * 57.3:+.2f}°")

if abs(pitch_deg) < 10 and abs(roll_deg) < 10:
    print("   → 姿态判断: 机器人处于站立/直立状态 ✅")
elif abs(pitch_deg) > 60 or abs(roll_deg) > 60:
    print("   → 姿态判断: 机器人可能已倒地或大幅度倾斜 ⚠️")
else:
    print("   → 姿态判断: 机器人处于倾斜状态（可能正在动作中）")

# 2. 电机活跃度分析
latest = samples[-1]
active_count = sum(1 for m in latest["motor_modes"] if m == 1)
print(f"\n🔧 电机状态:")
print(f"   活跃电机: {active_count}/35")

if active_count < 20:
    print("   → 活跃电机数异常偏低，可能处于保护状态 ⚠️")
elif active_count >= 25:
    print("   → 电机正常启用 ✅")

# 3. 力矩分析（判断 ZeroTorque 的关键指标）
tau_all = []
for s in samples:
    tau_all.extend([abs(t) for t in s["motor_tau"]])

avg_tau = sum(tau_all) / len(tau_all) if tau_all else 0
max_tau = max(tau_all) if tau_all else 0

print(f"\n⚡ 前12电机力矩分析（采样 {len(tau_all)} 个点）:")
print(f"   平均绝对力矩: {avg_tau:.4f} Nm")
print(f"   最大绝对力矩: {max_tau:.4f} Nm")

if max_tau < 0.5:
    print("   → 力矩极低，机器人很可能处于 ZeroTorque 或 Damp 状态 ⚠️")
elif avg_tau < 1.0:
    print("   → 力矩偏低，可能处于被动支撑状态")
else:
    print("   → 力矩正常，电机在主动控制中 ✅")

# 4. 前6关节位置（判断腿部是否展开/折叠）
latest_q = samples[-1]["motor_q"]
q_deg = [q * 57.3 for q in latest_q]
print(f"\n🦵 前6关节位置（度）:")
for i, deg in enumerate(q_deg):
    print(f"   Motor {i}: {deg:+.2f}°")

# 5. 数据稳定性
rpy_variance = [
    sum((s["rpy"][i] - avg_rpy[i]) ** 2 for s in samples) / len(samples)
    for i in range(3)
]
print(f"\n📊 姿态稳定性（方差）:")
print(f"   Roll  方差: {rpy_variance[0]:.6f}")
print(f"   Pitch 方差: {rpy_variance[1]:.6f}")

if rpy_variance[1] < 0.001:
    print("   → 姿态非常稳定（静止状态）")
elif rpy_variance[1] > 0.01:
    print("   → 姿态变化较大（可能在运动或晃动中）")
else:
    print("   → 姿态轻微晃动")

print("=" * 55)