"""Door robot arm: collision hulls and the FBX export for Unreal (Coder stage 3; stages 'ucx' and 'export').

ucx      UCX_ convex hulls per link (k-means clusters of the link's vertices, one hull each, k by the link's size),
         parented to the showroom's link objects, in the collection 'DoorArm | collision' (eye off in the view layer,
         hidden in renders). They go with the static link meshes (SM_DoorRobotArm_Links.fbx) for a Blueprint build; the
         skeletal mesh gets a Physics Asset in Unreal instead.
export   into fbx/ (everything from temporary copies in a temporary scene; nothing is saved):
  SK_DoorRobotArm_<A|B>.fbx, ..._LOD1.fbx, ..._LOD2.fbx   the arm per version: one skeleton (a root bone and one bone per
          moving link, at the link's pivot, the A axes), one skinned mesh (every link rigid on its bone; the elbow cables
          and cylinder hoses blended between their two links). LOD1 about 50 %, LOD2 about 20 % of the triangles: the
          markings and small parts dropped first (screws, clips), then decimated (collapse).
  A_DoorRobotArm_<A|B>_<AFT|FORE>_OpenClose.fbx   the skeleton with the door animation (frames 1-120, 30 fps) baked
          from the rig in the ship-fit scene (the version switched on first: a hidden version is not evaluated).
  SM_DoorRobotArm_Links.fbx    every link as a static mesh at its pivot with its UCX_ hulls.
  SM_DoorArm_RailKit_<AFT|FORE>.fbx, SM_DoorArm_Fixture<A|B>_<AFT|FORE>.fbx   static, ship coordinates (fixtures: origin on
          the door's attach point), with simple UCX_ hulls.
  SK_DoorArm_Chain_<AFT|FORE>.fbx + A_DoorArm_Chain_<AFT|FORE>_OpenClose.fbx   the energy chain: one bone per link.
  DoorRobotArm_layout.json (next to the .blend): every number a consumer needs.
The project's FBX convention (docking arm, pressure door): forward -Y, up Z, the data turned -90 deg about Z first, so
Unreal (import with "Convert Scene" and "Force Front X Axis") reads ue_cm = (x, -y, z) * 100 of the Blender data.
Statics: global scale 1, unit scale applied. Skeletal: data x100 and the exporter's own x100 cancelled by global
scale 0.01 (centimetres, no scale on any node), as the docking arm's hose.
"""
import math
import os
import json
import time
import numpy as np
import bpy
import bmesh
from mathutils import Vector, Matrix

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
FBX = os.path.join(ROOT, "fbx")
PRE = Matrix.Rotation(-math.pi / 2, 4, 'Z')
SKT = Matrix.Scale(100.0, 4) @ PRE
COLL = "DoorArm | collision"
LOD_RATIO = {1: 0.50, 2: 0.20}
LOD_SMALL = {1: 0.030 * 1.6, 2: 0.065 * 1.6}   # parts smaller than this (bounding box diagonal, m) are dropped first (x the scale-up)


def bone_name(link):
    return link.replace(" ", "_").lower()


def sm_name(link):
    return "SM_DoorArm_" + "".join(w.capitalize() for w in link.replace("J3", "j3").split())


def links(ver=None):
    import dra_rigdata as rd
    return [l for l in rd.LINKS if ver is None or not l["tool"] or l["tool"] == ver]


def showroom_objects():
    sc = bpy.data.scenes["DoorArm | Showroom"]
    out = {}
    for v in ("A", "B"):
        for o in sc.objects:
            if o.name.startswith("Showroom | arm %s |" % v) and o.type == 'MESH' and "link" in o:
                out.setdefault(o["link"], o)
    return out


def coll(sc, name):
    c = bpy.data.collections.get(name) or bpy.data.collections.new(name)
    if c.name not in [x.name for x in sc.collection.children]:
        sc.collection.children.link(c)
    return c


def layer_coll(vl, name):
    def walk(lc):
        if lc.name == name:
            return lc
        for ch in lc.children:
            r = walk(ch)
            if r:
                return r
        return None
    return walk(vl.layer_collection)


# ------------------------------------------------------------------ parts of a link mesh
def part_faces(me):
    """{part index: [face indices]}, and the decal part indices"""
    lay = me.attributes.get('part')
    dec = list(me.get("decals", []))
    out = {}
    if lay is None:
        return {0: list(range(len(me.polygons)))}, set()
    pv = np.empty(len(me.polygons), np.int64)
    lay.data.foreach_get("value", pv)
    for i, k in enumerate(pv.tolist()):
        out.setdefault(k, []).append(i)
    return out, set(i for i, d in enumerate(dec) if d)


def part_sizes(me):
    pf, dec = part_faces(me)
    co = np.empty(len(me.vertices) * 3, np.float64)
    me.vertices.foreach_get("co", co)
    co = co.reshape(-1, 3)
    sizes = {}
    for k, fs in pf.items():
        vi = sorted({v for f in fs for v in me.polygons[f].vertices})
        p = co[vi]
        sizes[k] = float(np.linalg.norm(p.max(0) - p.min(0)))
    return sizes, pf, dec


# ------------------------------------------------------------------ UCX
def kmeans(P, k, iters=25, seed=3):
    rng = np.random.default_rng(seed)
    C = P[rng.choice(len(P), k, replace=False)].copy()
    lab = np.zeros(len(P), np.int64)
    for _ in range(iters):
        d = ((P[:, None, :] - C[None, :, :]) ** 2).sum(-1)
        lab = d.argmin(1)
        for j in range(k):
            m = lab == j
            if m.any():
                C[j] = P[m].mean(0)
    return lab


def extreme_points(P, n=96):
    """the points of P that lie furthest along n directions spread over the sphere: a hull of them has at most n
    vertices and keeps the shape (Unreal simplifies big collision hulls anyway)"""
    k = np.arange(n) + 0.5
    phi = np.arccos(1 - 2 * k / n)
    th = math.pi * (1 + 5 ** 0.5) * k
    D = np.stack([np.cos(th) * np.sin(phi), np.sin(th) * np.sin(phi), np.cos(phi)], 1)
    idx = np.unique((P @ D.T).argmax(0))
    return P[idx]


def convex_err(me):
    """how far (m) the worst vertex lies outside the plane of any face (0 for a convex hull)"""
    co = np.empty(len(me.vertices) * 3, np.float64)
    me.vertices.foreach_get("co", co)
    P = co.reshape(-1, 3)
    worst = 0.0
    for p in me.polygons:
        n = np.array(p.normal, np.float64)
        d = float(n @ P[p.vertices[0]])
        worst = max(worst, float((P @ n - d).max()))
    return worst


def hull_mesh(name, pts):
    """a convex hull of the points, all triangles (fix round, CRITIQUE_2 nice 11: checked convex; when the hull op
    leaves a vertex outside a face it is built again from slightly jittered points)"""
    pts = extreme_points(np.asarray(pts, np.float64)) if len(pts) > 96 else np.asarray(pts, np.float64)
    me = bpy.data.meshes.get(name) or bpy.data.meshes.new(name)
    rng = np.random.default_rng(11)
    for attempt in range(8):
        P = pts if attempt == 0 else pts + rng.normal(0.0, 2e-5 * attempt, pts.shape)
        bm = bmesh.new()
        vs = [bm.verts.new(Vector(p)) for p in P]
        bmesh.ops.convex_hull(bm, input=vs)
        loose = [v for v in bm.verts if not v.link_faces]
        bmesh.ops.delete(bm, geom=loose, context='VERTS')
        bmesh.ops.triangulate(bm, faces=bm.faces[:])
        bmesh.ops.recalc_face_normals(bm, faces=bm.faces[:])
        me.clear_geometry()
        bm.to_mesh(me)
        bm.free()
        if convex_err(me) < 1e-4:
            break
    me["convex_err_mm"] = round(convex_err(me) * 1000, 4)
    return me


def box_mesh(name, lo, hi):
    """a box hull (8 corners, 12 triangles) from lo to hi"""
    bm = bmesh.new()
    vs = [bm.verts.new((x, y, z)) for x in (lo[0], hi[0]) for y in (lo[1], hi[1]) for z in (lo[2], hi[2])]
    bmesh.ops.convex_hull(bm, input=vs)
    bmesh.ops.triangulate(bm, faces=bm.faces[:])
    bmesh.ops.recalc_face_normals(bm, faces=bm.faces[:])
    me = bpy.data.meshes.get(name) or bpy.data.meshes.new(name)
    me.clear_geometry()
    bm.to_mesh(me)
    bm.free()
    return me


def rail_ucx(sc, name, me_src, transform=PRE):
    """the rail kit's collision (fix round, CRITIQUE_2 nice 11): one box per part and per run along the rail (the
    flanges, the tray, the rack, each end buffer); parts under 6 cm (bolts, markings) left out"""
    pf, dec = part_faces(me_src)
    co = np.empty(len(me_src.vertices) * 3, np.float64)
    me_src.vertices.foreach_get("co", co)
    co = co.reshape(-1, 3)
    out = []
    for k in sorted(pf):
        if k in dec:
            continue
        vi = np.array(sorted({v for f in pf[k] for v in me_src.polygons[f].vertices}))
        P = co[vi]
        order = np.argsort(P[:, 0])
        xs = P[order, 0]
        cuts = np.where(np.diff(xs) > 0.05)[0]
        for grp in np.split(order, cuts + 1):
            Q = P[grp]
            lo, hi = Q.min(0), Q.max(0)
            if np.linalg.norm(hi - lo) < 0.06:
                continue
            nm = "UCX_%s_%02d" % (name, len(out))
            me = box_mesh(nm, lo, hi)
            me.transform(transform)
            ob = bpy.data.objects.new(nm, me)
            sc.collection.objects.link(ob)
            out.append(ob)
    return out


def link_hulls(me, sm):
    """[(name, mesh)] convex hulls of the link's vertex clusters (markings left out)"""
    sizes, pf, dec = part_sizes(me)
    keep = sorted({v for k, fs in pf.items() if k not in dec for f in fs for v in me.polygons[f].vertices})
    co = np.empty(len(me.vertices) * 3, np.float64)
    me.vertices.foreach_get("co", co)
    P = co.reshape(-1, 3)[keep]
    diag = float(np.linalg.norm(P.max(0) - P.min(0)))
    k = int(max(1, min(6, round(diag / 0.42))))
    if len(P) > 6000:
        P = P[np.random.default_rng(7).choice(len(P), 6000, replace=False)]
    lab = kmeans(P, k) if k > 1 else np.zeros(len(P), np.int64)
    out = []
    for j in range(k):
        Q = P[lab == j]
        if len(Q) < 8:
            continue
        nm = "UCX_%s_%02d" % (sm, len(out))
        out.append((nm, hull_mesh(nm, Q)))
    return out


def make_ucx():
    t0 = time.time()
    sc = bpy.data.scenes["DoorArm | Showroom"]
    C = coll(sc, COLL)
    for o in list(C.objects):
        bpy.data.objects.remove(o)
    objs = showroom_objects()
    rep = {}
    for spec in links():
        ob = objs.get(spec["name"])
        if ob is None:
            continue
        sm = sm_name(spec["name"])
        hulls = link_hulls(ob.data, sm)
        for nm, me in hulls:
            h = bpy.data.objects.new(nm, me)
            C.objects.link(h)
            h.parent = ob
            h.matrix_parent_inverse = Matrix.Identity(4)
            h.matrix_basis = Matrix.Identity(4)
            h.display_type = 'WIRE'
            h.hide_render = True
            h["ucx_for"] = sm
        rep[sm] = [len(me.vertices) for _, me in hulls]
        rep[sm + " convex_err_mm"] = max([me.get("convex_err_mm", 0.0) for _, me in hulls] or [0.0])
    C["not_exported_directly"] = "UCX hulls: exported with SM_DoorRobotArm_Links.fbx"
    for vl in sc.view_layers:
        lc = layer_coll(vl, COLL)
        if lc is not None:
            lc.hide_viewport = True               # the eye off: hidden in the viewport too
    return dict(hulls=rep, summary=dict(links=len([k for k in rep if not k.endswith("convex_err_mm")]),
                                        hulls=sum(len(v) for k, v in rep.items() if not k.endswith("convex_err_mm")),
                                        worst_convex_err_mm=max([v for k, v in rep.items() if k.endswith("convex_err_mm")] or [0.0])),
                seconds=round(time.time() - t0, 1))


# ------------------------------------------------------------------ LODs
def lod_copy(me, level, tri_target):
    """a copy of a link mesh for LOD level: markings and small parts dropped, then decimated toward tri_target"""
    out = me.copy()
    out.name = "%s_LOD%d" % (me.name, level)
    if level == 0:
        return out
    sizes, pf, dec = part_sizes(me)
    drop = [f for k, fs in pf.items() if k in dec or sizes[k] < LOD_SMALL[level] for f in fs]
    if drop and len(drop) < len(me.polygons):
        bm = bmesh.new()
        bm.from_mesh(out)
        bm.faces.ensure_lookup_table()
        bmesh.ops.delete(bm, geom=[bm.faces[i] for i in drop], context='FACES')
        loose = [v for v in bm.verts if not v.link_faces]
        bmesh.ops.delete(bm, geom=loose, context='VERTS')
        bm.to_mesh(out)
        bm.free()
    out.calc_loop_triangles()
    have = len(out.loop_triangles)
    ratio = min(1.0, tri_target / max(1, have))
    if ratio < 0.98:
        sc = bpy.context.scene
        tmp = bpy.data.objects.new("_lod_tmp", out)
        sc.collection.objects.link(tmp)
        m = tmp.modifiers.new("lod", 'DECIMATE')
        m.decimate_type = 'COLLAPSE'
        m.ratio = ratio
        m.use_collapse_triangulate = False
        bpy.context.view_layer.update()
        ev = tmp.evaluated_get(bpy.context.evaluated_depsgraph_get())
        red = bpy.data.meshes.new_from_object(ev)
        bpy.data.objects.remove(tmp)
        red.name = out.name
        for mat in out.materials:
            pass
        bpy.data.meshes.remove(out)
        out = red
    return out


# ------------------------------------------------------------------ the skeleton and the skinned mesh
def new_scene():
    sc = bpy.data.scenes.get("_export") or bpy.data.scenes.new("_export")
    for o in list(sc.collection.objects):
        sc.collection.objects.unlink(o)
    wm = bpy.context.window_manager
    if wm.windows:
        wm.windows[0].scene = sc
    sc.render.fps = 30
    sc.frame_start, sc.frame_end = 1, 120         # the doors' animation (the default 1-250 would go into the FBX)
    return sc


def build_skeleton(sc, ver, name):
    import dra_rigdata as rd
    arm = bpy.data.armatures.new(name + "_skeleton")
    ob = bpy.data.objects.new(name, arm)
    sc.collection.objects.link(ob)
    vl = bpy.context.view_layer
    vl.objects.active = ob
    for o in vl.objects:
        o.select_set(o is ob)
    bpy.ops.object.mode_set(mode='EDIT')
    root = arm.edit_bones.new("root")
    root.head = (0, 0, 0)
    root.tail = (0, 0.1, 0)
    for spec in links(ver):
        b = arm.edit_bones.new(bone_name(spec["name"]))
        j = Vector(spec["joint_s"])
        b.head = j
        b.tail = j + Vector((0, 0.1, 0))
        b.roll = 0.0
        b.parent = arm.edit_bones[bone_name(spec["parent"])] if spec["parent"] else root
        b.use_connect = False
    bpy.ops.object.mode_set(mode='OBJECT')
    return ob


def curve_weights(cu_ob, R0inv, kind):
    """a cable / hose curve at rest -> (mesh in the A frame, per-vertex [(bone, weight)])"""
    dg = bpy.context.evaluated_depsgraph_get()
    ev = cu_ob.evaluated_get(dg)
    me = bpy.data.meshes.new_from_object(ev)
    me.transform(R0inv @ cu_ob.matrix_world)
    n_rings = len(cu_ob.data.splines[0].bezier_points)
    # the control points' links: cables 0,1 upper arm; 2 0.65/0.35; 3 0.30/0.70; 4,5 lower arm. hoses: 0,1 the cylinder,
    # 2 the lower arm
    if kind == "cable":
        w = [("upper_arm", 1.0), ("upper_arm", 1.0), ("upper_arm", 0.65), ("upper_arm", 0.30), ("upper_arm", 0.0),
             ("upper_arm", 0.0)]
        other = "lower_arm"
    else:
        w = [("j3_cylinder", 1.0), ("j3_cylinder", 1.0), ("j3_cylinder", 0.0)]
        other = "lower_arm"
    # the curve's mesh: rings in order along the spline; ring r is at parameter r / (rings - 1) * (points - 1)
    seg = cu_ob.data.resolution_u
    n_pts = len(w)
    ring_count = (n_pts - 1) * seg + 1
    nv = len(me.vertices)
    per = max(1, nv // ring_count)
    out = []
    for vi in range(nv):
        r = min(ring_count - 1, vi // per)
        t = r / max(1, ring_count - 1) * (n_pts - 1)
        a = int(math.floor(t))
        b = min(n_pts - 1, a + 1)
        f = t - a
        wa = w[a][1] * (1 - f) + w[b][1] * f
        out.append([(w[0][0], wa), (other, 1.0 - wa)])
    return me, out


def skinned_mesh(sc, ver, rig, level, name, budget):
    """every link of version ver (LOD level) joined into one mesh, rigid on its bone; the curves blended"""
    import dra_rigdata as rd
    objs = showroom_objects()
    root = bpy.data.objects["Showroom | arm %s" % ver]
    R0inv = root.matrix_world.inverted()
    pieces = []
    total0 = sum(len_tris(bpy.data.meshes[rd.mesh_name(s["name"])]) for s in links(ver))
    for spec in links(ver):
        src = bpy.data.meshes[rd.mesh_name(spec["name"])]
        target = len_tris(src) * (LOD_RATIO[level] if level else 1.0)
        me = lod_copy(src, level, target)
        me.transform(Matrix.Translation(Vector(spec["joint_s"])))
        pieces.append((me, [(bone_name(spec["name"]), 1.0)]))
    curves = [o for o in bpy.data.scenes["DoorArm | Showroom"].objects
              if o.name.startswith("Showroom | arm %s |" % ver) and o.type == 'CURVE']
    for cu in curves:
        kind = "cable" if "cable" in cu.name else "hose"
        me, wts = curve_weights(cu, R0inv, kind)
        if level:
            pass
        pieces.append((me, wts))
    bm = bmesh.new()
    mats = []
    dvert = bm.verts.layers.deform.verify()
    groups = {}
    for me, wts in pieces:
        idx_map = []
        for m in me.materials:
            m = merged_mat(m)
            if m is None:
                idx_map.append(0)
                continue
            if m.name not in [x.name for x in mats]:
                mats.append(m)
            idx_map.append([x.name for x in mats].index(m.name))
        base = len(bm.verts)
        tmp = bmesh.new()
        tmp.from_mesh(me)
        uv_names = [l.name for l in me.uv_layers]
        for f in tmp.faces:
            f.material_index = idx_map[f.material_index] if idx_map else 0
        tmp_me = bpy.data.meshes.new("_tmp")
        tmp.to_mesh(tmp_me)
        tmp.free()
        bm.from_mesh(tmp_me)           # appends
        bpy.data.meshes.remove(tmp_me)
        bm.verts.ensure_lookup_table()
        for vi in range(base, len(bm.verts)):
            local = vi - base
            ws = wts[local] if len(wts) == len(me.vertices) else wts[0] if len(wts) == 1 else wts[min(local, len(wts) - 1)]
            if isinstance(ws, tuple):
                ws = [ws]
            for bn, wv in (ws if isinstance(ws, list) else [ws]):
                if wv <= 0.0:
                    continue
                gi = groups.setdefault(bn, len(groups))
                bm.verts[vi][dvert][gi] = wv
        bpy.data.meshes.remove(me)
    out = bpy.data.meshes.new(name)
    bm.to_mesh(out)
    bm.free()
    for m in mats:
        out.materials.append(m)
    out.transform(SKT)
    ob = bpy.data.objects.new(name, out)
    sc.collection.objects.link(ob)
    for bn, gi in sorted(groups.items(), key=lambda kv: kv[1]):
        ob.vertex_groups.new(name=bn)
    ob.parent = rig
    mod = ob.modifiers.new("Armature", 'ARMATURE')
    mod.object = rig
    out.calc_loop_triangles()
    return ob, len(out.loop_triangles)


def settle(rig):
    """after the armature's rest changed (turned and scaled): every pose bone back to its rest, evaluated, so the FBX
    nodes' default pose is the bind pose"""
    for pb in rig.pose.bones:
        pb.matrix_basis = Matrix.Identity(4)
    rig.data.pose_position = 'POSE'
    bpy.context.view_layer.update()
    rig.data.pose_position = 'REST'
    bpy.context.view_layer.update()
    rig.data.pose_position = 'POSE'
    bpy.context.view_layer.update()


def len_tris(me):
    me.calc_loop_triangles()
    return len(me.loop_triangles)


EXPECT = {}


def expect_of(objs, skeletal, anim):
    """what the file should hold once imported again (metres, the turned frame)"""
    k = 0.01 if skeletal else 1.0
    pts, tris, mats, bones, meshes, ucx = [], 0, set(), [], 0, 0
    for o in objs:
        if o.type == 'ARMATURE':
            bones += [b.name for b in o.data.bones]
        elif o.type == 'MESH':
            if o.name.startswith("UCX_"):
                ucx += 1
                continue
            meshes += 1
            me = o.data
            me.calc_loop_triangles()
            tris += len(me.loop_triangles)
            mats |= {m.name for m in me.materials if m}
            M = o.matrix_world
            co = np.empty(len(me.vertices) * 3, np.float64)
            me.vertices.foreach_get("co", co)
            co = co.reshape(-1, 3)
            if len(co):
                A = np.array(M)
                w = co @ A[:3, :3].T + A[:3, 3]
                pts.append(w.min(0))
                pts.append(w.max(0))
    e = dict(bones=sorted(bones), meshes=meshes, tris=tris, materials=sorted(mats), ucx=ucx, anim=bool(anim))
    if pts:
        P = np.array(pts) * k
        e["bounds_m"] = [round(float(x), 4) for x in list(P.min(0)) + list(P.max(0))]
    return e


def fbx_out(path, objs, types, anim=False, skeletal=True):
    EXPECT[os.path.basename(path)] = expect_of(objs, skeletal, anim)
    vl = bpy.context.view_layer
    for o in vl.objects:
        o.select_set(False)
    for o in objs:
        o.select_set(True)
    vl.objects.active = objs[0]
    kw = dict(filepath=path, use_selection=True, object_types=types, use_mesh_modifiers=False, mesh_smooth_type="FACE",
              add_leaf_bones=False, use_armature_deform_only=False, axis_forward="-Y", axis_up="Z",
              apply_unit_scale=True, apply_scale_options="FBX_SCALE_NONE", use_space_transform=True,
              bake_space_transform=False, path_mode="STRIP", use_metadata=False, use_tspace=True)
    kw["global_scale"] = 0.01 if skeletal else 1.0
    if anim:
        kw.update(bake_anim=True, bake_anim_use_all_bones=True, bake_anim_use_nla_strips=False,
                  bake_anim_use_all_actions=False, bake_anim_force_startend_keying=True, bake_anim_step=1.0,
                  bake_anim_simplify_factor=0.0)
    else:
        kw["bake_anim"] = False
    bpy.ops.export_scene.fbx(**kw)


# ------------------------------------------------------------------ the door animation, baked from the rig
def bake_action(rig, ver, key):
    """keys on rig's pose bones for frames 1-120 from the ship-fit arm (key, ver): basis = rest^-1 pose, in the A frame,
    locations x100 (the data is in centimetres)"""
    import dra_rigdata as rd
    sf = bpy.data.scenes["DoorArm | Ship fit"]
    ctl = bpy.data.objects["DoorArm | CONTROLS"]
    wm = bpy.context.window_manager
    prev = wm.windows[0].scene if wm.windows else None
    if wm.windows:
        wm.windows[0].scene = sf
    ctl["version"] = 0 if ver == "A" else 1
    ctl.update_tag()
    prefix = "DoorArm_%s_%s" % (key, ver)
    root = bpy.data.objects[prefix]
    specs = links(ver)
    rest = {"root": Matrix.Identity(4)}
    for s in specs:
        rest[bone_name(s["name"])] = Matrix.Translation(Vector(s["joint_s"]))
    parent = {bone_name(s["name"]): (bone_name(s["parent"]) if s["parent"] else "root") for s in specs}
    frames = {}
    for f in range(1, 121):
        sf.frame_set(f)
        dg = bpy.context.evaluated_depsgraph_get()
        Rinv = root.evaluated_get(dg).matrix_world.inverted()
        pose = {"root": Matrix.Identity(4)}
        for s in specs:
            o = bpy.data.objects["%s | %s" % (prefix, s["name"])]
            pose[bone_name(s["name"])] = Rinv @ o.evaluated_get(dg).matrix_world
        frames[f] = pose
    ctl["version"] = 0
    ctl.update_tag()
    sf.frame_set(1)
    if wm.windows and prev is not None:
        wm.windows[0].scene = prev
    act = bpy.data.actions.new("%s_%s_OpenClose" % (ver, key))
    rig.animation_data_create()
    rig.animation_data.action = act
    for pb in rig.pose.bones:
        pb.rotation_mode = 'QUATERNION'
    for f, pose in frames.items():
        for pb in rig.pose.bones:
            b = pb.name
            if b == "root":
                basis = Matrix.Identity(4)
            else:
                p = parent[b]
                basis = (rest[p].inverted() @ rest[b]).inverted() @ (pose[p].inverted() @ pose[b])
            loc, rot, _sc = basis.decompose()
            pb.location = loc * 100.0
            pb.rotation_quaternion = rot
            pb.keyframe_insert("location", frame=f)
            pb.keyframe_insert("rotation_quaternion", frame=f)
    return act, frames


# ------------------------------------------------------------------ material sections
# fix round (CRITIQUE_2 nice 12): the cables share the rubber's section and the lamp the glass's, so the skeletal mesh
# has 7 material sections in Unreal (paint white, orange, gunmetal, steel, rubber, glass, black metal where used)
MERGE_MATS = {"M_DoorArm_Cable": "M_DoorArm_Rubber", "M_DoorArm_Lamp": "M_DoorArm_Glass"}


def merged_mat(m):
    if m is None:
        return None
    t = bpy.data.materials.get(MERGE_MATS.get(m.name, m.name))
    return t or m


def merge_slots(me):
    """map the mesh's slots through MERGE_MATS and join slots that end up with the same material"""
    if not me.materials:
        return
    target = [merged_mat(m) for m in me.materials]
    names = [m.name if m else None for m in target]
    uniq = []
    for n in names:
        if n not in uniq:
            uniq.append(n)
    if uniq == [m.name if m else None for m in me.materials]:
        return
    remap = np.array([uniq.index(n) for n in names], np.int64)
    idx = np.empty(len(me.polygons), np.int64)
    me.polygons.foreach_get("material_index", idx)
    keep = {n: m for n, m in zip(names, target)}
    me.materials.clear()
    for n in uniq:
        me.materials.append(keep[n])
    me.polygons.foreach_set("material_index", remap[np.clip(idx, 0, len(remap) - 1)])


# ------------------------------------------------------------------ statics
def static_obj(sc, name, me_src, matrix=None, transform=PRE):
    me = me_src.copy()
    me.name = name
    merge_slots(me)
    if matrix is not None:
        me.transform(matrix)
    me.transform(transform)
    ob = bpy.data.objects.new(name, me)
    sc.collection.objects.link(ob)
    return ob


def box_ucx(sc, name, me_src, transform=PRE):
    """one convex hull of a mesh (for the small statics)"""
    co = np.empty(len(me_src.vertices) * 3, np.float64)
    me_src.vertices.foreach_get("co", co)
    P = co.reshape(-1, 3)
    if len(P) > 4000:
        P = P[np.random.default_rng(5).choice(len(P), 4000, replace=False)]
    me = hull_mesh(name, P)
    me.transform(transform)
    ob = bpy.data.objects.new(name, me)
    sc.collection.objects.link(ob)
    return ob


# ------------------------------------------------------------------ the energy chain as a small skeletal mesh
def chain_export(sc, key, rep):
    sf = bpy.data.scenes["DoorArm | Ship fit"]
    wm = bpy.context.window_manager
    lk = sorted([o for o in sf.objects if o.name.startswith("DoorArm_%s | chain link" % key)], key=lambda o: o.name)
    if not lk:
        return
    if wm.windows:
        wm.windows[0].scene = sf
    frames = {}
    for f in range(1, 121):
        sf.frame_set(f)
        dg = bpy.context.evaluated_depsgraph_get()
        frames[f] = [o.evaluated_get(dg).matrix_world.copy() for o in lk]
    sf.frame_set(1)
    if wm.windows:
        wm.windows[0].scene = sc
    name = "SK_DoorArm_Chain_%s" % key
    arm = bpy.data.armatures.new(name + "_skeleton")
    rig = bpy.data.objects.new(name, arm)
    sc.collection.objects.link(rig)
    vl = bpy.context.view_layer
    vl.objects.active = rig
    for o in vl.objects:
        o.select_set(o is rig)
    bpy.ops.object.mode_set(mode='EDIT')
    root = arm.edit_bones.new("root")
    root.head, root.tail = (0, 0, 0), (0, 0.1, 0)
    rest = frames[1]
    for i, M in enumerate(rest):
        b = arm.edit_bones.new("link_%02d" % i)
        h = M.translation
        b.head = h
        b.tail = h + M.to_3x3() @ Vector((0, 0.05, 0))
        b.align_roll(M.to_3x3() @ Vector((0, 0, 1)))
        b.parent = root
    bpy.ops.object.mode_set(mode='OBJECT')
    rest_m = {bb.name: bb.matrix_local.copy() for bb in arm.bones}
    chain_me = lk[0].data
    bm = bmesh.new()
    dvert = bm.verts.layers.deform.verify()
    for i, M in enumerate(rest):
        tmp = chain_me.copy()
        tmp.transform(M)
        base = len(bm.verts)
        bm.from_mesh(tmp)
        bpy.data.meshes.remove(tmp)
        bm.verts.ensure_lookup_table()
        for vi in range(base, len(bm.verts)):
            bm.verts[vi][dvert][i] = 1.0
    me = bpy.data.meshes.new(name + "_mesh")
    bm.to_mesh(me)
    bm.free()
    for m in chain_me.materials:
        me.materials.append(m)
    merge_slots(me)
    mo = bpy.data.objects.new(name + "_mesh", me)
    sc.collection.objects.link(mo)
    for i in range(len(rest)):
        mo.vertex_groups.new(name="link_%02d" % i)
    act = bpy.data.actions.new("Chain_%s_OpenClose" % key)
    rig.animation_data_create()
    rig.animation_data.action = act
    for pb in rig.pose.bones:
        pb.rotation_mode = 'QUATERNION'
    for f, Ms in frames.items():
        for i, M in enumerate(Ms):
            pb = rig.pose.bones["link_%02d" % i]
            R = rest_m["link_%02d" % i]
            P = M @ (rest[i].inverted() @ R)            # the bone moves with its link
            basis = R.inverted() @ P
            loc, rot, _s = basis.decompose()
            pb.location = loc
            pb.rotation_quaternion = rot
            pb.keyframe_insert("location", frame=f)
            pb.keyframe_insert("rotation_quaternion", frame=f)
    # into centimetres and turned, as the arm (rest and keys conjugated: locations x100 after the turn)
    arm.transform(SKT)
    me.transform(SKT)
    for fc_owner in [act]:
        pass
    # the keys were made in metres in the untransformed frame: scale the location curves x100 and turn them with the bones
    # (a pure turn about Z of every bone's rest leaves rotation keys unchanged; locations are bone-local)
    from bpy_extras.anim_utils import action_get_channelbag_for_slot
    cb = action_get_channelbag_for_slot(act, rig.animation_data.action_slot)
    for fc in cb.fcurves:
        if fc.data_path.endswith(".location"):
            for kp in fc.keyframe_points:
                kp.co[1] *= 100.0
                kp.handle_left[1] *= 100.0
                kp.handle_right[1] *= 100.0
    mo.parent = rig
    mm = mo.modifiers.new("Armature", 'ARMATURE')
    mm.object = rig
    act_keep = rig.animation_data.action
    rig.animation_data.action = None
    settle(rig)
    fbx_out(os.path.join(FBX, name + ".fbx"), [rig, mo], {"ARMATURE", "MESH"})
    rig.animation_data.action = act_keep
    bpy.context.view_layer.update()
    fbx_out(os.path.join(FBX, "A_DoorArm_Chain_%s_OpenClose.fbx" % key), [rig], {"ARMATURE"}, anim=True)
    rep["chain_" + key] = dict(bones=len(arm.bones), tris=len_tris(me))


# ------------------------------------------------------------------ everything
def export(only=None):
    t0 = time.time()
    import dra_rigdata as rd
    os.makedirs(FBX, exist_ok=True)
    rep = {"files": [], "lods": {}, "skeleton": {}}
    src_scene = bpy.context.window_manager.windows[0].scene if bpy.context.window_manager.windows else None
    sc = new_scene()
    want = set(only) if only else {"links", "kit", "skeletal", "anim", "chain"}
    # 1. the links as statics with their UCX hulls
    if "links" in want:
        objs = []
        ucx_by = {}
        for o in bpy.data.objects:
            if o.name.startswith("UCX_") and o.get("ucx_for"):
                ucx_by.setdefault(o["ucx_for"], []).append(o)
        for spec in links():
            me = bpy.data.meshes[rd.mesh_name(spec["name"])]
            sm = sm_name(spec["name"])
            objs.append(static_obj(sc, sm, me))
            for h in sorted(ucx_by.get(sm, []), key=lambda o: o.name):
                objs.append(static_obj(sc, h.name, h.data))
        p = os.path.join(FBX, "SM_DoorRobotArm_Links.fbx")
        fbx_out(p, objs, {"MESH"}, skeletal=False)
        rep["files"].append(os.path.basename(p))
        rep["links_static"] = dict(meshes=len([o for o in objs if not o.name.startswith("UCX_")]),
                                   ucx=len([o for o in objs if o.name.startswith("UCX_")]))
    # 2. the kit: rail kits, fixtures (static, ship coordinates; fixtures at the attach point)
    if "kit" in want:
        for key in ("AFT", "FORE"):
            for nm, src in (("SM_DoorArm_RailKit_%s" % key, "DoorArm | rail kit %s" % key),
                            ("SM_DoorArm_RailBeam_%s" % key, "DoorArm | rail beam %s" % key),
                            ("SM_DoorArm_FixtureA_%s" % key, "DoorArm_%s | door fixture A" % key),
                            ("SM_DoorArm_FixtureB_%s" % key, "DoorArm_%s | door fixture B" % key)):
                me = bpy.data.meshes.get(src)
                if me is None:
                    continue
                o = static_obj(sc, nm, me)
                objs2 = [o]
                if "Fixture" in nm:
                    objs2.append(box_ucx(sc, "UCX_%s_00" % nm, me))
                else:
                    objs2 += rail_ucx(sc, nm, me)
                p = os.path.join(FBX, nm + ".fbx")
                fbx_out(p, objs2, {"MESH"}, skeletal=False)
                rep["files"].append(os.path.basename(p))
    # 3. the skeletal mesh per version and LOD, 4. the door animations
    for ver in ("A", "B"):
        if not ({"skeletal", "anim"} & want):
            break
        name = "SK_DoorRobotArm_%s" % ver
        rig = build_skeleton(sc, ver, name + "_rig")
        rig_rest = rig
        if "skeletal" in want:
            tris = {}
            for level in (0, 1, 2):
                mo, n = skinned_mesh(sc, ver, rig, level, name if level == 0 else "%s_LOD%d" % (name, level), None)
                tris["LOD%d" % level] = n
                if level == 0 and not rig.get("_scaled"):
                    rig.data.transform(SKT)
                    rig["_scaled"] = True
                settle(rig)
                p = os.path.join(FBX, (name if level == 0 else "%s_LOD%d" % (name, level)) + ".fbx")
                fbx_out(p, [rig, mo], {"ARMATURE", "MESH"})
                rep["files"].append(os.path.basename(p))
                bpy.data.objects.remove(mo)
            rep["lods"][ver] = tris
            rep["skeleton"][ver] = dict(bones=len(rig.data.bones), names=[b.name for b in rig.data.bones])
        if "anim" in want:
            if not rig.get("_scaled"):
                rig.data.transform(SKT)
                rig["_scaled"] = True
            settle(rig)
            for key in ("AFT", "FORE"):
                act, frames = bake_action(rig, ver, key)
                p = os.path.join(FBX, "A_DoorRobotArm_%s_%s_OpenClose.fbx" % (ver, key))
                fbx_out(p, [rig], {"ARMATURE"}, anim=True)
                rep["files"].append(os.path.basename(p))
                rig.animation_data.action = None
        bpy.data.objects.remove(rig)
    # 5. the energy chains
    if "chain" in want:
        for key in ("AFT", "FORE"):
            chain_export(sc, key, rep)
            rep["files"] += ["SK_DoorArm_Chain_%s.fbx" % key, "A_DoorArm_Chain_%s_OpenClose.fbx" % key]
    # back to the file's own scene; the temporary one goes
    wm = bpy.context.window_manager
    if wm.windows and src_scene is not None:
        wm.windows[0].scene = src_scene
    bpy.data.scenes.remove(sc)
    old = {}
    ep = os.path.join(HERE, "export_report.json")
    if os.path.exists(ep):
        try:
            old = json.load(open(ep)).get("expect", {})
        except Exception:
            old = {}
    old.update(EXPECT)
    rep["expect"] = old
    rep["layout"] = write_layout(rep)
    rep["summary"] = dict(files=len(rep["files"]), lods=rep["lods"], seconds=round(time.time() - t0, 1))
    return rep


# ------------------------------------------------------------------ the layout json
def write_layout(rep):
    import dra_rigdata as rd
    with open(os.path.join(ROOT, "design.json")) as f:
        D = json.load(f)
    rr = {}
    p = os.path.join(HERE, "rig_report.json")
    if os.path.exists(p):
        rr = json.load(open(p))
    K = D["kinematics"]
    ue = lambda v: [round(100 * v[0], 2), round(-100 * v[1], 2), round(100 * v[2], 2)]
    bones = []
    for s in rd.LINKS:
        j = list(Vector(s["joint_s"]))
        bones.append(dict(link=s["name"], bone=bone_name(s["name"]), parent=bone_name(s["parent"]) if s["parent"] else "root",
                          version=s["tool"] or "both", pivot_A_m=[round(c, 4) for c in j],
                          static_mesh=sm_name(s["name"])))
    chan = {"carriage": "location.x = c (ship x of the turntable centre)", "turret": "rotation_euler.z = q1",
            "lower arm": "rotation_euler.y = -q2", "upper arm": "rotation_euler.y = -q3", "forearm": "rotation_euler.x = q4",
            "wrist head": "rotation_euler.y = q5", "flange": "rotation_euler.x = q6", "J3 cylinder": "rotation_euler.y = -cyl",
            "J3 piston": "location.x/z = rest + u * (L - L0) along the rod",
            "carriage pinion": "rotation_euler.y = (c - c_closed) / %.4f" % (0.040 * rd.S)}
    for b in bones:
        b["channel"] = chan.get(b["link"], "fixed to its parent (fingers: grip)")
    lay = {
        "asset": "door robot arm (models/door-robot-arm); version B is installed in the ship by build/dra_ship_install.py "
                 "(README, 'In the ship')",
        "scale_up": {k: D.get("scale_up", {}).get(k) for k in ("date", "S", "rail_y", "lengths_m", "length_changes_pct",
                                                               "plate_B", "pad_B", "d_t", "j3_limits", "j1_limits")},
        "units": "metres in Blender, Z up; FBX in centimetres; Unreal = (x, -y, z) * 100 of the Blender data (forward -Y, "
                 "up Z, data turned -90 deg about Z; import with Convert Scene + Force Front X Axis)",
        "arm_frame_A": K["frame_A"],
        "roots": {key: dict(object="DoorArm_%s_<A|B>" % key, location_m=[0.0, D["carriage"]["centre_y"],
                                                                         D["carriage"]["underside_z"]],
                            rotation_deg=[180.0, 0.0, 0.0], location_ue_cm=ue([0.0, D["carriage"]["centre_y"],
                                                                               D["carriage"]["underside_z"]]),
                            note="the arm hangs: its frame A turned 180 deg about x under the rail")
                  for key in ("AFT", "FORE")},
        "bones": bones,
        "joints": {k: dict(axis=v["axis"], origin_A_m=v["origin"], limits_deg=v["limits"],
                           used_deg={a: rr.get("arms", {}).get(a, {}).get("joints_deg", {}).get(k)
                                     for a in ("AFT_A", "FORE_A", "AFT_B", "FORE_B")})
                   for k, v in K["joints"].items()},
        "carriage": dict(travel_x_m=D["carriage"]["travel_ends"], schedule="c = A + B * door_x",
                         A_B={k: [D["doors"][k]["carriage"]["A"], D["doors"][k]["carriage"]["B"]] for k in ("AFT", "FORE")},
                         rail_y_m=D["scene"]["rail"]["y"], rail_z_m=D["scene"]["rail"]["z"],
                         pinion_pitch_radius_m=round(0.040 * rd.S, 4)),
        "doors": {k: dict(object=D["doors"][k]["door"], closed_location_m=D["doors"][k]["closed_location"],
                          attach_local_m=D["doors"][k]["attach_local"], face_normal_local=D["doors"][k]["normal_local"],
                          tool_axes_world=D["doors"][k]["tool_axes_world"],
                          tool_axes_world_A=D["doors"][k].get("tool_axes_world_A"),
                          tool_roll_A_deg=D["doors"][k].get("tool_roll_A_deg"),
                          wrist_centre_off_face_m={v: D["tools"][v]["d_t"] for v in ("A", "B")},
                          frames=D["scene"]["door_frames"]) for k in ("AFT", "FORE")},
        "ik_closed_form": {"note": "DESIGN.md section 5; the rig's drivers and build/dra_ik.py compute it",
                           "L2": K["L2"], "L3": K["L3"], "beta_rad": K["beta_rad"], "phi_l0_rad": K["phi_l0_rad"],
                           "j2_r": K["j2_r"], "j2_h": K["j2_h"], "rod": K["rod"]},
        "version_switch": dict(control="DoorArm | CONTROLS", property="version", values={"0": "A gripper (grab limbs)",
                                                                                        "1": "B plate (red sketch)"},
                               grip="CONTROLS['grip'] 0 open .. 1 closed on the knob (default 1)",
                               hides="hide_viewport and hide_render of every object of the other version (arms, curves, "
                                     "fixtures); collections 'DoorArm | version A' / 'DoorArm | version B'",
                               unreal="two skeletal meshes on one skeleton layout: SK_DoorRobotArm_A / _B; show one"),
        "fbx": rep.get("files", []),
        "lods_tris": rep.get("lods", {}),
        "materials": None,
    }
    sp = os.path.join(HERE, "surf_report.json")
    if os.path.exists(sp):
        lay["materials"] = json.load(open(sp)).get("table")
    out = os.path.join(ROOT, "DoorRobotArm_layout.json")
    with open(out, "w") as f:
        json.dump(lay, f, indent=1)
    return os.path.basename(out)