#!/usr/bin/env python3
"""复刻 free-dog-sdk example_position(lowlevel).py 的 PRO 版本.

跟原版 EDU 例程逻辑一致:
  0..10 步:    记录初始关节角 qInit
  10..400 步:  Kp=5 Kd=1 插值到 sin_mid_q = [0.0, 1.2, -2.0]
  400+ 步:     1Hz sin 摆动 (FR_1 += 0.6*sin, FR_2 += -0.9*sin)

差异:
  - 用 Go1 PRO SDK (Blowfish 加密 / PRO 格式 / 安全保护)
  - 集成 safety 层 (power_factor=1, 位置硬限位)

⚠️ 狗必须悬空 (脚架/吊带). Ctrl+C 紧急停止.
"""
import argparse
import math
import signal
import sys
import time

from go1_pro_sdk import (
    MCUClient, LowCmd, MotorCmd, MotorMode,
    apply_safety, PowerProtectViolation,
    JOINT_NAMES,
)

# 跟原版一致的常量
SIN_MID = {'hip': 0.0, 'thigh': 1.2, 'knee': -2.0}
DT = 0.002

running = True


def sigint(s, f):
    global running
    running = False


def interp(a, b, t):
    t = max(0.0, min(1.0, t))
    return a * (1 - t) + b * t


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--state', default=None)
    p.add_argument('--max-steps', type=int, default=5000)
    p.add_argument('--freq-hz', type=float, default=1.0)
    p.add_argument('--power-factor', type=int, default=1, choices=range(1, 11),
                   help='PowerProtect 1-10, 越大越宽松')
    args = p.parse_args()
    signal.signal(signal.SIGINT, sigint)

    print('🐕 PRO example_position (Go1 PRO SDK)')
    print(f'  控制周期 {DT*1000:.0f}ms ({1/DT:.0f} Hz)')
    print(f'  sin 中点: {SIN_MID}')
    print(f'  sin 频率: {args.freq_hz} Hz')
    print(f'  power_factor: {args.power_factor} (限 {args.power_factor*10}% 力矩)')
    print(f'  最大步数: {args.max_steps} (≈{args.max_steps*DT:.1f}s)')

    with MCUClient(state_path=args.state) as client:
        print(f'\n[Phase 0] 唤醒 MCU...')
        recv = client.wake_mcu(50)
        if recv == 0:
            print('❌ 无回包. 先停 Pi 上的 keep_sport_alive/Legged_sport/appTransit')
            return 1
        state = client.last_state
        print(f'  当前 FR_0={state.motorState[0].q:.4f} '
              f'FR_1={state.motorState[1].q:.4f} '
              f'FR_2={state.motorState[2].q:.4f}')

        # 主循环
        qInit = [0.0, 0.0, 0.0]
        qDes = [0.0, 0.0, 0.0]
        Kp = [0.0, 0.0, 0.0]
        Kd = [0.0, 0.0, 0.0]
        sin_count = 0
        rate_count = 0
        motiontime = 0
        freq_rad = args.freq_hz * 2 * math.pi

        print(f'\n[主循环]')
        loop_start_t = time.time()
        loop_times = []
        last_print = time.time()

        try:
            while running and motiontime < args.max_steps:
                t0 = time.time()
                motiontime += 1
                state = client.recv_latest() or state

                if motiontime < 10:
                    for i in range(3):
                        qInit[i] = state.motorState[i].q

                if 10 <= motiontime < 400:
                    rate_count += 1
                    rate = rate_count / 200.0
                    Kp = [5, 5, 5]
                    Kd = [1, 1, 1]
                    qDes[0] = interp(qInit[0], SIN_MID['hip'], rate)
                    qDes[1] = interp(qInit[1], SIN_MID['thigh'], rate)
                    qDes[2] = interp(qInit[2], SIN_MID['knee'], rate)

                if motiontime >= 400:
                    sin_count += 1
                    t = DT * sin_count
                    sin_j1 = 0.6 * math.sin(t * freq_rad)
                    sin_j2 = -0.9 * math.sin(t * freq_rad)
                    qDes[0] = SIN_MID['hip']
                    qDes[1] = SIN_MID['thigh'] + sin_j1
                    qDes[2] = SIN_MID['knee'] + sin_j2

                # 构造命令: 只控 FR (索引 0,1,2)
                cmd = LowCmd()
                for i in range(3):
                    cmd.motorCmd[i] = MotorCmd(
                        mode=MotorMode.Servo,
                        q=qDes[i], dq=0,
                        tau=-0.65 if i == 0 else 0.0,   # FR_0 预紧
                        Kp=Kp[i], Kd=Kd[i],
                    )

                # 安全保护
                try:
                    apply_safety(cmd, state, power_factor=args.power_factor,
                                 position_limit_on=True,
                                 position_protect_limit=0.5)
                except PowerProtectViolation as e:
                    print(f'\n⚠️ {e}')
                    break

                client.send(cmd)
                loop_times.append(time.time() - t0)

                now = time.time()
                if now - last_print >= 0.5:
                    phase = 'qInit' if motiontime < 10 else 'ramp' if motiontime < 400 else 'sin '
                    print(f'  t={motiontime*DT:5.2f}s [{phase}] '
                          f'des=[{qDes[0]:+.3f} {qDes[1]:+.3f} {qDes[2]:+.3f}] '
                          f'act=[{state.motorState[0].q:+.3f} '
                          f'{state.motorState[1].q:+.3f} {state.motorState[2].q:+.3f}]')
                    last_print = now

                # 维持 DT
                sleep_t = DT - (time.time() - t0)
                if sleep_t > 0:
                    time.sleep(sleep_t)

            # 时序统计
            if loop_times:
                total = time.time() - loop_start_t
                avg_ms = sum(loop_times) / len(loop_times) * 1000
                print(f'\n实际频率: {len(loop_times)/total:.1f} Hz, '
                      f'每步均值 {avg_ms:.2f}ms')
        finally:
            print('\n安全停机...')
            client.safe_stop()
            print('✅ 退出')


if __name__ == '__main__':
    sys.exit(main() or 0)