"""Build assets/microduck.glb from Pollen Robotics' Microduck MuJoCo model (Apache-2.0).

Assembles the visual meshes of github.com/pollen-robotics/microduck_rl in their default pose, skips hidden internals,
merges them by color (optionally simplified with --keep), and saves one y-up glTF facing +z with the feet on y=0. Each leg is
a separate node (`leg_left`, `leg_right`) pivoting at the hip, so the game can animate a walk.

    pip install trimesh numpy fast-simplification && python make_duck.py
"""

import argparse
import io
import urllib.request
import xml.etree.ElementTree as ET
from collections import defaultdict
from pathlib import Path

import fast_simplification
import numpy as np
import trimesh


REPO = "https://raw.githubusercontent.com/pollen-robotics/microduck_rl/develop/src/mjlab_microduck/robot/microduck"
KEEP = 1.0  # default fraction of each part's triangles kept (--keep); simplifying tears the CAD meshes
# Internals hidden by the shells: bearings, boards, battery and supports.
HIDDEN = {"seeed_bearing__configuration__22x16x4", "seeed_bearing__configuration_default", "pcb__raspberry_pi_zero_2_w",
          "elec_rpi_robot_hat_pcb", "np_f970", "power_support", "banana_pcb_locker", "speaker", "motor_support", "m12_lens_holder"}


def simplify(mesh, keep, passes):
    """Keep about `keep` of the triangles; one pass stops early to preserve shape, so `passes` repeats it."""
    if keep >= 1:
        return mesh
    goal = max(150, int(len(mesh.faces) * keep))
    for _ in range(passes):
        if len(mesh.faces) <= goal:
            break
        vertices, faces = fast_simplification.simplify(mesh.vertices, mesh.faces, target_reduction=1 - goal / len(mesh.faces))
        if len(faces) >= len(mesh.faces) * 0.95:
            break
        mesh = trimesh.Trimesh(vertices, faces)
    return mesh


def fetch(path):
    with urllib.request.urlopen(f"{REPO}/{path}", timeout=60) as response:
        return response.read()


def transform(element):
    """4x4 matrix from an MJCF element's pos and quat (w x y z)."""
    matrix = np.eye(4)
    if "quat" in element.attrib:
        w, x, y, z = map(float, element.get("quat").split())
        matrix[:3, :3] = trimesh.transformations.quaternion_matrix([w, x, y, z])[:3, :3]
    if "pos" in element.attrib:
        matrix[:3, 3] = list(map(float, element.get("pos").split()))
    return matrix


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--keep", type=float, default=KEEP, help="Fraction of triangles kept per part (1 = no simplification).")
    parser.add_argument("--passes", type=int, default=1, help="Simplification passes; more get closer to --keep but add artifacts.")
    parser.add_argument("--out", default=str(Path(__file__).parent / "assets" / "microduck.glb"))
    args = parser.parse_args()
    root = ET.fromstring(fetch("robot_walk.xml"))
    colors = {m.get("name"): [float(v) for v in m.get("rgba").split()] for m in root.iter("material")}
    # parts: group ("body", "leg_left", "leg_right") -> rgba -> [(mesh name, mesh in the robot's frame)]
    meshes, parts, pivots = {}, defaultdict(lambda: defaultdict(list)), {}

    def visit(body, parent, group):
        world = parent @ transform(body)
        for geom in body.findall("geom"):
            name = geom.get("mesh")
            if geom.get("class") != "visual" or name in HIDDEN:
                continue
            if name not in meshes:
                meshes[name] = simplify(trimesh.load(io.BytesIO(fetch(f"assets/{name}.stl")), file_type="stl"), args.keep, args.passes)
            mesh = meshes[name].copy()
            mesh.apply_transform(world @ transform(geom))
            parts[group][tuple(colors[geom.get("material")])].append((name, mesh))
        for child in body.findall("body"):
            child_group = group
            if group == "body":  # each leg is its own node, pivoting where it attaches to the trunk
                below = {g.get("mesh") for g in child.iter("geom")}
                side = "left" if "foot_left" in below else "right" if "foot_right" in below else None
                if side:
                    child_group = f"leg_{side}"
                    pivots[child_group] = (world @ transform(child))[:3, 3]
            visit(child, world, child_group)

    for body in root.find("worldbody").findall("body"):
        visit(body, np.eye(4), "body")
    every = [mesh for colors_ in parts.values() for group in colors_.values() for _, mesh in group]
    print(f"{len(meshes)} meshes, {len(every)} visible parts, {sum(len(m.faces) for m in every):,} triangles, groups {list(parts)}")

    # MuJoCo is z-up; glTF is y-up. Then turn the duck so its face looks along +z, and stand it on y=0, centered.
    z_up_to_y_up = trimesh.transformations.rotation_matrix(-np.pi / 2, [1, 0, 0])
    everything = trimesh.util.concatenate(every)
    face = trimesh.util.concatenate([m for colors_ in parts.values() for group in colors_.values() for n, m in group if n == "face_part"])
    forward = (face.centroid - everything.centroid) @ z_up_to_y_up[:3, :3].T
    frame = trimesh.transformations.rotation_matrix(-np.arctan2(forward[0], forward[2]), [0, 1, 0]) @ z_up_to_y_up
    low, high = everything.copy().apply_transform(frame).bounds
    frame = trimesh.transformations.translation_matrix([-(low[0] + high[0]) / 2, -low[1], -(low[2] + high[2]) / 2]) @ frame

    scene = trimesh.Scene()
    for group, by_color in parts.items():
        pivot = (frame @ [*pivots[group], 1])[:3] if group in pivots else np.zeros(3)
        scene.graph.update(frame_from=scene.graph.base_frame, frame_to=group, matrix=trimesh.transformations.translation_matrix(pivot))
        for i, (rgba, group_parts) in enumerate(by_color.items()):
            mesh = trimesh.util.concatenate([m for _, m in group_parts])
            mesh.apply_transform(trimesh.transformations.translation_matrix(-pivot) @ frame)  # relative to the group's pivot
            mesh.visual = trimesh.visual.TextureVisuals(material=trimesh.visual.material.PBRMaterial(
                baseColorFactor=[int(c * 255) for c in rgba], metallicFactor=0.1, roughnessFactor=0.6))
            scene.add_geometry(mesh, node_name=f"{group}_{i}", parent_node_name=group)

    out = Path(args.out)
    out.parent.mkdir(exist_ok=True)
    out.write_bytes(scene.export(file_type="glb"))
    size = scene.bounds[1] - scene.bounds[0]
    print(f"wrote {out} ({out.stat().st_size / 1e6:.1f} MB), size {size[0]:.3f} x {size[1]:.3f} x {size[2]:.3f} m, "
          f"leg pivots {({g: np.round((frame @ [*p, 1])[:3], 3).tolist() for g, p in pivots.items()})}")


if __name__ == "__main__":
    main()