"""Isolate contact geometry using prescribed link poses and a free bottle."""
from pathlib import Path
import xml.etree.ElementTree as E
import copy,json
import numpy as np
import mujoco

p=Path(__file__).resolve().parents[1]/'output/l20_2047635068/bottle_grasp'
tree=E.parse(p/'physics_scene.xml');root=tree.getroot();world=root.find('worldbody')
hand=world.find("body[@name='hand_base_link']")
original=mujoco.MjModel.from_xml_path(str(p/'physics_scene.xml'));d=mujoco.MjData(original)
a=np.load(p/'physics_motion.npz');d.qpos[:]=a['qpos'][0]
d.mocap_pos[0]=a['wrist'][0,:3];d.mocap_quat[0]=a['wrist'][0,3:];mujoco.mj_forward(original,d)
poses=[]
for b in hand.iter('body'):
    bid=original.body(b.get('name')).id
    new=E.SubElement(world,'body',name=b.get('name'),mocap='true')
    poses.append((d.xpos[bid].copy(),d.xquat[bid].copy()))
    for g in b.findall('geom'):new.append(copy.deepcopy(g))
world.remove(hand)
for tag in ['equality','actuator']:
    e=root.find(tag)
    if e is not None:root.remove(e)
path=p/'kinematic_grasp_scene.xml';tree.write(path)
m=mujoco.MjModel.from_xml_path(str(path));s=mujoco.MjData(m)
oi=original.jnt_qposadr[original.joint('object_free').id]
s.qpos[:]=d.qpos[oi:oi+7]
pos=np.array([x[0] for x in poses]);quat=np.array([x[1] for x in poses])
s.mocap_pos[:]=pos;s.mocap_quat[:]=quat
start=s.qpos[:3].copy();rows=[];mp=[];counts=[];drift=[]
for k in range(3000):
    t=k*.001;f=np.clip(t-.5,0,1);lift=.1*f*f*(3-2*f)
    s.mocap_pos[:]=pos+[0,0,lift];mujoco.mj_step(m,s)
    if k%33==0:
        rows.append(s.qpos.copy());mp.append(s.mocap_pos.copy());counts.append(s.ncon)
        drift.append(np.linalg.norm(s.qpos[:3]-start-[0,0,lift]))
np.savez_compressed(p/'kinematic_grasp_motion.npz',qpos=rows,mocap_pos=mp,mocap_quat=quat,fps=1/.033)
r={'free_object':True,'hand_prescribed_kinematics':True,
   'final_object_displacement':(s.qpos[:3]-start).tolist(),'expected_lift_m':.1,
   'relative_drift_m':float(np.linalg.norm(s.qpos[:3]-start-[0,0,.1])),
   'maximum_drift_m':float(max(drift)),'retained':bool(max(drift)<.02),
   'contact_frames':int(np.sum(np.array(counts)>0))}
(p/'kinematic_grasp_validation.json').write_text(json.dumps(r,indent=2));print(json.dumps(r,indent=2))