"""Matched wrist-local diagnostic view of the worst contact-jump interval."""
import os
os.environ.setdefault('MUJOCO_GL','osmesa')
import json
import numpy as np
import mujoco,imageio.v2 as imageio
from PIL import Image,ImageDraw
from l20_calibrated import ROOT
base=ROOT/'output/l20_2047635068';out=base/'jitter_audit'
sources=[base/'motion.npz',out/'temporal/motion.npz',base/'bottle_grasp/registered_motion.npz']
motion=[np.load(p) for p in sources]
model=mujoco.MjModel.from_xml_path(str(base/'l20_moving.xml'));data=mujoco.MjData(model)
addr=[model.jnt_qposadr[model.joint(str(n)).id] for n in motion[0]['joint_names']]
model.geom_rgba[model.geom('floor').id,3]=0;model.vis.headlight.ambient[:]=.8
model.vis.quality.offsamples=1;renderer=mujoco.Renderer(model,height=400,width=400)
cam=mujoco.MjvCamera();cam.lookat[:]=[.025,0,.105];cam.distance=.48;cam.azimuth=130;cam.elevation=5
opt=mujoco.MjvOption();opt.sitegroup[:]=0
writer=imageio.get_writer(out/'comparison_12s_18s.mp4',fps=30,codec='libx264',quality=8)
labels=['Original L20 (hand only)','Corrected + offline temporal smoothing','Previous contact fit (known jumps)']
for t in range(360,540):
    panels=[]
    for m in motion:
        data.qpos[addr]=m['qpos'][t];mujoco.mj_forward(model,data);renderer.update_scene(data,camera=cam,scene_option=opt);panels.append(renderer.render().copy())
    img=Image.fromarray(np.concatenate(panels,axis=1));draw=ImageDraw.Draw(img);draw.rectangle([0,0,1200,64],fill=(20,25,35))
    for j,label in enumerate(labels):
        step=np.rad2deg(np.abs(motion[j]['qpos'][t]-motion[j]['qpos'][t-1]).max())
        draw.text((j*400+8,7),label,fill='white');draw.text((j*400+8,24),f'frame {t} | maximum joint step {step:.2f} deg',fill='white')
    draw.text((8,45),'Wrist-local diagnostic view only; saved trajectories retain moving wrists. No physical grasp claim.',fill='white')
    writer.append_data(np.asarray(img))
    if t==461:img.save(out/'comparison_frame461.png')
writer.close();renderer.close()
print('COMPLETE 180 frames, source interval [360,540), 30 fps')