"""Convert ANYbotics' ANYmal D simple description into replay display parts.

Usage: python anymal_visuals.py DESCRIPTION_DIR OUT.json

Source: github.com/ANYbotics/anymal_d_simple_description (BSD-3-Clause, (c) 2023 ANYbotics AG;
see ANYBOTICS-LICENSE.txt). Each visual mesh is expressed in the frame of the nearest simulated
body (base, *_HIP, *_THIGH, *_SHANK, *_FOOT) by composing the URDF's fixed joints. Those frames
match the Newton-compiled Isaac Lab model exactly at the zero pose (checked when building).
Textures are replaced by flat colours; "engine" marks parts tinted with each panel's colour.
The inspection payload (pan/tilt joints not simulated in Isaac Lab) is omitted.
"""
import json
import math
from pathlib import Path
import sys
import xml.etree.ElementTree as ET

import numpy as np

NS = {'c': 'http://www.collada.org/2005/11/COLLADASchema'}
BODIES = {'base'} | {f'{leg}_{part}' for leg in ('LF', 'LH', 'RF', 'RH') for part in ('HIP', 'THIGH', 'SHANK', 'FOOT')}
COLORS = {'top_shell': 'engine', 'face_shell': 'engine', 'foot': 0x262626, 'drive': 0x4a4a4a, 'hip': 0x5c5c5c,
          'thigh': 0x6e6e6e, 'shank_l': 0x6e6e6e, 'shank_r': 0x6e6e6e, 'base': 0x3a3a3a, 'bottom_shell': 0x2f2f2f,
          'battery': 0x3a3a3a, 'docking_socket': 0x3a3a3a, 'lidar': 0x2b2b2b, 'lidar_guard': 0x555555}
SKIP_LINKS = ('inspection_payload',)


def rpy(r, p, y):
    cr, sr, cp, sp, cy, sy = math.cos(r), math.sin(r), math.cos(p), math.sin(p), math.cos(y), math.sin(y)
    return np.array([[cy*cp, cy*sp*sr-sy*cr, cy*sp*cr+sy*sr], [sy*cp, sy*sp*sr+cy*cr, sy*sp*cr-cy*sr], [-sp, cp*sr, cp*cr]])


def transform(origin):
    T = np.eye(4)
    if origin is not None:
        T[:3, :3] = rpy(*[float(v) for v in origin.get('rpy', '0 0 0').split()])
        T[:3, 3] = [float(v) for v in origin.get('xyz', '0 0 0').split()]
    return T


def collada(path):
    """Triangles of a single-node Collada file in its own frame (metres, Z up)."""
    root = ET.parse(path).getroot()
    assert root.find('c:asset/c:up_axis', NS).text == 'Z_UP'
    scale = float(root.find('c:asset/c:unit', NS).get('meter'))
    node = root.find('.//c:node', NS)
    M = np.array([float(v) for v in node.find('c:matrix', NS).text.split()]).reshape(4, 4)
    vertices, faces = [], []
    for mesh in root.iterfind('.//c:mesh', NS):
        sources = {s.get('id'): np.array([float(v) for v in s.find('c:float_array', NS).text.split()])
                   for s in mesh.iterfind('c:source', NS)}
        position = mesh.find('c:vertices/c:input[@semantic="POSITION"]', NS).get('source')[1:]
        points = sources[position].reshape(-1, 3)*scale
        points = (np.c_[points, np.ones(len(points))]@M.T)[:, :3]
        for tri in mesh.iterfind('c:triangles', NS):
            inputs = tri.findall('c:input', NS)
            stride = max(int(i.get('offset')) for i in inputs)+1
            offset = int(next(i for i in inputs if i.get('semantic') == 'VERTEX').get('offset'))
            index = np.array([int(v) for v in tri.find('c:p', NS).text.split()]).reshape(-1, stride)[:, offset]
            base = len(vertices)
            used = np.unique(index)
            remap = {int(old): base+k for k, old in enumerate(used)}
            vertices += points[used].tolist()
            faces += [[remap[int(a)], remap[int(b)], remap[int(c)]] for a, b, c in index.reshape(-1, 3)]
    return np.array(vertices), faces


def main(description, out):
    description = Path(description)
    urdf = ET.parse(description/'urdf'/'anymal.urdf').getroot()
    parent = {j.find('child').get('link'): (j.find('parent').get('link'), transform(j.find('origin')))
              for j in urdf.iter('joint')}

    def to_body(link):
        T = np.eye(4)
        while link not in BODIES:
            link, J = parent[link]
            T = J@T
        return link, T

    parts = {b: [] for b in BODIES}
    cache = {}
    for link in urdf.iter('link'):
        name = link.get('name')
        if name.startswith(SKIP_LINKS):
            continue
        for visual in link.findall('visual'):
            mesh = visual.find('geometry/mesh')
            if mesh is None:
                continue
            stem = Path(mesh.get('filename')).stem
            if stem not in cache:
                cache[stem] = collada(description/'meshes'/f'{stem}.dae')
            v, f = cache[stem]
            body, T = to_body(name)
            T = T@transform(visual.find('origin'))
            v = (np.c_[v, np.ones(len(v))]@T.T)[:, :3]
            parts[body].append(dict(name=f'{name}:{stem}', color=COLORS.get(stem, 0x555555),
                                    vertices=v.round(5).tolist(), faces=f))
    Path(out).write_text(json.dumps(dict(
        schema='locomotion.anymal-d-visuals.v1',
        source='github.com/ANYbotics/anymal_d_simple_description', license='BSD-3-Clause, (c) 2023 ANYbotics AG',
        note='Display geometry only; simulation used the Isaac Lab ANYmal-D asset. Frames verified identical at the zero pose.',
        triangles=sum(len(p['faces']) for ps in parts.values() for p in ps), bodies=parts), separators=(',', ':')))


if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2])