"""Play a hand-only L20 motion file with a free wrist in MuJoCo."""
import argparse
import time
import mujoco
import mujoco.viewer
import numpy as np

p = argparse.ArgumentParser()
p.add_argument('--directory', required=True)
a = p.parse_args()
from pathlib import Path
directory = Path(a.directory)
with np.load(directory / 'motion.npz') as data:
    motion = {name:data[name] for name in data.files}
model = mujoco.MjModel.from_xml_path(str(directory / 'l20_moving.xml'))
data = mujoco.MjData(model)
addr = [model.jnt_qposadr[model.joint(str(n)).id] for n in motion['joint_names']]
wrist = model.jnt_qposadr[model.joint('wrist_free').id]
fps = float(motion['fps'])
count = len(motion['qpos'])
state = dict(frame=0., paused=False)

def key(k):
    if k == 32:
        state['paused'] = not state['paused']
    elif k in [82,114]:
        state['frame'] = 0.
    elif k in [262,263]:
        state['frame'] = (int(state['frame']) + (1 if k == 262 else -1)) % count
        state['paused'] = True

with mujoco.viewer.launch_passive(model, data, key_callback=key) as viewer:
    viewer.cam.lookat[:] = motion['wrist_pos'].mean(0)
    viewer.cam.distance = .9
    viewer.cam.azimuth, viewer.cam.elevation = 135, 10
    viewer.opt.sitegroup[:] = 0
    last = time.perf_counter()
    while viewer.is_running():
        now = time.perf_counter()
        if not state['paused']:
            state['frame'] = (state['frame'] + (now-last)*fps) % count
        last = now
        t = int(state['frame'])
        with viewer.lock():
            data.qpos[addr] = motion['qpos'][t]
            data.qpos[wrist:wrist+3] = motion['wrist_pos'][t]
            data.qpos[wrist+3:wrist+7] = motion['wrist_quat_wxyz'][t]
            data.time = t/fps
            mujoco.mj_forward(model,data)
        viewer.sync()
        time.sleep(.005)