"""Door robot arm: real surfaces (Coder stage 3; stages 'surf' and 'bake' of dra_build.py).

What the pictures show (ref1-ref3): satin white powder-coat paint with a fine texture, dark gunmetal and black metal on
the bearings, inner plates and the carriage's stacked plates, black rubber cables, bright steel screws and pins, small
orange accent rings and markings. The ship is a clean, grayish-white vessel: wear is low.

Made from real texture data, as the project's other models (docking arm surfaces.py, pressure door pd_look.py):
  * Poly Haven `blue_metal_plate` (CC0, `source_polyhaven/`, fetched with curl): a scan of a painted steel plate. Its
    relief, roughness and scuffs lie under the white and orange paint (its colour is not used: the tint is), and under
    the dark metals as light and dark;
  * the user's own pack (models/vessel-interior/textures/ModSci_Engineer, read in place): StainedSteel for the bright
    steel, ScratchedMetal's roughness for the rubber, Glass_DirtRough as smudges, T_Scratch_Norm for fine scratches;
  * maps baked from the model itself into one atlas on UV0 (`UVMap`), shared by every arm mesh of both versions:
    textures/T_DoorArm_Masks.png (R ambient occlusion, G convex edges, B crevices) and textures/T_DoorArm_BevelNormal.png
    (hard edges rounded by 8 mm, tangent space, OpenGL). They put a little shade in crevices, chip the paint on exposed
    edges (sparingly) and let edges catch the light;
  * a fine orange-peel bump (procedural) on the powder coat.
UV1 (`UV_Tile`, metres, box projection) carries the tiling sets.

The rail kits, chain link, door fixtures and the showroom rail wear the same recipes without the baked atlas (their own
materials, M_DoorArmKit_*); the cables and hoses (curves) a rubber in object space (M_DoorArm_Cable).

    surf   pack the atlas (UV0 of every arm mesh in one square), dress every material, the stand-in colours for the ship
           context's missing images, the light rig '90 STUDIO | lights (not part of the asset)' in the ship-fit scene
    bake   bake the two atlas maps (Cycles, GPU) into textures/ (run after surf; each in its own call)
"""
import math
import os
import time
import json
import numpy as np
import bpy
import bmesh
from mathutils import Vector

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
TEX = os.path.join(ROOT, "textures")
SCAN = os.path.join(ROOT, "source_polyhaven", "textures", "blue_metal_plate")
PACK = os.path.normpath(os.path.join(ROOT, "..", "vessel-interior", "textures", "ModSci_Engineer", "Textures"))
ATLAS = 4096
MASKS, BEVEL = "T_DoorArm_Masks", "T_DoorArm_BevelNormal"
EDGE_R, BEVEL_R = 0.010, 0.008
LIGHTS = "90 STUDIO | lights (not part of the asset)"
KIT_MATS = {"M_DoorArm_Gunmetal": "M_DoorArmKit_Gunmetal", "M_DoorArm_Steel": "M_DoorArmKit_Steel",
            "M_DoorArm_Rubber": "M_DoorArmKit_Rubber", "M_DoorArm_Orange": "M_DoorArmKit_Orange"}


def scan(m):
    return os.path.join(SCAN, "blue_metal_plate_%s_2k.jpg" % m)


def pack(*p):
    return os.path.join(PACK, *p)


# ------------------------------------------------------------------ which meshes
def arm_meshes():
    import dra_rigdata as rd
    return [bpy.data.meshes[rd.mesh_name(l["name"])] for l in rd.LINKS if bpy.data.meshes.get(rd.mesh_name(l["name"]))]


def kit_meshes():
    names = ["DoorArm | rail kit AFT", "DoorArm | rail kit FORE", "DoorArm | chain link", "Showroom | floor rail"]
    names += ["DoorArm_%s | door fixture %s" % (k, v) for k in ("AFT", "FORE") for v in ("A", "B")]
    return [bpy.data.meshes[n] for n in names if bpy.data.meshes.get(n)]


def decal_faces(me):
    """indices of the floating marking faces (they are not baked: no shade from them, none on them)"""
    dec = list(me.get("decals", []))
    lay = me.attributes.get('part')
    if not dec or lay is None:
        return set()
    pv = np.empty(len(me.polygons), np.int64)
    lay.data.foreach_get("value", pv)
    ids = [i for i, d in enumerate(dec) if d]
    return set(np.nonzero(np.isin(pv, ids))[0].tolist())


# ------------------------------------------------------------------ the atlas: UV0 of every arm mesh in one square
def pack_atlas(force=False):
    """once per build: a second pack would move the islands away from the baked maps"""
    t0 = time.time()
    meshes = arm_meshes()
    if not force and meshes and all(me.get("atlas_packed") for me in meshes):
        return dict(meshes=len(meshes), method="kept (already packed; the bake matches it)", seconds=0.0)
    sc = bpy.context.scene
    vl = bpy.context.view_layer
    tmp = []
    for me in meshes:
        me.uv_layers.active = me.uv_layers["UVMap"]
        o = bpy.data.objects.new("_uv_" + me.name, me)
        sc.collection.objects.link(o)
        tmp.append(o)
    vl.update()
    for o in vl.objects:
        o.select_set(False)
    for o in tmp:
        o.select_set(True)
    vl.objects.active = tmp[0]
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='SELECT')
    bpy.ops.uv.select_all(action='SELECT')
    bpy.ops.uv.average_islands_scale()        # one texel density on every mesh (the smart projects were per mesh)
    try:
        bpy.ops.uv.pack_islands(rotate=True, rotate_method='CARDINAL', scale=True, margin_method='FRACTION',
                                margin=6.0 / ATLAS, shape_method='CONCAVE')
        how = "concave"
    except Exception as ex:
        print("[dra] concave pack failed (%s), box pack" % ex)
        bpy.ops.uv.pack_islands(rotate=True, scale=True, margin_method='FRACTION', margin=6.0 / ATLAS, shape_method='AABB')
        how = "box"
    bpy.ops.object.mode_set(mode='OBJECT')
    for o in tmp:
        bpy.data.objects.remove(o)
    for me in meshes:
        me["atlas_packed"] = True
    # texel density (px per metre on the atlas) and fill
    a3 = a2 = 0.0
    for me in meshes:
        uv = np.empty(len(me.loops) * 2, np.float32)
        me.uv_layers["UVMap"].data.foreach_get("uv", uv)
        uv = uv.reshape(-1, 2)
        for p in me.polygons:
            q = uv[p.loop_start:p.loop_start + p.loop_total]
            a2 += 0.5 * abs(float(np.sum(q[:, 0] * np.roll(q[:, 1], -1) - np.roll(q[:, 0], -1) * q[:, 1])))
            a3 += p.area
    return dict(meshes=len(meshes), method=how, seconds=round(time.time() - t0, 1), fill=round(a2, 3),
                px_per_m=round(math.sqrt(a2 * ATLAS * ATLAS / max(a3, 1e-9)), 1))


# ------------------------------------------------------------------ node helpers (the docking arm's / pressure door's)
_RANGE = {}


def tex_image(path, colorspace="sRGB"):
    path = os.path.normpath(path)
    for im in bpy.data.images:
        if im.filepath and os.path.normcase(os.path.normpath(bpy.path.abspath(im.filepath))) == os.path.normcase(path):
            return im
    im = bpy.data.images.load(path, check_existing=True)
    im.colorspace_settings.name = colorspace
    return im


def tex_range(im, lo=3, hi=97):
    if im.name not in _RANGE:
        n, ch = im.size[0] * im.size[1], im.channels
        px = np.empty(n * ch, np.float32)
        im.pixels.foreach_get(px)
        px = px.reshape(n, ch)
        if ch >= 3:
            c = px[:, :3]
            if im.colorspace_settings.name == 'sRGB':
                c = np.where(c <= 0.04045, c / 12.92, ((c + 0.055) / 1.055) ** 2.4)
            v = c @ np.array([0.2126, 0.7152, 0.0722], np.float32)
        else:
            v = px[:, 0]
        _RANGE[im.name] = (round(float(np.percentile(v, lo)), 4), round(float(np.percentile(v, hi)), 4))
    return _RANGE[im.name]


class Tree:
    def __init__(self, mat):
        mat.use_nodes = True
        self.nt = mat.node_tree
        self.nt.nodes.clear()

    def node(self, kind, label=None, **props):
        n = self.nt.nodes.new(kind)
        if label:
            n.label = label
        for k, v in props.items():
            setattr(n, k, v)
        i = len(self.nt.nodes) - 1
        n.location = (-2400 + 240 * (i // 5), 600 - 300 * (i % 5))
        return n

    def link(self, a, b):
        self.nt.links.new(a, b)

    def put(self, sock, v):
        if isinstance(v, bpy.types.NodeSocket):
            self.nt.links.new(v, sock)
        else:
            sock.default_value = v

    def math(self, op, a, b=None, clamp=False, label=None):
        n = self.node('ShaderNodeMath', label, operation=op, use_clamp=clamp)
        self.put(n.inputs[0], a)
        if b is not None:
            self.put(n.inputs[1], b)
        return n.outputs[0]

    def ranged(self, v, lo, hi, label=None, f0=0.0, f1=1.0, smooth=False):
        n = self.node('ShaderNodeMapRange', label)
        if smooth:
            n.interpolation_type = 'SMOOTHSTEP'
        self.put(n.inputs["Value"], v)
        n.inputs["From Min"].default_value = f0
        n.inputs["From Max"].default_value = f1
        n.inputs["To Min"].default_value = lo
        n.inputs["To Max"].default_value = hi
        return n.outputs["Result"]

    def mixc(self, fac, a, b, label=None):
        n = self.node('ShaderNodeMix', label, data_type='RGBA')
        self.put(n.inputs[0], fac)
        self.put(n.inputs[6], a if isinstance(a, bpy.types.NodeSocket) else tuple(a) + (1.0,))
        self.put(n.inputs[7], b if isinstance(b, bpy.types.NodeSocket) else tuple(b) + (1.0,))
        return n.outputs[2]

    def mixf(self, fac, a, b, label=None):
        n = self.node('ShaderNodeMix', label, data_type='FLOAT')
        self.put(n.inputs[0], fac)
        self.put(n.inputs[2], a)
        self.put(n.inputs[3], b)
        return n.outputs[0]

    def coords(self, space):
        """'UV_Tile' / 'UVMap' (a UV layer) or 'OBJECT' (texture coordinates in object space, for the curves)"""
        if space == 'OBJECT':
            return self.node('ShaderNodeTexCoord').outputs["Object"]
        return self.node('ShaderNodeUVMap', uv_map=space).outputs["UV"]

    def mapped(self, src, scale, label=None):
        n = self.node('ShaderNodeMapping', label)
        self.nt.links.new(src, n.inputs["Vector"])
        n.inputs["Scale"].default_value = (scale, scale, scale)
        return n.outputs["Vector"]

    def tex(self, im, vec, label=None):
        n = self.node('ShaderNodeTexImage', label, image=im)
        self.nt.links.new(vec, n.inputs["Vector"])
        return n

    def up(self):
        g = self.node('ShaderNodeNewGeometry', "where dust settles")
        z = self.node('ShaderNodeSeparateXYZ')
        self.link(g.outputs["Normal"], z.inputs["Vector"])
        return self.ranged(z.outputs["Z"], 0.0, 1.0, "up-facing", 0.35, 0.95, smooth=True), g

    def noise(self, g_or_vec, scale, detail=3.0, rough=0.55, label=None):
        n = self.node('ShaderNodeTexNoise', label)
        n.inputs["Scale"].default_value = scale
        n.inputs["Detail"].default_value = detail
        n.inputs["Roughness"].default_value = rough
        self.link(g_or_vec, n.inputs["Vector"])
        return n.outputs["Fac"]


def masks_of(t, masks):
    if masks is None:
        return None, None, None
    sep = t.node('ShaderNodeSeparateColor')
    t.link(t.tex(masks, t.coords("UVMap"), "baked: R occlusion, G edges, B crevices").outputs["Color"], sep.inputs[0])
    return sep.outputs[0], sep.outputs[1], sep.outputs[2]


def bevel_normal(t, bevel):
    if bevel is None:
        return None
    nb = t.node('ShaderNodeNormalMap', "baked: rounded edges", uv_map="UVMap")
    t.link(t.tex(bevel, t.coords("UVMap"), "baked: bevel normal").outputs["Color"], nb.inputs["Color"])
    return nb.outputs["Normal"]


def join_normals(t, a, b):
    """two tangent-space normal results on one surface: add their deviations from the true normal"""
    if a is None or b is None:
        return a if a is not None else b
    g = t.node('ShaderNodeNewGeometry')
    s = t.node('ShaderNodeVectorMath', operation='ADD')
    t.link(a, s.inputs[0])
    t.link(b, s.inputs[1])
    d = t.node('ShaderNodeVectorMath', operation='SUBTRACT')
    t.link(s.outputs[0], d.inputs[0])
    t.link(g.outputs["Normal"], d.inputs[1])
    nz = t.node('ShaderNodeVectorMath', "both normals", operation='NORMALIZE')
    t.link(d.outputs[0], nz.inputs[0])
    return nz.outputs[0]


# ------------------------------------------------------------------ the recipes
def painted(mat, masks, bevel, tint, tile=1.25, shade=(0.93, 1.04), rough=(0.40, 0.56), coat=0.14, relief=0.30,
            wear=0.18, wear_col=(0.36, 0.36, 0.37), grime=0.18, dust=0.05, peel=0.035, metallic=0.0, space="UV_Tile",
            label="paint", grain=0.0, peel_scale=260.0, peel_dist=0.0004):
    """powder-coated steel: the tint shaded by the scan's light and dark, the scan's roughness and relief, a fine
    orange-peel bump, a satin clear coat; chips on exposed edges (low), shade and matte grime in crevices, a little dust
    on top"""
    t = Tree(mat)
    b = t.node('ShaderNodeBsdfPrincipled')
    out = t.node('ShaderNodeOutputMaterial')
    t.link(b.outputs[0], out.inputs["Surface"])
    u = t.coords(space)
    mp = t.mapped(u, 1.0 / tile, "one repeat = %.2f m" % tile)
    dim = tex_image(scan("diff"))
    lo, hi = tex_range(dim)
    bw = t.node('ShaderNodeRGBToBW')
    t.link(t.tex(dim, mp, "scan: painted steel (light and dark only)").outputs["Color"], bw.inputs[0])
    vm = t.node('ShaderNodeVectorMath', "tint x the scan's shade", operation='SCALE')
    vm.inputs[0].default_value = tint
    t.link(t.ranged(bw.outputs[0], shade[0], shade[1], "shade", lo, hi), vm.inputs["Scale"])
    col = vm.outputs["Vector"]
    if grain > 0:                             # the powder coat's fine speckle (ref2, ref3 close up)
        gv = t.node('ShaderNodeVectorMath', "fine grain", operation='SCALE')
        t.link(col, gv.inputs[0])
        t.link(t.ranged(t.noise(t.node('ShaderNodeTexCoord').outputs["Object"], 380.0, detail=2.0, rough=0.6,
                                label="grain"), 1.0 - grain, 1.0 + 0.4 * grain, None, 0.3, 0.7), gv.inputs["Scale"])
        col = gv.outputs["Vector"]
    rim = tex_image(scan("rough"), "Non-Color")
    rlo, rhi = tex_range(rim)
    ro = t.tex(rim, mp, "scan: roughness").outputs["Color"]
    rr = t.ranged(ro, rough[0], rough[1], "roughness range", rlo, rhi)
    dirt = t.tex(tex_image(pack("Base", "Glass", "T_Glass_DirtRough.png"), "Non-Color"), t.mapped(u, 1.0 / 2.4),
                 "smudges").outputs["Color"]
    rr = t.math('ADD', rr, t.math('MULTIPLY', dirt, 0.06), clamp=True, label="smudges dull the sheen")
    col = t.mixc(t.math('MULTIPLY', dirt, 0.10, clamp=True), col, (0.30, 0.29, 0.27), "faint smudges")
    ao, edge, crev = masks_of(t, masks)
    coat_w = coat
    metal = metallic
    if crev is not None and grime > 0:
        g = t.math('MULTIPLY', t.math('POWER', t.math('SUBTRACT', 1.0, crev), 1.4), grime, clamp=True, label="grime")
        col = t.mixc(g, col, (0.05, 0.048, 0.045), "shade in crevices")
        rr = t.math('ADD', rr, t.math('MULTIPLY', g, 0.35), clamp=True, label="grime is matte")
        coat_w = t.mixf(g, coat, 0.0)
    if ao is not None:
        col = t.mixc(t.ranged(ao, 0.30, 0.0, "baked occlusion", 0.0, 1.0), col, (0.02, 0.02, 0.02), "occlusion")
    upf, geo = t.up()
    if dust > 0:
        d = t.math('MULTIPLY', t.math('MULTIPLY', upf, t.ranged(t.noise(geo.outputs["Position"], 2.2, label="dust patches"),
                                                                0.0, 1.0, None, 0.35, 0.75, smooth=True)), dust, clamp=True)
        col = t.mixc(d, col, (0.50, 0.49, 0.46), "dust")
        rr = t.math('ADD', rr, t.math('MULTIPLY', d, 0.6), clamp=True)
    if edge is not None and wear > 0:
        brk = t.ranged(ro, 1.3, 0.5, "scratches break the chips up", rlo, rhi)
        exposed = t.ranged(ao, 0.0, 1.0, "only exposed edges", 0.72, 0.94, smooth=True)
        w = t.math('MULTIPLY', t.math('MULTIPLY', t.math('MULTIPLY', edge, brk), exposed), wear)
        w = t.ranged(w, 0.0, 1.0, "chipped edge", 0.30, 0.70, smooth=True)
        col = t.mixc(w, col, wear_col, "metal under the coat")
        rr = t.mixf(w, rr, 0.40)
        metal = t.mixf(w, metallic, 0.85)
        coat_w = t.mixf(w, coat_w, 0.0)
    t.put(b.inputs["Base Color"], col)
    t.put(b.inputs["Roughness"], rr)
    t.put(b.inputs["Metallic"], metal)
    t.put(b.inputs["Coat Weight"], coat_w)
    b.inputs["Coat Roughness"].default_value = 0.30
    b.inputs["Specular IOR Level"].default_value = 0.5
    nd = t.node('ShaderNodeNormalMap', "scan: relief", uv_map=space if space != 'OBJECT' else "UV_Tile")
    nd.inputs["Strength"].default_value = relief
    t.link(t.tex(tex_image(scan("nor_gl"), "Non-Color"), mp, "scan: normal").outputs["Color"], nd.inputs["Color"])
    n = nd.outputs["Normal"]
    if peel > 0:                              # the powder coat's fine orange peel
        bp = t.node('ShaderNodeBump', "orange peel")
        bp.inputs["Strength"].default_value = peel
        bp.inputs["Distance"].default_value = peel_dist
        t.link(t.noise(t.node('ShaderNodeTexCoord').outputs["Object"], peel_scale, detail=3.0, rough=0.55, label="peel"),
               bp.inputs["Height"])
        t.link(n, bp.inputs["Normal"])
        n = bp.outputs["Normal"]
    n = join_normals(t, n, bevel_normal(t, bevel))
    t.link(n, b.inputs["Normal"])
    mat.diffuse_color = tuple(tint) + (1.0,)
    mat.metallic = float(metallic)
    mat.roughness = 0.5
    mat["dra_surf"] = "%s: tint %s, scan blue_metal_plate (relief %.2f, tile %.2f m), coat %.2f, wear %.2f" % (
        label, tuple(round(c, 3) for c in tint), relief, tile, coat, wear)


def steel(mat, masks, bevel, tint=(0.80, 0.80, 0.81), rough=(0.18, 0.34), tile=0.5, grime=0.20, space="UV_Tile"):
    """bright brushed steel: the pack's StainedSteel set, fine scratches, fingerprints"""
    t = Tree(mat)
    b = t.node('ShaderNodeBsdfPrincipled')
    out = t.node('ShaderNodeOutputMaterial')
    t.link(b.outputs[0], out.inputs["Surface"])
    u = t.coords(space)
    mp = t.mapped(u, 1.0 / tile, "one repeat = %.2f m" % tile)
    dim = tex_image(pack("Base", "StainedSteel", "T_StainSteeled_Diff.png"))
    lo, hi = tex_range(dim)
    bw = t.node('ShaderNodeRGBToBW')
    t.link(t.tex(dim, mp, "pack: stained steel").outputs["Color"], bw.inputs[0])
    vm = t.node('ShaderNodeVectorMath', "tint x shade", operation='SCALE')
    vm.inputs[0].default_value = tint
    t.link(t.ranged(bw.outputs[0], 0.84, 1.06, "shade", lo, hi), vm.inputs["Scale"])
    col = vm.outputs["Vector"]
    rim = tex_image(pack("Base", "StainedSteel", "T_StainSteeled_Rough.png"), "Non-Color")
    rlo, rhi = tex_range(rim)
    rr = t.ranged(t.tex(rim, mp, "pack: roughness").outputs["Color"], rough[0], rough[1], "roughness range", rlo, rhi)
    dirt = t.tex(tex_image(pack("Base", "Glass", "T_Glass_DirtRough.png"), "Non-Color"), t.mapped(u, 1.0 / 0.9),
                 "fingerprints").outputs["Color"]
    rr = t.math('ADD', rr, t.math('MULTIPLY', dirt, 0.14), clamp=True)
    ao, edge, crev = masks_of(t, masks)
    if crev is not None:
        g = t.math('MULTIPLY', t.math('POWER', t.math('SUBTRACT', 1.0, crev), 1.4), grime, clamp=True)
        col = t.mixc(g, col, (0.03, 0.03, 0.03), "grime in crevices")
        rr = t.math('ADD', rr, t.math('MULTIPLY', g, 0.35), clamp=True)
    if ao is not None:
        col = t.mixc(t.ranged(ao, 0.30, 0.0, "baked occlusion", 0.0, 1.0), col, (0.02, 0.02, 0.02), "occlusion")
    t.put(b.inputs["Base Color"], col)
    t.put(b.inputs["Roughness"], rr)
    b.inputs["Metallic"].default_value = 1.0
    b.inputs["Anisotropic"].default_value = 0.3
    nd = t.node('ShaderNodeNormalMap', "fine scratches", uv_map=space if space != 'OBJECT' else "UV_Tile")
    nd.inputs["Strength"].default_value = 0.45
    t.link(t.tex(tex_image(pack("Base", "Glass", "T_Scratch_Norm.png"), "Non-Color"), mp, "pack: scratch normal").outputs["Color"],
           nd.inputs["Color"])
    t.link(join_normals(t, nd.outputs["Normal"], bevel_normal(t, bevel)), b.inputs["Normal"])
    mat.diffuse_color = tuple(tint) + (1.0,)
    mat.metallic = 1.0
    mat.roughness = 0.3
    mat["dra_surf"] = "bright steel: pack StainedSteel + scratch normal, tint %s" % (tuple(tint),)


def rubber(mat, masks, bevel, tint=(0.016, 0.016, 0.017), rough=(0.62, 0.84), tile=0.4, dust=0.05, space="UV_Tile"):
    """black rubber (cables, hoses, belt, boots, buffers): a dull sheen broken by the pack's scratched-metal roughness,
    a little dust on top"""
    t = Tree(mat)
    b = t.node('ShaderNodeBsdfPrincipled')
    out = t.node('ShaderNodeOutputMaterial')
    t.link(b.outputs[0], out.inputs["Surface"])
    u = t.coords(space)
    mp = t.mapped(u, 1.0 / tile)
    rim = tex_image(pack("Base", "ScratchedMetal", "T_ScratchedMetal_Rough.png"), "Non-Color")
    rlo, rhi = tex_range(rim)
    rr = t.ranged(t.tex(rim, mp, "pack: roughness").outputs["Color"], rough[1], rough[0], "roughness range", rlo, rhi)
    col = tuple(tint)
    upf, geo = t.up()
    if dust > 0:
        d = t.math('MULTIPLY', t.math('MULTIPLY', upf, t.ranged(t.noise(geo.outputs["Position"], 2.0), 0.0, 1.0, None, 0.4,
                                                                0.8, smooth=True)), dust, clamp=True)
        col = t.mixc(d, col, (0.30, 0.29, 0.27), "dust")
        rr = t.math('ADD', rr, t.math('MULTIPLY', d, 0.3), clamp=True)
    ao, edge, crev = masks_of(t, masks)
    if ao is not None:
        col = t.mixc(t.ranged(ao, 0.30, 0.0, None, 0.0, 1.0), col, (0.004, 0.004, 0.004), "occlusion")
    t.put(b.inputs["Base Color"], col)
    t.put(b.inputs["Roughness"], rr)
    b.inputs["Metallic"].default_value = 0.0
    b.inputs["Specular IOR Level"].default_value = 0.35
    bn = bevel_normal(t, bevel)
    if bn is not None:
        t.link(bn, b.inputs["Normal"])
    mat.diffuse_color = tuple(tint) + (1.0,)
    mat.metallic = 0.0
    mat.roughness = 0.75
    mat["dra_surf"] = "black rubber: pack ScratchedMetal roughness"


def glass(mat):
    t = Tree(mat)
    b = t.node('ShaderNodeBsdfPrincipled')
    out = t.node('ShaderNodeOutputMaterial')
    t.link(b.outputs[0], out.inputs["Surface"])
    dirt = t.tex(tex_image(pack("Base", "Glass", "T_Glass_DirtRough.png"), "Non-Color"),
                 t.mapped(t.coords("UV_Tile"), 1.0 / 0.6), "smudges").outputs["Color"]
    b.inputs["Base Color"].default_value = (0.010, 0.011, 0.013, 1)
    t.put(b.inputs["Roughness"], t.ranged(dirt, 0.10, 0.22))
    b.inputs["Coat Weight"].default_value = 0.6
    mat.diffuse_color = (0.02, 0.022, 0.026, 1.0)
    mat["dra_surf"] = "smoked lens glass"


def lamp(mat):
    t = Tree(mat)
    b = t.node('ShaderNodeBsdfPrincipled')
    out = t.node('ShaderNodeOutputMaterial')
    t.link(b.outputs[0], out.inputs["Surface"])
    b.inputs["Base Color"].default_value = (0.60, 0.66, 0.70, 1)
    b.inputs["Roughness"].default_value = 0.25
    b.inputs["Emission Color"].default_value = (0.80, 0.90, 1.0, 1)
    b.inputs["Emission Strength"].default_value = 0.8
    mat.diffuse_color = (0.85, 0.9, 1.0, 1.0)
    mat["dra_surf"] = "indicator: faint cool white (not a light source)"


# the recipes per material (tints linear). The whites: 0.66 linear (sRGB about 0.84) so the coat reads white, not
# blown out, under the studio and the pocket's lights
RECIPES = {
    # (stage 3 look pass: the first whites read clean and smooth like plastic beside ref2/ref3; now a grey-white with
    # the scan's scuffs as light and dark, a visible relief and orange peel, a fine grain, satin not glossy, a little
    # more shade in crevices and a few more chips on exposed edges, still a clean, kept machine)
    "M_DoorArm_PaintWhite": ("painted", dict(tint=(0.62, 0.62, 0.605), shade=(0.86, 1.04), rough=(0.46, 0.64), coat=0.10,
                                             relief=0.55, wear=0.24, wear_col=(0.36, 0.36, 0.37), grime=0.32, dust=0.06,
                                             peel=0.22, peel_scale=140.0, peel_dist=0.0011, grain=0.10,
                                             label="satin white powder coat")),       # fix round: coarser grain (ref1)
    "M_DoorArm_Orange": ("painted", dict(tint=(0.70, 0.15, 0.018), shade=(0.88, 1.05), rough=(0.45, 0.60), coat=0.10,
                                         relief=0.40, wear=0.20, grime=0.25, dust=0.03, peel=0.04, grain=0.05,
                                         label="safety orange")),
    "M_DoorArm_Gunmetal": ("painted", dict(tint=(0.075, 0.078, 0.085), shade=(0.78, 1.25), rough=(0.32, 0.52),
                                           coat=0.08, relief=0.50, wear=0.30, wear_col=(0.52, 0.52, 0.54), grime=0.25,
                                           dust=0.06, peel=0.0, metallic=0.75, tile=0.70, label="dark gunmetal")),
    "M_DoorArm_BlackMetal": ("painted", dict(tint=(0.028, 0.029, 0.032), shade=(0.76, 1.32), rough=(0.40, 0.62),
                                             coat=0.06, relief=0.55, wear=0.40, wear_col=(0.42, 0.42, 0.44), grime=0.20,
                                             dust=0.06, peel=0.0, metallic=0.55, tile=0.70,
                                             label="black metal (the carriage's stacked plates)")),
    "M_DoorArm_Steel": ("steel", dict()),
    "M_DoorArm_Rubber": ("rubber", dict()),
    "M_DoorArm_Glass": ("glass", None),
    "M_DoorArm_Lamp": ("lamp", None),
}


def dress(name, kind, kw, masks, bevel, space="UV_Tile"):
    m = bpy.data.materials.get(name) or bpy.data.materials.new(name)
    if kind == "painted":
        painted(m, masks, bevel, space=space, **kw)
    elif kind == "steel":
        steel(m, masks, bevel, space=space, **kw)
    elif kind == "rubber":
        rubber(m, masks, bevel, space=space, **kw)
    elif kind == "glass":
        glass(m)
    elif kind == "lamp":
        lamp(m)
    return m


def atlas_images():
    """the two baked maps: loaded from textures/ when baked, else blank placeholders the bake fills"""
    out = []
    for nm, rgba in ((MASKS, (1.0, 0.0, 1.0, 1.0)), (BEVEL, (0.5, 0.5, 1.0, 1.0))):
        path = os.path.join(TEX, nm + ".png")
        im = bpy.data.images.get(nm)
        if im is None and os.path.exists(path):
            im = bpy.data.images.load(path)
            im.name = nm
        if im is None:
            im = bpy.data.images.new(nm, ATLAS, ATLAS, alpha=False)
            im.pixels.foreach_set(np.tile(np.array(rgba, np.float32), ATLAS * ATLAS))
            im.filepath_raw = path
            im.file_format = 'PNG'
        im.colorspace_settings.name = 'Non-Color'
        out.append(im)
    return out


def dress_all():
    masks, bevel = atlas_images()
    done = {}
    for name, (kind, kw) in RECIPES.items():
        m = dress(name, kind, kw or {}, masks, bevel)
        done[name] = m.get("dra_surf")
    # the carriage's stacked plates in black metal (the other gunmetal parts stay gunmetal)
    me = bpy.data.meshes.get("DoorArm | carriage")
    if me is not None:
        for i, m in enumerate(me.materials):
            if m and m.name == "M_DoorArm_Gunmetal":
                me.materials[i] = bpy.data.materials["M_DoorArm_BlackMetal"]
    # the kit: the same recipes without the arm's atlas
    for src, dst in KIT_MATS.items():
        kind, kw = RECIPES[src]
        done[dst] = dress(dst, kind, kw or {}, None, None).get("dra_surf")
    for me in kit_meshes():
        for i, m in enumerate(me.materials):
            if m and m.name in KIT_MATS:
                me.materials[i] = bpy.data.materials[KIT_MATS[m.name]]
    # the cables and hoses (curves): rubber in object space
    cab = dress("M_DoorArm_Cable", "rubber", dict(tile=0.35), None, None, space="OBJECT")
    done["M_DoorArm_Cable"] = cab.get("dra_surf") + " (object space, for the curves)"
    for cu in bpy.data.curves:
        if cu.name.startswith(("DoorArm_", "Showroom | arm")):
            cu.materials.clear()
            cu.materials.append(cab)
    return done


# ------------------------------------------------------------------ the ship context's missing images
def stand_ins():
    """materials of the ship context whose image files are missing (an old absolute path in the ship): plain colours,
    so the reference shows in Rendered view (the context is not part of the asset)"""
    looks = {"PBR | Brushed titanium": ((0.50, 0.51, 0.53, 1), 1.0, 0.32),
             "PBR | Graphite lower hull": ((0.06, 0.065, 0.07, 1), 0.6, 0.45)}
    fixed = []
    ctx_mats = {s.material for o in bpy.data.objects if o.get("reference_only") for s in o.material_slots if s.material}
    for m in ctx_mats:
        if not m.node_tree or m.name.startswith("M_DoorArm"):
            continue
        missing = [n for n in m.node_tree.nodes if n.type == 'TEX_IMAGE' and n.image is not None and not n.image.has_data
                   and not os.path.exists(bpy.path.abspath(n.image.filepath))]
        if not missing:
            continue
        bsdf = next((n for n in m.node_tree.nodes if n.type == 'BSDF_PRINCIPLED'), None)
        for n in missing:
            for o in n.outputs:
                for lk in list(o.links):
                    m.node_tree.links.remove(lk)
        if bsdf is not None:
            col, met, rough = looks.get(m.name, ((0.35, 0.36, 0.38, 1), 0.5, 0.4))
            bsdf.inputs["Base Color"].default_value = col
            bsdf.inputs["Metallic"].default_value = met
            bsdf.inputs["Roughness"].default_value = rough
        m["stand_in"] = "the ship's image files for this material are missing; plain colour for the reference only"
        fixed.append(m.name)
    return fixed


# ------------------------------------------------------------------ the pocket's light rig, saved in the file
def mx(key, x):
    return x if key == "AFT" else 6.1 - x


def ship_view(sf):
    """fix round (CRITIQUE_2 nice 13): the solver empties hidden in the viewport by their collection's eye (they stay
    evaluated: the drivers read them), and a camera of the Ship fit scene framing the AFT arm and its door under the
    deck (version A, frame 1), set as the scene's camera; it lives in the studio collection, not part of the asset"""
    for vl in sf.view_layers:
        def walk(lc):
            if lc.collection.name == "DoorArm | solver (hidden)":
                lc.hide_viewport = True
            for ch in lc.children:
                walk(ch)
        walk(vl.layer_collection)
    wm = bpy.context.window_manager
    if wm.windows:
        wm.windows[0].scene = sf
    sf.frame_set(1)
    dg = bpy.context.evaluated_depsgraph_get()
    obs = [o for o in sf.objects if o.name.startswith("DoorArm_AFT_A | ") and o.type == 'MESH']
    obs.append(bpy.data.objects["Moonpool_Door_Aft"])
    pts = [o.evaluated_get(dg).matrix_world @ Vector(c) for o in obs for c in o.evaluated_get(dg).bound_box]
    lo = Vector([min(p[i] for p in pts) for i in range(3)])
    hi = Vector([max(p[i] for p in pts) for i in range(3)])
    ctr = (lo + hi) / 2
    name = "Cam | Ship fit (AFT arm and door)"
    cd = bpy.data.cameras.get(name) or bpy.data.cameras.new(name)
    cd.clip_start, cd.clip_end, cd.sensor_fit = 0.05, 60.0, 'HORIZONTAL'
    cam = bpy.data.objects.get(name) or bpy.data.objects.new(name, cd)
    c = bpy.data.collections[LIGHTS]
    if cam.name not in c.objects:
        c.objects.link(cam)
    cam["not_exported"] = True
    loc = Vector((ctr.x + 1.6, -1.6, -4.35))          # inboard of the pocket, just under the deck (as the rig renders)
    cam.location = loc
    cam.rotation_euler = (ctr - loc).normalized().to_track_quat('-Z', 'Y').to_euler()
    bpy.context.view_layer.update()
    M = cam.matrix_world.to_3x3()
    r, u, f = M.col[0], M.col[1], -M.col[2]
    mx_ = my_ = 1e-6
    for p in pts:
        v = p - loc
        d = v.dot(f)
        if d > 0.05:
            mx_, my_ = max(mx_, abs(v.dot(r)) / d), max(my_, abs(v.dot(u)) / d)
    sw = cd.sensor_width
    cd.lens = min(sw / 2 / mx_, sw * 9 / 16 / 2 / my_) / 1.08
    sf.camera = cam
    sf.render.resolution_x, sf.render.resolution_y = 1600, 900
    return cam.name


def light_rig(sf):
    """area lights under the deck along both pockets and a fill from the moon pool, in a collection that is not part of
    the asset (export skips it): the pocket under the deck is dark in the ship"""
    c = bpy.data.collections.get(LIGHTS) or bpy.data.collections.new(LIGHTS)
    if c.name not in [x.name for x in sf.collection.children]:
        sf.collection.children.link(c)
    c["not_exported"] = True

    def look(o, target):
        d = (Vector(target) - o.location).normalized()
        o.rotation_euler = d.to_track_quat('-Z', 'Y').to_euler()

    def area(name, loc, size, energy, target, color=(1.0, 0.97, 0.93)):
        ld = bpy.data.lights.get(name) or bpy.data.lights.new(name, 'AREA')
        ld.shape = 'RECTANGLE'
        ld.size, ld.size_y = size
        ld.energy = energy
        ld.color = color
        o = bpy.data.objects.get(name) or bpy.data.objects.new(name, ld)
        if o.name not in c.objects:
            c.objects.link(o)
        o.location = loc
        look(o, target)
        o["not_exported"] = True
        return o
    for key in ("AFT", "FORE"):
        x0 = mx(key, -5.3)
        area("Light | %s under deck inboard" % key, (x0, -4.55, -3.90), (7.5, 1.0), 750, (x0, -6.6, -6.2))
        area("Light | %s under deck outboard" % key, (x0, -8.60, -3.95), (7.5, 0.6), 300, (x0, -6.0, -6.0),
             color=(0.90, 0.95, 1.0))
        area("Light | %s fill from the well" % key, (mx(key, -2.1), -2.6, -5.6), (3.0, 3.0), 380, (mx(key, -5.0), -7.0, -6.0),
             color=(0.85, 0.92, 1.0))
        area("Light | %s fill over the moon pool" % key, (mx(key, 0.5), -1.5, -4.6), (3.0, 2.0), 280,
             (mx(key, -1.5), -7.0, -6.5), color=(0.85, 0.92, 1.0))
    w = bpy.data.worlds.get("DoorArm | pocket world") or bpy.data.worlds.new("DoorArm | pocket world")
    w.use_nodes = True
    bg = w.node_tree.nodes.get("Background")
    if bg:
        bg.inputs[0].default_value = (0.018, 0.021, 0.026, 1)
        bg.inputs[1].default_value = 1.0
    sf.world = w
    return c


# ------------------------------------------------------------------ the showroom's studio, kept in the file
HDRI = os.path.join(ROOT, "source_polyhaven", "hdris", "studio_small_09_2k.hdr")


def studio_look():
    """the showroom as the references are lit: a mid-grey backdrop, a soft studio HDRI (Poly Haven, CC0) under the key,
    fill and rim lights, the fill low for some contrast"""
    m = bpy.data.materials.get("M_Studio_Floor")
    if m and m.node_tree:
        b = m.node_tree.nodes.get("Principled BSDF")
        if b:
            b.inputs["Base Color"].default_value = (0.40, 0.40, 0.40, 1)
            b.inputs["Roughness"].default_value = 0.7
        m.diffuse_color = (0.40, 0.40, 0.40, 1)
    for nm, e in (("Showroom | key", 800.0), ("Showroom | fill", 150.0), ("Showroom | rim", 500.0)):
        ld = bpy.data.lights.get(nm)
        if ld:
            ld.energy = e
    w = bpy.data.worlds.get("Showroom | world")
    if w and os.path.exists(HDRI):
        # the HDRI lights and reflects; the camera sees a plain grey (past the backdrop's edge it would show the studio)
        w.use_nodes = True
        nt = w.node_tree
        nt.nodes.clear()
        out = nt.nodes.new('ShaderNodeOutputWorld')
        env = nt.nodes.new('ShaderNodeTexEnvironment')
        env.image = bpy.data.images.load(HDRI, check_existing=True)
        lit = nt.nodes.new('ShaderNodeBackground')
        lit.inputs["Strength"].default_value = 0.30
        nt.links.new(env.outputs["Color"], lit.inputs["Color"])
        plain = nt.nodes.new('ShaderNodeBackground')
        plain.inputs["Color"].default_value = (0.21, 0.21, 0.21, 1)
        plain.inputs["Strength"].default_value = 1.0
        lp = nt.nodes.new('ShaderNodeLightPath')
        mix = nt.nodes.new('ShaderNodeMixShader')
        nt.links.new(lp.outputs["Is Camera Ray"], mix.inputs[0])
        nt.links.new(lit.outputs[0], mix.inputs[1])
        nt.links.new(plain.outputs[0], mix.inputs[2])
        nt.links.new(mix.outputs[0], out.inputs["Surface"])
    return True


# ------------------------------------------------------------------ the bake (the docking arm's method)
def use_gpu(sc):
    sc.render.engine = 'CYCLES'
    try:
        prefs = bpy.context.preferences.addons['cycles'].preferences
        for kind in ('OPTIX', 'CUDA', 'HIP', 'ONEAPI'):
            try:
                prefs.compute_device_type = kind
            except Exception:
                continue
            (prefs.refresh_devices if hasattr(prefs, "refresh_devices") else prefs.get_devices)()
            if any(d.type == kind for d in prefs.devices):
                for d in prefs.devices:
                    d.use = d.type == kind
                sc.cycles.device = 'GPU'
                return kind
    except Exception as ex:
        print("[dra] GPU not set:", ex)
    sc.cycles.device = 'CPU'
    return 'CPU'


def bake(only=None):
    t0 = time.time()
    sc = bpy.context.scene
    kind = use_gpu(sc)
    masks, bevel = atlas_images()
    if only is None:                          # a fresh bake starts from the blank fill
        for im, rgba in ((masks, (1.0, 0.0, 1.0, 1.0)), (bevel, (0.5, 0.5, 1.0, 1.0))):
            if tuple(im.size) != (ATLAS, ATLAS):
                im.scale(ATLAS, ATLAS)
            im.pixels.foreach_set(np.tile(np.array(rgba, np.float32), ATLAS * ATLAS))
    mm = bpy.data.materials.new("_bake_masks")
    mm.use_nodes = True
    nt = mm.node_tree
    nt.nodes.clear()
    ao = nt.nodes.new('ShaderNodeAmbientOcclusion')
    ao.samples = 16
    ao.only_local = True
    ao.inputs["Distance"].default_value = 0.35
    cv = nt.nodes.new('ShaderNodeAmbientOcclusion')
    cv.samples = 16
    cv.only_local = True
    cv.inputs["Distance"].default_value = 0.035
    bv = nt.nodes.new('ShaderNodeBevel')
    bv.samples = 8
    bv.inputs["Radius"].default_value = EDGE_R
    geo = nt.nodes.new('ShaderNodeNewGeometry')
    dot = nt.nodes.new('ShaderNodeVectorMath')
    dot.operation = 'DOT_PRODUCT'
    nt.links.new(bv.outputs["Normal"], dot.inputs[0])
    nt.links.new(geo.outputs["Normal"], dot.inputs[1])
    inv = nt.nodes.new('ShaderNodeMath')
    inv.operation = 'SUBTRACT'
    inv.inputs[0].default_value = 1.0
    nt.links.new(dot.outputs["Value"], inv.inputs[1])
    edg = nt.nodes.new('ShaderNodeMath')
    edg.operation = 'MULTIPLY'
    edg.inputs[1].default_value = 3.5
    edg.use_clamp = True
    nt.links.new(inv.outputs[0], edg.inputs[0])
    cc = nt.nodes.new('ShaderNodeCombineColor')
    nt.links.new(ao.outputs["AO"], cc.inputs[0])
    nt.links.new(edg.outputs[0], cc.inputs[1])
    nt.links.new(cv.outputs["AO"], cc.inputs[2])
    em = nt.nodes.new('ShaderNodeEmission')
    nt.links.new(cc.outputs[0], em.inputs["Color"])
    om = nt.nodes.new('ShaderNodeOutputMaterial')
    nt.links.new(em.outputs[0], om.inputs["Surface"])
    ti = nt.nodes.new('ShaderNodeTexImage')
    ti.image = masks
    nt.nodes.active = ti
    mn = bpy.data.materials.new("_bake_bevel")
    mn.use_nodes = True
    nt2 = mn.node_tree
    b2 = nt2.nodes["Principled BSDF"]
    bv2 = nt2.nodes.new('ShaderNodeBevel')
    bv2.samples = 8
    bv2.inputs["Radius"].default_value = BEVEL_R
    nt2.links.new(bv2.outputs["Normal"], b2.inputs["Normal"])
    ti2 = nt2.nodes.new('ShaderNodeTexImage')
    ti2.image = bevel
    nt2.nodes.active = ti2
    sc.cycles.samples = 16
    sc.cycles.use_denoising = False
    sc.render.bake.margin = 6
    sc.render.bake.use_clear = False
    rows = []
    meshes = [me for me in arm_meshes() if only is None or me.name in only]
    for src in meshes:
        t1 = time.time()
        me = src.copy()
        bm = bmesh.new()
        bm.from_mesh(me)
        drop = decal_faces(src)
        if drop:
            bm.faces.ensure_lookup_table()
            bmesh.ops.delete(bm, geom=[bm.faces[i] for i in drop], context='FACES')
        for f in bm.faces:
            f.material_index = 0
        bm.to_mesh(me)
        bm.free()
        me.materials.clear()
        me.materials.append(mm)
        me.uv_layers.active = me.uv_layers["UVMap"]
        me.uv_layers["UVMap"].active_render = True
        tmp = bpy.data.objects.new("_bake_" + src.name, me)
        sc.collection.objects.link(tmp)
        for o2 in bpy.context.view_layer.objects:
            o2.select_set(False)
        tmp.select_set(True)
        bpy.context.view_layer.objects.active = tmp
        try:
            bpy.ops.object.bake(type='EMIT', margin=6, use_clear=False)
            me.materials.clear()
            me.materials.append(mn)
            bpy.ops.object.bake(type='NORMAL', normal_space='TANGENT', margin=6, use_clear=False)
        finally:
            bpy.data.objects.remove(tmp)
            bpy.data.meshes.remove(me)
        rows.append([src.name, round(time.time() - t1, 1)])
        print("[dra] baked", src.name, rows[-1][1], "s", flush=True)
    os.makedirs(TEX, exist_ok=True)
    for im, nm in ((masks, MASKS), (bevel, BEVEL)):
        im.filepath_raw = os.path.join(TEX, nm + ".png")
        im.file_format = 'PNG'
        im.save()
        try:
            im.filepath = bpy.path.relpath(os.path.join(TEX, nm + ".png")) if bpy.data.filepath else im.filepath_raw
        except ValueError:                    # a file on another drive (a scratch copy): keep the absolute path
            im.filepath = os.path.join(TEX, nm + ".png")
    bpy.data.materials.remove(mm)
    bpy.data.materials.remove(mn)
    return dict(device=kind, seconds=round(time.time() - t0, 1), size=ATLAS, meshes=rows)


def surface_table():
    """what each material is made of (for the README and whoever rebuilds the materials in Unreal)"""
    rows = []
    for name, (kind, kw) in RECIPES.items():
        r = dict(material=name, kind=kind)
        if kw:
            r.update({k: (list(v) if isinstance(v, tuple) else v) for k, v in kw.items()})
        rows.append(r)
    return rows