"""Render fixed-world and wrist-following views of a moving L20 trajectory."""
import os
os.environ.setdefault('MUJOCO_GL','osmesa')
import argparse,json
from pathlib import Path
import numpy as np
import mujoco,imageio.v2 as imageio
from PIL import Image,ImageDraw
from scipy.spatial.transform import Rotation

def main():
    parser=argparse.ArgumentParser()
    parser.add_argument('--directory',type=Path,default=Path(__file__).resolve().parents[1]/'output/l20_2047635068')
    parser.add_argument('--frames',type=int,default=0)
    args=parser.parse_args();out=args.directory
    with np.load(out/'motion.npz') as archive:
        motion={name:archive[name] for name in archive.files}
    q=motion['qpos'][:args.frames or None];positions=motion['wrist_pos'][:len(q)]
    model=mujoco.MjModel.from_xml_path(str(out/'l20_moving.xml'));data=mujoco.MjData(model)
    addr=[model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,str(n))] for n in motion['joint_names']]
    wrist=model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,'wrist_free')]
    rotations=Rotation.from_quat(motion['wrist_quat_wxyz'][:,[1,2,3,0]]).as_matrix()
    world=np.einsum('tij,tkj->tki',rotations,motion['actual'])+motion['wrist_pos'][:,None]
    center=(world.min((0,1))+world.max((0,1)))/2
    span=np.ptp(world,axis=(0,1)).max()
    floor=mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_GEOM,'floor')
    model.geom_pos[floor,2]=world[:,:,2].min()-.035
    model.vis.quality.shadowsize=1024;model.vis.quality.offsamples=1
    model.vis.headlight.ambient[:]=.65;model.vis.headlight.diffuse[:]=.7
    renderer=mujoco.Renderer(model,height=480,width=640)
    cameras=[]
    for azimuth in [135,25]:
        cam=mujoco.MjvCamera();cam.lookat[:]=center;cam.distance=max(.65,span*2.3);cam.azimuth=azimuth;cam.elevation=10
        cameras.append(cam)
    option=mujoco.MjvOption();option.sitegroup[:]=0
    fps=float(motion['fps'])
    coupling=str(motion['coupling']) if 'coupling' in motion else 'JSON-calibrated'
    writer=imageio.get_writer(out/'l20_moving_replay.mp4',fps=fps,codec='libx264',quality=8)
    max_wrist_error=0.
    try:
        for t,pose in enumerate(q):
            data.qpos[addr]=pose;data.qpos[wrist:wrist+3]=positions[t];data.qpos[wrist+3:wrist+7]=motion['wrist_quat_wxyz'][t]
            mujoco.mj_forward(model,data)
            base=mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_BODY,'hand_base_link')
            max_wrist_error=max(max_wrist_error,float(np.abs(data.xpos[base]-positions[t]).max()))
            panels=[]
            cameras[1].lookat[:]=world[t].mean(0);cameras[1].distance=.58
            for vi,cam in enumerate(cameras):
                renderer.update_scene(data,camera=cam,scene_option=option)
                if vi==0:
                    points=positions[np.unique(np.linspace(0,t,min(t+1,60),dtype=int))]
                    for a,b in zip(points[:-1],points[1:]):
                        scene=renderer.scene
                        mujoco.mjv_initGeom(scene.geoms[scene.ngeom],mujoco.mjtGeom.mjGEOM_CAPSULE,np.ones(3)*.001,a,np.eye(3).ravel(),np.array([.2,.8,1,1]))
                        mujoco.mjv_connector(scene.geoms[scene.ngeom],mujoco.mjtGeom.mjGEOM_CAPSULE,.0015,a,b)
                        scene.ngeom+=1
                panels.append(renderer.render().copy())
            img=Image.fromarray(np.concatenate(panels,axis=1));draw=ImageDraw.Draw(img)
            draw.rectangle([0,0,1280,45],fill=(15,22,30))
            draw.text((10,8),f'L20 right | frame {t+1}/{len(q)} | moving 6D wrist + {coupling} fingers',fill='white')
            draw.text((10,27),'Fixed world camera / blue wrist trail',fill='white')
            draw.text((650,27),'Following camera / finger detail',fill='white')
            writer.append_data(np.asarray(img))
            if t in [0,len(q)//2,len(q)-1]:img.save(out/f'preview_{t:04d}.png')
            if t%200==0:print(f'render L20 {t}/{len(q)}',flush=True)
    finally:writer.close();renderer.close()
    (out/'render_validation.json').write_text(json.dumps(dict(frames=len(q),fps=fps,max_wrist_position_error_m=max_wrist_error),indent=2))
    print('COMPLETE',flush=True)

if __name__=='__main__':main()