#!/usr/bin/env python3
"""Small-step arm tracking + order test via move_arm_joint_servo(group=2).

Stronger than the zero-motion gate: command ONE arm joint to move a small delta
(slew-limited), and verify (a) that joint moves ~delta and (b) the OTHER 13 joints
stay put. If a different joint moves => order wrong. If the target joint doesn't
reach => tracking problem. Then ramp back to the original pose.

Safe defaults: left wrist idx27_arm_l_joint7, +5 deg, 0.01 rad/tick (~0.3 rad/s).

Run:
  source /home/agi/workspace/env.sh
  export PYTHONPATH=/home/agi/workspace:$PYTHONPATH
  /data/agi/venvs/go1_torch/bin/python /data/agi/smallstep_test_0617.py [--live] [--joint 6] [--delta-deg 5]
"""
import argparse
import math
import signal
import sys
import time

import go1_env  # 路径来自 go1_env(GO1_* 环境变量, 仓库相对默认)
import agibot_gdk as g

ARM = ["idx21_arm_l_joint1", "idx22_arm_l_joint2", "idx23_arm_l_joint3", "idx24_arm_l_joint4",
       "idx25_arm_l_joint5", "idx26_arm_l_joint6", "idx27_arm_l_joint7",
       "idx61_arm_r_joint1", "idx62_arm_r_joint2", "idx63_arm_r_joint3", "idx64_arm_r_joint4",
       "idx65_arm_r_joint5", "idx66_arm_r_joint6", "idx67_arm_r_joint7"]
GROUP = 2
# smooth servo: 100Hz, control_period == tick period (demo_03 pattern)
DT = 1.0 / 100.0
CONTROL_PERIOD = 0.01

_stop = False
def _on_sigint(s, f):
    global _stop; _stop = True
    print("\n[!] Ctrl-C", flush=True)


def read_arm(robot):
    d = {s["name"]: s["motor_position"] for s in robot.get_joint_states()["states"]}
    return [d[n] for n in ARM]


def ramp(robot, start, target, slew, live):
    cmd = list(start)
    while True:
        if _stop: return "Ctrl-C"
        done = True
        for j in range(14):
            diff = target[j] - cmd[j]
            if abs(diff) > slew:
                cmd[j] += math.copysign(slew, diff); done = False
            else:
                cmd[j] = target[j]
        if live:
            try:
                rc = robot.move_arm_joint_servo(cmd, CONTROL_PERIOD, GROUP)
                if rc not in (0, None): return "rc=%r" % rc
            except Exception as e:
                return "%s: %s" % (type(e).__name__, e)
        time.sleep(DT)
        if done: return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--live", action="store_true")
    ap.add_argument("--joint", type=int, default=6, help="arm index 0..13 (default 6=idx27_arm_l_joint7)")
    ap.add_argument("--delta-deg", type=float, default=20.0)
    ap.add_argument("--slew", type=float, default=0.005, help="rad/tick @100Hz (0.005=0.5 rad/s)")
    ap.add_argument("--priority", type=int, default=20)
    args = ap.parse_args()
    signal.signal(signal.SIGINT, _on_sigint)
    delta = math.radians(args.delta_deg)

    print("=" * 60)
    print("Small-step arm test (%s)  joint=%d %s  delta=%+.1f deg  slew=%.3f" % (
        "LIVE" if args.live else "READ-ONLY", args.joint, ARM[args.joint], args.delta_deg, args.slew))
    print("=" * 60)
    if g.gdk_init() != g.GDKRes.kSuccess:
        print("gdk_init failed"); return 1
    robot = g.Robot(); time.sleep(0.5)
    start = read_arm(robot)
    target = list(start); target[args.joint] += delta
    print("start[%d]=%.3f -> target=%.3f" % (args.joint, start[args.joint], target[args.joint]))
    if not args.live:
        print("[read-only] add --live to move."); g.gdk_release(); return 0

    mc = g.MotionControlMode()
    mc.control_mode = g.ControlMode.CTRL_JOINT_POSITION
    mc.input_source = g.InputSource.INPUT_GDK
    mc.safe_mode = g.SafeMode.SAFE_NORMAL
    mc.target = g.Target.TARGET_ARMS
    mc.priority = int(args.priority)
    try:
        ret = robot.set_control_mode(mc)
    except Exception as e:
        print("set_control_mode raised: %r" % e); g.gdk_release(); return 1
    print("set_control_mode(prio=%d) -> %r" % (args.priority, ret))
    if ret != 0:
        print("could not acquire control (teleop?). refusing."); g.gdk_release(); return 1

    err = ramp(robot, start, target, args.slew, True)
    time.sleep(0.2)
    mid = read_arm(robot)
    moved = mid[args.joint] - start[args.joint]
    other = max((abs(mid[k] - start[k]) for k in range(14) if k != args.joint), default=0.0)
    owj = max((k for k in range(14) if k != args.joint), key=lambda k: abs(mid[k] - start[k]))
    print("\n[result] forward ramp err=%s" % err)
    print("  target joint moved : %+.2f deg (commanded %+.2f)" % (math.degrees(moved), args.delta_deg))
    print("  worst OTHER joint  : %.2f deg @%s" % (math.degrees(other), ARM[owj]))
    track_ok = abs(moved - delta) < math.radians(2.0)
    order_ok = other < math.radians(2.0)
    print("  tracking: %s   order: %s" % ("OK" if track_ok else "BAD", "OK" if order_ok else "BAD(other joint moved!)"))

    # ramp back
    print("[return] ramping back to start...")
    ramp(robot, mid, start, args.slew, True)
    time.sleep(0.2)
    fin = read_arm(robot)
    back = max(abs(fin[k] - start[k]) for k in range(14))
    print("  back to start within %.2f deg" % math.degrees(back))

    try:
        mc.priority = 1; print("[release] ->", robot.set_control_mode(mc))
    except Exception as e:
        print("[release] raised: %r" % e)
    g.gdk_release()
    print("\n[verdict] %s" % ("PASS: arm tracks small target, correct joint" if (track_ok and order_ok)
                              else "FAIL: see above"))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())