"""Side-by-side interactive playback: the same scene duplicated in one MuJoCo world.
Left copy = trajectory A (e.g. kinematic reference), right copy = trajectory B (e.g. physics rollout), time-synchronised."""
import argparse, os, time, json, copy
from pathlib import Path
import xml.etree.ElementTree as E
p = argparse.ArgumentParser(description=__doc__)
p.add_argument('--scene', type=Path, required=True); p.add_argument('--a', type=Path, required=True); p.add_argument('--b', type=Path, required=True)
p.add_argument('--label-a', default='A'); p.add_argument('--label-b', default='B'); p.add_argument('--offset', type=float, default=.7); p.add_argument('--check', action='store_true'); p.add_argument('--out', type=Path)
a = p.parse_args(); os.environ['MUJOCO_GL'] = 'osmesa' if a.check else 'glfw'
import numpy as np, mujoco
root = E.parse(a.scene).getroot(); world = root.find('worldbody')
NAMED = ['name', 'joint', 'joint1', 'joint2', 'site', 'body', 'target', 'geom1', 'geom2', 'body1', 'body2']
def prefix_tree(el, pre):
    for e in el.iter():
        for k in NAMED:
            v = e.get(k)
            if v and k != 'mesh' and k != 'material' and k != 'class': e.set(k, pre + v)
originals = list(world)
copyB = E.Element('body', name='B_root', pos=f'0 {a.offset} 0')
for e in originals:
    if e.tag == 'camera' or e.tag == 'light': continue
    c = copy.deepcopy(e); prefix_tree(c, 'B_'); copyB.append(c)
for e in originals:
    if e.tag not in ('camera', 'light'): prefix_tree(e, 'A_')
world.append(copyB)
for sec in ['actuator', 'equality', 'contact', 'sensor', 'tendon']:
    s = root.find(sec)
    if s is None: continue
    items = list(s)
    for e in items:
        c = copy.deepcopy(e); prefix_tree(c, 'B_'); s.append(c); prefix_tree(e, 'A_')
out = a.out or (a.a.parent / 'compare_scene.xml'); E.ElementTree(root).write(out)
m = mujoco.MjModel.from_xml_path(str(out)); d = mujoco.MjData(m)
def load(path):
    z = np.load(path); return z['qpos'], z['time']
qa, ta = load(a.a); qb, tb = load(a.b); n = qa.shape[1]; assert qb.shape[1] == n and m.nq == 2 * n, (m.nq, n)
# qpos layout: joints of A come first (original order), then B (same order) -> verify via joint names
ja = [i for i in range(m.njnt) if m.joint(i).name.startswith('A_')]; jb = [i for i in range(m.njnt) if m.joint(i).name.startswith('B_')]
adrA = np.concatenate([np.arange(m.jnt_qposadr[i], m.jnt_qposadr[i] + (7 if m.jnt_type[i] == 0 else 4 if m.jnt_type[i] == 1 else 1)) for i in ja]); adrB = np.concatenate([np.arange(m.jnt_qposadr[i], m.jnt_qposadr[i] + (7 if m.jnt_type[i] == 0 else 4 if m.jnt_type[i] == 1 else 1)) for i in jb])
assert len(adrA) == n and len(adrB) == n
tmax = float(min(ta[-1], tb[-1]))
def frame(tq, q, t): i = int(np.clip(np.searchsorted(tq, t, side='right') - 1, 0, len(tq) - 1)); return q[i]
def set_time(t):
    d.qpos[adrA] = frame(ta, qa, t); d.qpos[adrB] = frame(tb, qb, t); d.time = t; mujoco.mj_forward(m, d)
if a.check:
    for t in [0, tmax / 2, tmax]: set_time(t)
    print(json.dumps({'nq': m.nq, 'duration_s': tmax, 'finite': bool(np.isfinite(d.qpos).all()), 'scene': str(out)})); raise SystemExit
import mujoco.viewer
state = {'paused': False, 'time': 0., 'speed': 1., 'reset_camera': False}
def key(k):
    if k == 32: state['paused'] = not state['paused']
    elif k in [82, 114]: state['time'] = 0.
    elif k in [262, 263]: state['time'] = float(np.clip(state['time'] + (1 if k == 262 else -1) / 30, 0, tmax)); state['paused'] = True
    elif k == 265: state['speed'] = min(4., state['speed'] * 2)
    elif k == 264: state['speed'] = max(.125, state['speed'] / 2)
    elif k in [67, 99]: state['reset_camera'] = True
m.vis.headlight.ambient[:] = .7; set_time(0)
centers = [np.mean([d.site_xpos[m.site(pre + s + '_object_track').id] for s in ['right', 'left']], axis=0) for pre in ['A_', 'B_']]; center = np.mean(centers, axis=0)
cid = m.camera('A_source_camera').id if any(m.camera(i).name == 'A_source_camera' for i in range(m.ncam)) else 0; offset = d.cam_xpos[cid] - centers[0]
camera = {'distance': float(np.linalg.norm(offset)) * 1.6, 'azimuth': float(np.degrees(np.arctan2(offset[1], offset[0]))), 'elevation': float(-np.degrees(np.arctan2(offset[2], np.linalg.norm(offset[:2]))))}
status = out.parent / 'compare_playback_status.json'
print(f'LEFT (y=0): {a.label_a} | RIGHT (y=+{a.offset}): {a.label_b} | SPACE pause | R restart | LEFT/RIGHT frame step | UP/DOWN speed | C camera', flush=True)
with mujoco.viewer.launch_passive(m, d, key_callback=key) as viewer:
    def reset_camera():
        viewer.cam.type = mujoco.mjtCamera.mjCAMERA_FREE; viewer.cam.lookat[:] = center
        for k, x in camera.items(): setattr(viewer.cam, k, x)
    with viewer.lock():
        reset_camera(); viewer.opt.sitegroup[:] = 0; viewer.opt.geomgroup[3:] = 0
        # colour markers: green sphere above A, red above B
        for i, (c, rgba) in enumerate([(centers[0], (0, .8, 0, 1)), (centers[1], (.9, .1, .1, 1))]):
            g = viewer.user_scn.geoms[i]; mujoco.mjv_initGeom(g, mujoco.mjtGeom.mjGEOM_SPHERE, np.array([.02, 0, 0]), np.array(c) + np.array([0, 0, .35]), np.eye(3).ravel(), np.array(rgba, dtype=np.float32))
        viewer.user_scn.ngeom = 2
    last = time.perf_counter(); last_status = 0
    while viewer.is_running():
        now = time.perf_counter(); el = now - last; last = now
        if not state['paused']: state['time'] = (state['time'] + el * state['speed']) % tmax
        with viewer.lock():
            if state['reset_camera']: reset_camera(); state['reset_camera'] = False
            set_time(state['time'])
        viewer.sync()
        if now - last_status > 1:
            status.write_text(json.dumps({'running': True, 'pid': os.getpid(), 'left': a.label_a, 'right': a.label_b, 'time_s': state['time'], 'speed': state['speed'], 'paused': state['paused'], 'updated_unix': time.time()}, indent=2)); last_status = now
        time.sleep(.008)
status.write_text(json.dumps({'running': False, 'pid': os.getpid(), 'updated_unix': time.time()}))