#!/usr/bin/env python3
"""Snapshot / restore G2 pose.

  --save [--out PATH]    : read current 22 joints (qpos order) -> save .npy (read-only)
  --restore [--in PATH] --live [--speed R] [--no-waist] [--no-head]
                         : smoothly ramp to the saved pose (100Hz, verified servo APIs)

Default file: $GO1_POSE_FILE or <deploy>/saved_init_pose.npy   (qpos order = DEFAULT_JOINT_NAMES_Go2:
  [0:5] body idx01-05 | [5:8] head idx11-13 | [8:15] left arm | [15:22] right arm)
"""
import argparse, math, signal, sys, time
import go1_env  # 路径来自 go1_env(GO1_* 环境变量, 仓库相对默认)
import numpy as np
import agibot_gdk as g

ALLNAMES=["idx01_body_joint1","idx02_body_joint2","idx03_body_joint3","idx04_body_joint4","idx05_body_joint5",
    "idx11_head_joint1","idx12_head_joint2","idx13_head_joint3",
    "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"]
DEFAULT=go1_env.POSE_FILE
SERVO_HZ=100.0; DT=1.0/SERVO_HZ; CP=0.01
HEAD_PITCH=(-0.10,0.65)

_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 ALLNAMES]

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("--save",action="store_true")
    ap.add_argument("--restore",action="store_true")
    ap.add_argument("--live",action="store_true")
    ap.add_argument("--out",default=DEFAULT)
    ap.add_argument("--in",dest="inp",default=DEFAULT)
    ap.add_argument("--speed",type=float,default=0.4,help="rad/s (restore)")
    ap.add_argument("--no-waist",action="store_true")
    ap.add_argument("--no-head",action="store_true")
    ap.add_argument("--priority",type=int,default=20)
    args=ap.parse_args()
    if not (args.save or args.restore):
        print("specify --save or --restore"); return 1

    g.gdk_init(); robot=g.Robot(); time.sleep(0.5)
    cur=read(robot)

    if args.save:
        np.save(args.out, np.array(cur,dtype=np.float32))
        print("[save] current pose -> %s"%args.out)
        print("  body:",[round(x,3) for x in cur[0:5]])
        print("  head:",[round(x,3) for x in cur[5:8]])
        print("  armL:",[round(x,3) for x in cur[8:15]])
        print("  armR:",[round(x,3) for x in cur[15:22]])
        g.gdk_release(); return 0

    # restore
    tgt=np.load(args.inp).astype(float).tolist()
    tgt[7]=min(max(tgt[7],HEAD_PITCH[0]),HEAD_PITCH[1])
    maxd=max(abs(tgt[i]-cur[i]) for i in range(22)); wj=max(range(22),key=lambda i:abs(tgt[i]-cur[i]))
    print("[restore] target <- %s"%args.inp)
    print("  current:",[round(x,2) for x in cur])
    print("  target :",[round(x,2) for x in tgt])
    print("  max travel %.2f rad @%s -> ~%.1fs @ %.2f rad/s"%(maxd,ALLNAMES[wj],maxd/args.speed,args.speed))
    if not args.live:
        print("[dry-run] 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)
    ret=robot.set_control_mode(mc); print("  set_control_mode ->",ret)
    if ret!=0: print("  no control (teleop?). refusing."); g.gdk_release(); return 1

    slew=args.speed*DT; cmd=list(cur); t0=time.monotonic(); tick=0; err=None
    while True:
        if _stop: err="Ctrl-C"; break
        done=True
        for i in range(22):
            diff=tgt[i]-cmd[i]
            if abs(diff)>slew: cmd[i]+=math.copysign(slew,diff); done=False
            else: cmd[i]=tgt[i]
        try:
            rc=robot.move_arm_joint_servo(cmd[8:22],CP,2)
            if rc not in (0,None): err="arm rc=%r"%rc; break
            if not args.no_head:
                rc=robot.move_head_joint_servo([cmd[5],cmd[6],cmd[7]],CP,True)
                if rc not in (0,None): err="head rc=%r"%rc; break
            if not args.no_waist:
                rc=robot.move_waist_joint_servo(cmd[0:5],CP,True)
                if rc not in (0,None): err="waist rc=%r"%rc; break
        except Exception as e: err="%s:%s"%(type(e).__name__,e); break
        tick+=1; rem=t0+tick*DT-time.monotonic()
        if rem>0: time.sleep(rem)
        if done: break
        if tick>int((maxd/args.speed+8)*SERVO_HZ): err="cap"; break

    time.sleep(0.3); fin=read(robot)
    be=max(abs(fin[i]-tgt[i]) for i in range(22)); bj=max(range(22),key=lambda i:abs(fin[i]-tgt[i]))
    print("[restore] err=%s  final max|err|=%.3f rad @%s"%(err,be,ALLNAMES[bj]))
    try:
        mc.priority=1; print("[release] ->",robot.set_control_mode(mc))
    except Exception as e: print("[release]",e)
    g.gdk_release(); print("[done]"); return 0

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