"""Build locomotion replay bundles for the existing trace viewer.

Usage: python replay_export.py --mjcf MODEL.xml [--visuals anymal-d-visuals.json] --recordings ROOT --output DIR
(ROOT holds the recording folders named in SHOWCASES, e.g. replays-01/, replays-02/.)
A showcase with torque_plot also writes a synchronized line plot of the largest absolute applied joint torque per
policy step (recordings with joint_torque_Nm, replays-05 onward), matched to its panels by engine and checkpoint hash.

Display geometry is ANYbotics' ANYmal D visual meshes (anymal_visuals.py, BSD-3-Clause) when
--visuals is given; otherwise ANYmal-D's collision model as compiled by Newton (capsule base,
box/cylinder legs, sphere feet), tessellated per body in the body frame. Poses are the
native body link poses recorded at the 50 Hz policy rate; foot arrows are the native
net contact force on each foot, drawn at the foot. A robot that falls holds its last
recorded pose (Isaac Lab resets it in the same step; the reset is not shown).
"""
import argparse
import gzip
import hashlib
import json
import math
from pathlib import Path
import xml.etree.ElementTree as ET

import numpy as np

COORDINATES = dict(frame='world', up='z', quaternion='wxyz', units='SI')
LABELS = dict(isaacsim_physx='PhysX', newton_mjwarp='Newton')
COLORS = dict(isaacsim_physx=0x315d9e, newton_mjwarp=0xb65335)
# Revolute joints: (child body = joint location, body whose origin the joint swings, axis in the child frame).
# Axes from ANYbotics' anymal.urdf (joint frame = child frame at zero angle; frames match the simulated model).
JOINTS = {f'{leg}_{j}': (f'{leg}_{child}', f'{leg}_{end}', axis)
          for leg, signs in dict(LF=(1, -1, 1), RF=(1, 1, -1), LH=(-1, -1, 1), RH=(-1, 1, -1)).items()
          for (j, child, end), sign in zip((('HAA', 'HIP', 'SHANK'), ('HFE', 'THIGH', 'SHANK'), ('KFE', 'SHANK', 'FOOT')), signs)
          for axis in [(sign, 0., 0.)]}
COMMAND_GAIN_M_PER_RAD = .5  # exaggeration: 0.2 rad of commanded joint motion -> 10 cm arrow
SHOWCASES = {
    'locomotion-transfer-forward1.json.gz': dict(
        folder='replays-05', env='22', color_by='train', torque_plot='locomotion-transfer-torque.json', panels=[('c0b-P0-on-P', 'PhysX-trained · run in PhysX'), ('c0b-P0-on-N', 'PhysX-trained · run in Newton'),
                          ('c0b-N0-on-N', 'Newton-trained · run in Newton'), ('c0b-N0-on-P', 'Newton-trained · run in PhysX')],
        notice='Seed-0 policies trained at Isaac Lab defaults, 1 m/s forward command, one fixed randomized starting state (robot 22 of 64). Selected after inspecting outcomes: one of 4 PhysX-to-Newton falls for this seed.'),
    'locomotion-joint-order.json.gz': dict(
        folder='replays-03', env='22', panels=[('c0b-N0-on-P', 'Newton-trained in PhysX · joints matched by name'),
                          ('c0b-N0-on-P-unmapped', 'Same policy · native joint order (wrong)')],
        notice='The same Newton-trained policy and starting state in PhysX. Without name matching, joint observations and leg commands are silently permuted.'),
    'locomotion-coarse-steps.json.gz': dict(
        folder='replays-03', env='0', panels=[('c0b-N0-on-N-steps', 'Newton · 5 ms step (default) · walks'), ('c4-N0-on-N', 'Newton · 10 ms step · never learned to walk'),
                         ('c5-N0-on-N', 'Newton · 20 ms step · training collapsed'), ('c5-P0-on-P', 'PhysX · 20 ms step · never learned to walk')],
        notice='Seed-0 policies, each trained and run in its own engine and step size, 1 m/s forward command, robot 0 of 16. At 20 ms, Newton feet sink about 1 cm into the soft contact the step forces, and PhysX feet jitter up to 3.3 cm into the floor. Selected after inspecting outcomes.'),
    'locomotion-stance-c1.json.gz': dict(
        folder='replays-03', env='49', panels=[('c1-P1-on-P', 'PhysX-trained (seed 1) · standing in PhysX'), ('c1-P1-on-N', 'Same policy · standing in Newton')],
        notice='Friction-matched cell, seed 1, the only seed with standing falls (8/64 in PhysX, 22/64 in Newton in the evaluation). Stand command, robot 49, which fell in both engines in the evaluation and again in this recording run. Selected after inspecting outcomes.'),
    'locomotion-cone.json.gz': dict(
        folder='replays-04', env='22', panels=[('c0b-P0-on-N', 'Newton · elliptic cone, impratio 100 (default)'),
                          ('c8-P0-on-N', 'Newton · pyramidal cone, impratio 1')],
        notice='The PhysX-trained seed-0 policy at default friction, 1 m/s forward command, robot 22 of 64 (the opening replay\'s starting state), recorded in one session. Only the Newton friction cone and impratio differ. Of the four seed-0 robots that fall with the default cone (21, 22, 44, 48), 22, 44 and 48 walk with the pyramidal cone, in the evaluation and in this recording run. Selected after inspecting outcomes.'),
}


def quat_matrix(w, x, y, z):
    return np.array([[1-2*(y*y+z*z), 2*(x*y-z*w), 2*(x*z+y*w)],
                     [2*(x*y+z*w), 1-2*(x*x+z*z), 2*(y*z-x*w)],
                     [2*(x*z-y*w), 2*(y*z+x*w), 1-2*(x*x+y*y)]])


def primitive(kind, size, n=16):
    """Vertices/faces of an MJCF primitive centred at the origin, axis along z."""
    if kind == 'box':
        hx, hy, hz = size
        v = [[sx*hx, sy*hy, sz*hz] for sx in (-1, 1) for sy in (-1, 1) for sz in (-1, 1)]
        f = [[0, 2, 3], [0, 3, 1], [4, 5, 7], [4, 7, 6], [0, 1, 5], [0, 5, 4],
             [2, 6, 7], [2, 7, 3], [0, 4, 6], [0, 6, 2], [1, 3, 7], [1, 7, 5]]
        return v, f
    if kind == 'sphere':
        r = size[0]; rings = n//2
        v = [[r*math.sin(math.pi*i/rings)*math.cos(2*math.pi*j/n), r*math.sin(math.pi*i/rings)*math.sin(2*math.pi*j/n),
              r*math.cos(math.pi*i/rings)] for i in range(rings+1) for j in range(n)]
        f = []
        for i in range(rings):
            for j in range(n):
                a, b, c, d = i*n+j, i*n+(j+1) % n, (i+1)*n+j, (i+1)*n+(j+1) % n
                f += [[a, c, b], [b, c, d]]
        return v, f
    r, h = size[0], size[1]
    if kind == 'cylinder':
        v = [[r*math.cos(2*math.pi*j/n), r*math.sin(2*math.pi*j/n), z] for z in (-h, h) for j in range(n)]+[[0, 0, -h], [0, 0, h]]
        f = []
        for j in range(n):
            k = (j+1) % n
            f += [[j, k, n+k], [j, n+k, n+j], [2*n, k, j], [2*n+1, n+j, n+k]]
        return v, f
    if kind == 'capsule':  # cylinder with hemispherical caps
        half = n//4  # rings per hemisphere; the equator ring is duplicated to form the straight band
        rows = [(math.pi/2*i/half, h) for i in range(half+1)]+[(math.pi/2*(1+i/half), -h) for i in range(half+1)]
        v = []
        for phi, offset in rows:
            v += [[r*math.sin(phi)*math.cos(2*math.pi*j/n), r*math.sin(phi)*math.sin(2*math.pi*j/n), r*math.cos(phi)+offset] for j in range(n)]
        rings = len(rows)-1
        f = []
        for i in range(rings):
            for j in range(n):
                a, b, c, d = i*n+j, i*n+(j+1) % n, (i+1)*n+j, (i+1)*n+(j+1) % n
                f += [[a, c, b], [b, c, d]]
        return v, f
    raise ValueError(kind)


def body_meshes(mjcf):
    meshes = {}
    for body in ET.parse(mjcf).getroot().iter('body'):
        name = body.get('name').split('_Robot_')[-1]
        vertices, faces = [], []
        for g in body.findall('geom'):
            kind = g.get('type', 'sphere')
            v, f = primitive(kind, [float(x) for x in g.get('size').split()])
            R = quat_matrix(*[float(x) for x in (g.get('quat') or '1 0 0 0').split()])
            p = np.array([float(x) for x in (g.get('pos') or '0 0 0').split()])
            offset = len(vertices)
            vertices += (np.array(v)@R.T+p).round(6).tolist()
            faces += [[a+offset, b+offset, c+offset] for a, b, c in f]
        meshes[name] = dict(id=name, shape='mesh', vertices=vertices, faces=faces)
    return meshes


def trace(recording_path, env, label, meshes, color=None):
    rec = json.loads(gzip.decompress(Path(recording_path).read_bytes()))
    names, feet = rec['body_names'], rec['foot_bodies']
    data = rec['envs'][env]
    scene = dict(task='quadruped', floor=True, dt=rec['dt_s'], duration=data['frames'][-1]['time_s'],
                 bodies=[meshes[n] for n in names])
    scene_sha = hashlib.sha256(json.dumps(scene, sort_keys=True).encode()).hexdigest()
    frames, previous = [], None
    joint_names = rec.get('joint_names')
    for raw in data['frames']:
        poses = np.array(raw['poses'])
        commands = {}
        if joint_names and 'joint_target' in raw and raw['time_s'] > 0:
            # Policy command per joint: (target - current) angle, drawn at the joint along the motion it drives.
            for jn, target, current in zip(joint_names, raw['joint_target'], raw['joint_pos']):
                child, end, axis = JOINTS[jn]
                ci, ei = names.index(child), names.index(end)
                w = poses[ci, 3:] / np.linalg.norm(poses[ci, 3:])
                axis_world = quat_matrix(*w) @ np.array(axis)
                tangent = np.cross(axis_world, poses[ei, :3] - poses[ci, :3])
                if np.linalg.norm(tangent) < 1e-6:
                    continue
                vector = COMMAND_GAIN_M_PER_RAD * (target - current) * tangent / np.linalg.norm(tangent)
                commands.setdefault(child, []).append(dict(on_body=child, kind='command', point_m=poses[ci, :3].tolist(),
                                                           vector_N=vector.tolist(), joint=jn, command_rad=target - current))
        bodies = []
        for i, n in enumerate(names):
            p, q = poses[i, :3], poses[i, 3:]
            q = q/np.linalg.norm(q)
            v = np.zeros(3) if previous is None else (p-previous[0][i])/(raw['time_s']-previous[1])
            forces = []
            if n in feet:
                f = raw['foot_force_N'][feet.index(n)]
                if np.linalg.norm(f) > 1e-6:
                    forces.append(dict(on_body=n, kind='contact', point_m=p.tolist(), vector_N=f))
            forces += commands.get(n, [])
            bodies.append(dict(id=n, position_m=p.tolist(), quaternion_wxyz=q.tolist(), velocity_m_s=v.tolist(),
                               angular_velocity_rad_s=[0., 0., 0.], forces=forces))
        frames.append(dict(time_s=raw['time_s'], bodies=bodies))
        previous = (poses[:, :3], raw['time_s'])
    torque = [float(np.abs(raw['joint_torque_Nm']).max()) for raw in data['frames'] if 'joint_torque_Nm' in raw]
    return dict(schema='simcompare.trace.v2', engine=label[0], label=label[1], color=color or COLORS[label[0]],
                coordinates=COORDINATES, scene=scene, scene_sha256=scene_sha, reference_sha256=label[2],
                frames=frames, status='fell' if data['fell'] else 'complete', stop_time_s=frames[-1]['time_s'],
                outcome=dict(fell=data['fell'], fall_time_s=data['fall_time_s']),
                display_note='Native body link poses at 50 Hz; display geometry only (see provenance), not the simulated collision shapes.',
                _torque=torque)


def main(mjcf, recordings, output, visuals=None):
    meshes = body_meshes(mjcf)
    geometry = Path(visuals or mjcf)
    if visuals:
        parts = json.loads(geometry.read_text())['bodies']
        meshes = {name: dict(id=name, shape='parts', parts=parts.get(name, [])) for name in meshes}
    output.mkdir(parents=True, exist_ok=True)
    provenance = dict(schema='locomotion.study-c.replay-provenance.v1', geometry=geometry.name,
                      geometry_sha256=hashlib.sha256(geometry.read_bytes()).hexdigest(),
                      geometry_note=('ANYbotics anymal_d_simple_description visual meshes (BSD-3-Clause), display only'
                                     if visuals else 'Newton-compiled collision model'), bundles={})
    for bundle_name, spec in SHOWCASES.items():
        traces, sources, trained_by = [], [], []
        for run, text in spec['panels']:
            meta = json.loads((recordings/spec['folder']/f'{run}.json').read_text())
            command = next(iter(meta['results']))
            rec = recordings/spec['folder']/f'{run}-{command}-recording.json.gz'
            backend = 'newton_mjwarp' if '-on-N' in run else 'isaacsim_physx'
            trained = 'newton_mjwarp' if '-N' in run.split('-on-')[0] else 'isaacsim_physx'
            ckpt_sha = hashlib.sha256(meta['checkpoint'].encode()).hexdigest()
            color = COLORS[trained] if spec.get('color_by') == 'train' else None
            trained_by.append(trained)
            traces.append(trace(str(rec), spec['env'], (backend, text, ckpt_sha), meshes, color))
            sources.append(dict(folder=spec['folder'], run=run, command=command, env=spec['env'], checkpoint=meta['checkpoint'], hydra=meta['hydra'],
                                recording_sha256=hashlib.sha256(rec.read_bytes()).hexdigest(),
                                fell=traces[-1]['outcome']['fell'], fall_time_s=traces[-1]['outcome']['fall_time_s']))
        torques = [t.pop('_torque') for t in traces]
        if spec.get('torque_plot'):
            dt = traces[0]['frames'][1]['time_s'] - traces[0]['frames'][0]['time_s']
            plot = dict(schema='locomotion.torque-plot.v1', title='Largest joint torque (Nm)', unit='Nm', decimals=1,
                        span='union', dt_s=round(dt, 6),
                        definition=('Largest absolute applied joint torque over the 12 joints at each 50 Hz policy step '
                                    '(the actuator torque of the last physics substep, the same sample as torque_rms). '
                                    'Colour is the training engine; dashed lines run in the other engine. '
                                    'A trace ends when its robot falls.'),
                        traces=[dict(engine=t['engine'], label=t['label'], color=f"#{t['color']:06x}",
                                     dash=None if t['engine'] == trained else '6 4',
                                     start_s=t['frames'][1]['time_s'], values=v[1:],
                                     after_end='fell' if t['outcome']['fell'] else None,
                                     scene_sha256=t['scene_sha256'], reference_sha256=t['reference_sha256'])
                                for t, v, trained in zip(traces, torques, trained_by)])
            for t in plot['traces']:
                if len(t['values']) == 0:
                    raise SystemExit(f'{bundle_name}: recording has no joint_torque_Nm')
            plot['traces'] = [{k: v for k, v in t.items() if v is not None} for t in plot['traces']]
            (output/spec['torque_plot']).write_text(json.dumps(plot, separators=(',', ':'))+'\n')
            provenance.setdefault('plots', {})[spec['torque_plot']] = dict(bundle=bundle_name, definition=plot['definition'])
        path = output/bundle_name
        path.write_bytes(gzip.compress(json.dumps(dict(schema='simcompare.bundle.v1', traces=traces, notice=spec['notice']),
                                                  separators=(',', ':')).encode(), mtime=0))
        provenance['bundles'][bundle_name] = dict(sha256=hashlib.sha256(path.read_bytes()).hexdigest(), notice=spec['notice'], sources=sources)
    (output/'replay-provenance.json').write_text(json.dumps(provenance, indent=1)+'\n')


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('--mjcf', type=Path, required=True)
    p.add_argument('--recordings', type=Path, required=True)
    p.add_argument('--output', type=Path, required=True)
    p.add_argument('--visuals', type=Path)
    a = p.parse_args(); main(a.mjcf, a.recordings, a.output, a.visuals)