#!/usr/bin/env python3
"""Squat-down / stand-up SMOOTHNESS test (coordinated waist chain).

Moves the 5 waist joints together from the current pose to the demonstrated
deepest config and back, smoothly (100Hz, control_period=0.01). The target is
taken from real teleop data (episode_0000 deepest frame) so it is COM-valid.
Stays WITHIN the demonstrated waist range -> safe.
"""
import argparse, math, signal, sys, time
import go1_env  # 路径来自 go1_env(GO1_* 环境变量, 仓库相对默认)
import agibot_gdk as g

WAIST = ["idx01_body_joint1","idx02_body_joint2","idx03_body_joint3","idx04_body_joint4","idx05_body_joint5"]
# demonstrated deepest waist config (episode_0000 frame 226) -- COM-valid
DEEP = [-0.545, 1.410, -0.556, -0.033, -0.027]
SERVO_HZ = 100.0; DT = 1.0/SERVO_HZ; CP = 0.01

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

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

def ramp(robot, start, target, slew):
    cmd = list(start)
    while True:
        if _stop: return "Ctrl-C"
        done = True
        for j in range(5):
            diff = target[j]-cmd[j]
            if abs(diff) > slew:
                cmd[j] += math.copysign(slew, diff); done = False
            else:
                cmd[j] = target[j]
        try:
            rc = robot.move_waist_joint_servo(cmd, CP, True)
            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

ap = argparse.ArgumentParser()
ap.add_argument("--live", action="store_true")
ap.add_argument("--speed", type=float, default=0.4, help="rad/s")
ap.add_argument("--cycles", type=int, default=2)
ap.add_argument("--priority", type=int, default=20)
ap.add_argument("--scale", type=float, default=1.0, help="fraction of (deep-start) travel, <=1 stays in demo range")
args = ap.parse_args()
slew = args.speed * DT

g.gdk_init(); robot = g.Robot(); time.sleep(0.5)
start = read(robot)
deep = [start[j] + args.scale*(DEEP[j]-start[j]) for j in range(5)]
print("start waist:", [round(x,3) for x in start])
print("deep  waist:", [round(x,3) for x in deep])
print("travel(deg):", [round(math.degrees(deep[j]-start[j]),1) for j in range(5)])
if not args.live:
    print("[read-only] add --live"); g.gdk_release(); raise SystemExit(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)
ret = robot.set_control_mode(mc)
print("set_control_mode ->", ret)
if ret != 0:
    print("no control (teleop?). refusing."); g.gdk_release(); raise SystemExit(1)

for c in range(args.cycles):
    if _stop: break
    print("\n-- cycle %d: DOWN (squat) --" % (c+1), flush=True)
    e = ramp(robot, read(robot), deep, slew)
    if e: print("down ERR:", e); break
    time.sleep(0.6)
    print("-- cycle %d: UP (stand) --" % (c+1), flush=True)
    e = ramp(robot, read(robot), start, slew)
    if e: print("up ERR:", e); break
    time.sleep(0.6)

fin = read(robot)
print("\nfinal waist:", [round(x,3) for x in fin], " back-err(deg):",
      round(math.degrees(max(abs(fin[j]-start[j]) for j in range(5))),2))
try:
    mc.priority = 1; print("[release] ->", robot.set_control_mode(mc))
except Exception as e:
    print("[release]", e)
g.gdk_release()
print("[done]")