"""Surfaces of the robot cell, headless, in place:

    blender.exe -b ../RobotCell.blend --python rc_surfaces.py [-- --nobake]

The canister station's recipe (models/canister-station/build/cs_surfaces.py), for the user's rule "realistic with good
textures" from the first version:
  * tiling sets from the user's own Season 2 pack (already in the Unreal project), read in place by relative path from
    models/vessel-interior/textures/ModSci_Engineer/Textures: ScratchedMetal under the paints, BrushedMetal for bare
    metal and copper, Glass_DirtRough as a smudge mask, T_Scratch_Norm for fine scratches;
  * the floor's diamond plate from Poly Haven's metal_plate scan (CC0, source_polyhaven/): its relief and its lugs, my
    own colour (the scan is green and rusty, the cell is clean);
  * mask maps baked in Cycles (R occlusion, G edges, B crevices), one per group of meshes:
        T_RC_Body_Masks    4096  SM_RobotCell_Body
        T_RC_Robot_Masks   4096  SK_RobotCell_Robot (bind pose)
        T_RC_Moving_Masks  2048  the two leaves, the side door, the turntable plate, the e-stop (one atlas, each mesh
                                 in its own square of it)
    A material used by several groups has a copy per group ("<name>_Robot", "<name>_Moving") reading that group's map;
  * drawn textures (rc_textures.py): the decal atlas, the screen.
Tints are the picture's colours (sampled from reference/ref_user.jpg, linear): the frame's charcoal about 0.035, the
back wall 0.028, the robot's orange about 0.78 0.12 0.006, copper, the turntable's bright steel.
UV layers (order matters for Unreal):
    UVMap   (UV0) a unique unwrap into the group's mask map; decals and the screen take their textures' 0..1
    UV_Tile (UV1) metres, box-projected, for the tiling sets
"""
import bpy, bmesh, math, os, sys, time, json
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
FOLDER = os.path.dirname(HERE)
TEX = os.path.join(FOLDER, "textures")
PACK = os.path.normpath(os.path.join(FOLDER, "..", "vessel-interior", "textures", "ModSci_Engineer", "Textures"))
PH = os.path.join(FOLDER, "source_polyhaven", "metal_plate")
SCR = ("Base/ScratchedMetal/T_ScratchedMetal_Diff.png", "Base/ScratchedMetal/T_ScratchedMetal_Rough.png",
       "Base/ScratchedMetal/T_ScratchedMetal_Normal.png")
BRU = ("Base/BrushedMetal/T_BrushedMetal_Diff.png", "Base/BrushedMetal/T_BrushedMetal_Rough.png",
       "Base/BrushedMetal/T_BrushedMetal_Normal.png")
DIRT = "Base/Glass/T_Glass_DirtRough.png"
SCRATCH_N = "Base/Glass/T_Scratch_Norm.png"
EDGE_R = 0.004

GROUPS = {
    "body": dict(objects=["SM_RobotCell_Body"], image="T_RC_Body_Masks", size=4096, suffix=""),
    "robot": dict(objects=["SK_RobotCell_Robot"], image="T_RC_Robot_Masks", size=4096, suffix="_Robot"),
    "moving": dict(objects=["SM_RobotCell_DoorLeft", "SM_RobotCell_DoorRight", "SM_RobotCell_DoorSide",
                            "SM_RobotCell_TurntablePlate", "SM_RobotCell_EStop"], image="T_RC_Moving_Masks", size=2048,
                   suffix="_Moving",
                   squares={"SM_RobotCell_DoorLeft": (0.0, 0.5, 0.5), "SM_RobotCell_DoorRight": (0.5, 0.5, 0.5),
                            "SM_RobotCell_DoorSide": (0.0, 0.0, 0.5), "SM_RobotCell_TurntablePlate": (0.5, 0.125, 0.375),
                            "SM_RobotCell_EStop": (0.875, 0.0, 0.125)}),
}

PAINT = dict(set=SCR, metal=0.0, tile=0.60, shade=(0.96, 1.03), rough=(0.48, 0.36), normal=0.14, dirt=0.04, dirt_tile=1.4,
             grime=0.18, coat=0.0)
METAL = dict(set=BRU, metal=1.0, tile=0.12, shade=(0.95, 1.04), rough=(0.24, 0.34), normal=0.25, normal_map=SCRATCH_N,
             dirt=0.03, dirt_tile=0.6, grime=0.12, coat=0.0)
SURF = {
    "M_RC_Frame":  dict(PAINT, tint=(0.034, 0.033, 0.034), rough=(0.58, 0.46)),           # charcoal powder coat
    "M_RC_Panel":  dict(PAINT, tint=(0.026, 0.028, 0.031), rough=(0.50, 0.40)),           # back wall, ceiling
    "M_RC_Black":  dict(PAINT, tint=(0.010, 0.010, 0.011), rough=(0.52, 0.40)),
    "M_RC_Orange": dict(PAINT, tint=(0.66, 0.095, 0.005), rough=(0.34, 0.26), shade=(0.97, 1.02), grime=0.12, coat=0.35),
    "M_RC_Copper": dict(METAL, tint=(0.80, 0.40, 0.24), rough=(0.26, 0.36)),
    "M_RC_Chrome": dict(METAL, tint=(0.82, 0.82, 0.83), rough=(0.10, 0.18), tile=0.10),
    "M_RC_Steel":  dict(METAL, tint=(0.62, 0.62, 0.64), rough=(0.28, 0.40), tile=0.35, normal=0.35),
    "M_RC_Rubber": dict(PAINT, tint=(0.008, 0.008, 0.009), rough=(0.72, 0.62), normal=0.05, grime=0.0),
    "M_RC_Gap":    dict(set=(None, SCR[1], None), tint=(0.004, 0.004, 0.005), metal=0.0, tile=0.40, shade=(1.0, 1.0),
                        rough=(0.70, 0.55), normal=0.0, dirt=0.0, dirt_tile=1.0, grime=0.0, coat=0.0),
}
PLASTIC = {"M_RC_Red": ((0.55, 0.02, 0.005), 0.22), "M_RC_Yellow": ((0.75, 0.42, 0.01), 0.38)}
EMIT = {"M_RC_LED": ((1.0, 0.97, 0.93), 14.0), "M_RC_Lamp": ((0.92, 0.95, 1.0), 2.5)}
SPECIAL = {"M_RC_Screen", "M_RC_Decal"}                      # keep their own 0..1 UVs
NOATLAS = {"M_RC_Glass", "M_RC_LED", "M_RC_Lamp"}            # no mask map: UVMap = their metres
NOBAKE = SPECIAL | NOATLAS


# ------------------------------------------------------------------ UV layers
def uv_tile_layer(me):
    """metres, by the face's dominant axis"""
    uv = me.uv_layers.get("UV_Tile") or me.uv_layers.new(name="UV_Tile")
    co = np.zeros(len(me.vertices) * 3, np.float32)
    me.vertices.foreach_get("co", co)
    co = co.reshape(-1, 3)
    out = np.zeros((len(me.loops), 2), np.float32)
    for p in me.polygons:
        n = p.normal
        ax = 0 if abs(n.x) >= abs(n.y) and abs(n.x) >= abs(n.z) else (1 if abs(n.y) >= abs(n.z) else 2)
        for li in range(p.loop_start, p.loop_start + p.loop_total):
            c = co[me.loops[li].vertex_index]
            out[li] = (c[1], c[2]) if ax == 0 else ((c[0], c[2]) if ax == 1 else (c[0], c[1]))
    uv.data.foreach_set("uv", out.ravel())


def uv_atlas_layer(ob, size, square=(0.0, 0.0, 1.0)):
    """UVMap: a unique unwrap of the mesh's baked faces packed into `square` (u0, v0, side) of its group's map;
    decals and the screen keep their own 0..1 (from the build's UV_Special), glass, LED and lamps their metres"""
    me = ob.data
    keep = {}
    for l in me.uv_layers:
        a = np.zeros(len(me.loops) * 2, np.float32)
        l.data.foreach_get("uv", a)
        keep[l.name] = a
    while len(me.uv_layers):
        me.uv_layers.remove(me.uv_layers[0])
    for nm in ("UVMap", "UV_Tile"):
        me.uv_layers.new(name=nm).data.foreach_set("uv", keep.get(nm, np.zeros(len(me.loops) * 2, np.float32)))
    me.uv_layers.active = me.uv_layers["UVMap"]
    me.uv_layers["UVMap"].active_render = True
    names = [m.name for m in me.materials]
    special = {i for i, n in enumerate(names) if n in NOBAKE}
    vl = bpy.context.view_layer
    for o in list(vl.objects):
        if o is not None:
            o.select_set(False)
    ob.select_set(True)
    vl.objects.active = ob
    for p in me.polygons:
        p.select = p.material_index not in special
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_mode(type='FACE')
    margin = 6.0 / (size * square[2])
    bpy.ops.uv.smart_project(angle_limit=math.radians(50), island_margin=margin, area_weight=0.0, correct_aspect=True,
                             scale_to_bounds=False)
    bpy.ops.uv.select_all(action='SELECT')
    try:
        bpy.ops.uv.pack_islands(rotate=True, scale=True, margin_method='FRACTION', margin=margin, shape_method='CONCAVE')
    except Exception as ex:
        print("pack skipped", ob.name, ex)
    bpy.ops.object.mode_set(mode='OBJECT')
    a = np.zeros(len(me.loops) * 2, np.float32)
    me.uv_layers["UVMap"].data.foreach_get("uv", a)
    a = a.reshape(-1, 2)
    pad = 4.0 / (size * square[2])
    a = pad + np.clip(a, 0, 1) * (1 - 2 * pad)
    a = np.array([square[0], square[1]], np.float32) + a * square[2]
    sp = keep.get("UV_Special")             # first run: the build's layer; a re-run keeps what UVMap already holds
    sp = (sp if sp is not None else keep["UVMap"]).reshape(-1, 2)
    tile = keep["UV_Tile"].reshape(-1, 2)
    for p in me.polygons:
        n = names[p.material_index]
        if n not in NOBAKE:
            continue
        for li in range(p.loop_start, p.loop_start + p.loop_total):
            a[li] = sp[li] if n in SPECIAL else tile[li]
    me.uv_layers["UVMap"].data.foreach_set("uv", a.ravel())


# ------------------------------------------------------------------ bake
def use_gpu():
    sc = bpy.context.scene
    sc.render.engine = 'CYCLES'
    try:
        prefs = bpy.context.preferences.addons['cycles'].preferences
        prefs.compute_device_type = 'HIP'
        prefs.get_devices()
        for d in prefs.devices:
            d.use = d.type == 'HIP'
        sc.cycles.device = 'GPU'
        return 'HIP'
    except Exception as ex:
        print("GPU not set:", ex)
        sc.cycles.device = 'CPU'
        return 'CPU'


def mask_material():
    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 = 24; ao.only_local = True; ao.inputs["Distance"].default_value = 0.20
    cv = nt.nodes.new('ShaderNodeAmbientOcclusion'); cv.samples = 16; cv.only_local = True; cv.inputs["Distance"].default_value = 0.018
    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"])
    return mm


def bake_group(g):
    """an emit bake of each mesh of the group alone (occlusion counts only that mesh) into the group's map"""
    t0 = time.time()
    kind = use_gpu()
    sc = bpy.context.scene
    name, N = g["image"], g["size"]
    for im in [i for i in bpy.data.images if i.name.startswith(name)]:
        bpy.data.images.remove(im)
    masks = bpy.data.images.new(name, N, N, alpha=False)
    masks.colorspace_settings.name = 'Non-Color'
    masks.pixels.foreach_set(np.tile(np.array((1.0, 0.0, 1.0, 1.0), np.float32), N * N))
    mm = mask_material()
    ti = mm.node_tree.nodes.new('ShaderNodeTexImage'); ti.image = masks
    mm.node_tree.nodes.active = ti
    sc.cycles.samples = 24
    sc.cycles.use_denoising = False
    sc.render.bake.margin = 6
    sc.render.bake.use_clear = False
    for on in g["objects"]:
        ob = bpy.data.objects[on]
        me = ob.data.copy()                      # rest geometry: the copy carries no armature modifier
        bm = bmesh.new()
        bm.from_mesh(me)
        names = [m.name for m in me.materials]
        bmesh.ops.delete(bm, geom=[f for f in bm.faces if names[f.material_index] in NOBAKE], 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_" + on, me)
        sc.collection.objects.link(tmp)
        tmp.matrix_world = ob.matrix_world
        for o2 in list(bpy.context.view_layer.objects):
            if o2 is not None:
                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)
        finally:
            bpy.data.objects.remove(tmp)
            bpy.data.meshes.remove(me)
    os.makedirs(TEX, exist_ok=True)
    masks.filepath_raw = os.path.join(TEX, name + ".png")
    masks.file_format = 'PNG'
    masks.save()
    masks.filepath = bpy.path.relpath(masks.filepath_raw)
    bpy.data.materials.remove(mm)
    return masks, dict(device=kind, seconds=round(time.time() - t0, 1), size=N)


# ------------------------------------------------------------------ materials (node recipe from the docking arm)
def tex_image(path, colorspace="sRGB"):
    path = os.path.normpath(path)
    for im in bpy.data.images:
        if im.filepath and os.path.normpath(bpy.path.abspath(im.filepath)) == path:
            return im
    im = bpy.data.images.load(path)
    im.colorspace_settings.name = colorspace
    try:
        im.filepath = bpy.path.relpath(path)
    except Exception:
        pass
    return im


_RANGE = {}
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 = (-2000 + 240 * (i // 4), 500 - 320 * (i % 4))
        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 uv(self, name):
        return self.node('ShaderNodeUVMap', uv_map=name).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 bsdf_out(t):
    b = t.node('ShaderNodeBsdfPrincipled')
    out = t.node('ShaderNodeOutputMaterial')
    t.link(b.outputs[0], out.inputs["Surface"])
    return b


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


def dress(mat, r, masks=None):
    """a coated or bare metal surface: a tiling set + smudges; with the baked masks also shade in crevices"""
    t = Tree(mat)
    b = bsdf_out(t)
    u0 = t.uv("UV_Tile")
    mp = t.mapped(u0, 1.0 / r["tile"], "one repeat = %.2f m" % r["tile"])
    diff, rough, normal = r["set"]
    dirt = t.tex(tex_image(os.path.join(PACK, DIRT), "Non-Color"), t.mapped(u0, 1.0 / r["dirt_tile"]), "smudges and dust").outputs["Color"]
    ao = edge = crev = None
    if masks is not None:
        ao, edge, crev = masks_of(t, masks)
    tint = tuple(r["tint"])
    col = tint
    if diff:
        im = tex_image(os.path.join(PACK, diff))
        lo, hi = tex_range(im)
        bw = t.node('ShaderNodeRGBToBW')
        t.link(t.tex(im, mp, "surface, diffuse").outputs["Color"], bw.inputs[0])
        sh = t.ranged(bw.outputs[0], r["shade"][0], r["shade"][1], "shade the tint", lo, hi)
        vm = t.node('ShaderNodeVectorMath', "tint x shade", operation='SCALE')
        vm.inputs[0].default_value = tint
        t.link(sh, vm.inputs["Scale"])
        col = vm.outputs["Vector"]
    grime = None
    if crev is not None and r["grime"] > 0:
        grime = t.math('MULTIPLY', t.math('POWER', t.math('SUBTRACT', 1.0, crev), 1.4), r["grime"], clamp=True, label="shade in crevices")
        col = t.mixc(grime, col, (0.02, 0.019, 0.018), "darken crevices")
    if ao is not None:
        col = t.mixc(t.ranged(ao, 0.25, 0.0, "occlusion", 0.0, 1.0), col, (0.004, 0.005, 0.006), "occlusion")
    rim = tex_image(os.path.join(PACK, rough), "Non-Color")
    rlo, rhi = tex_range(rim)
    ro = t.tex(rim, mp, "surface, roughness").outputs["Color"]
    rr = t.ranged(ro, r["rough"][0], r["rough"][1], "roughness range", rlo, rhi)
    rr = t.math('ADD', rr, t.math('MULTIPLY', dirt, r["dirt"]), clamp=True, label="smudges dull the sheen")
    if grime is not None:
        rr = t.math('ADD', rr, t.math('MULTIPLY', grime, 0.30), clamp=True, label="crevices are matte")
    t.put(b.inputs["Base Color"], col)
    t.put(b.inputs["Roughness"], rr)
    t.put(b.inputs["Metallic"], r["metal"])
    if r.get("coat", 0.0) > 0:
        b.inputs["Coat Weight"].default_value = r["coat"]
        b.inputs["Coat Roughness"].default_value = 0.08
    nm_file = r.get("normal_map") or normal
    if nm_file and r["normal"] > 0:
        nd = t.node('ShaderNodeNormalMap', "fine relief", uv_map="UV_Tile")
        nd.inputs["Strength"].default_value = r["normal"]
        t.link(t.tex(tex_image(os.path.join(PACK, nm_file), "Non-Color"), mp, "surface, normal").outputs["Color"], nd.inputs["Color"])
        t.link(nd.outputs["Normal"], b.inputs["Normal"])
    mat.diffuse_color = tuple(r["tint"]) + (1.0,)
    mat.metallic = r["metal"]


def dress_tread(mat, masks):
    """diamond plate: the scan's relief and lugs (Poly Haven metal_plate), a clean dark steel of my own; the lugs'
    tops worn brighter and smoother"""
    t = Tree(mat)
    b = bsdf_out(t)
    u0 = t.uv("UV_Tile")
    mp = t.mapped(u0, 1.0 / 1.10, "one repeat = 1.10 m (lugs about 4 cm apart, as the picture's)")
    disp = t.tex(tex_image(os.path.join(PH, "metal_plate_disp_2k.jpg"), "Non-Color"), mp, "scan: height").outputs["Color"]
    lug = t.ranged(disp, 0.0, 1.0, "lugs (the height's top)", 0.55, 0.75, smooth=True)
    ao, edge, crev = masks_of(t, masks)
    dirt = t.tex(tex_image(os.path.join(PACK, DIRT), "Non-Color"), t.mapped(u0, 1.0 / 1.6), "smudges").outputs["Color"]
    col = t.mixc(lug, (0.050, 0.056, 0.066), (0.26, 0.27, 0.29), "flats / worn lug tops")
    col = t.mixc(t.ranged(ao, 0.3, 0.0, "occlusion", 0.0, 1.0), col, (0.01, 0.01, 0.012), "occlusion")
    ro = t.tex(tex_image(os.path.join(PH, "metal_plate_rough_2k.jpg"), "Non-Color"), mp, "scan: roughness").outputs["Color"]
    rr = t.mixf(lug, t.ranged(ro, 0.48, 0.62, "flats' roughness"), 0.17, "lugs polished by feet")
    rr = t.math('ADD', rr, t.math('MULTIPLY', dirt, 0.06), clamp=True, label="smudges")
    t.put(b.inputs["Base Color"], col)
    t.put(b.inputs["Roughness"], rr)
    b.inputs["Metallic"].default_value = 0.85
    nd = t.node('ShaderNodeNormalMap', "scan: relief", uv_map="UV_Tile")
    nd.inputs["Strength"].default_value = 1.6
    t.link(t.tex(tex_image(os.path.join(PH, "metal_plate_nor_gl_2k.jpg"), "Non-Color"), mp, "scan: normal (GL)").outputs["Color"],
           nd.inputs["Color"])
    t.link(nd.outputs["Normal"], b.inputs["Normal"])
    mat.diffuse_color = (0.11, 0.115, 0.125, 1.0)
    mat.metallic = 0.85


def dress_glass(mat):
    t = Tree(mat)
    b = bsdf_out(t)
    b.inputs["Base Color"].default_value = (0.64, 0.83, 0.95, 1.0)
    b.inputs["Transmission Weight"].default_value = 1.0
    b.inputs["Roughness"].default_value = 0.02
    b.inputs["IOR"].default_value = 1.52
    dirt = t.tex(tex_image(os.path.join(PACK, DIRT), "Non-Color"), t.mapped(t.uv("UV_Tile"), 1.0 / 1.2), "smudges").outputs["Color"]
    t.put(b.inputs["Roughness"], t.ranged(dirt, 0.01, 0.07, "smudges on the glass"))
    mat.diffuse_color = (0.55, 0.72, 0.80, 0.25)
    try:
        mat.surface_render_method = 'BLENDED'
    except Exception:
        pass


def dress_plastic(mat, tint, rough):
    t = Tree(mat)
    b = bsdf_out(t)
    b.inputs["Base Color"].default_value = tuple(tint) + (1.0,)
    b.inputs["Roughness"].default_value = rough
    b.inputs["Coat Weight"].default_value = 0.2
    mat.diffuse_color = tuple(tint) + (1.0,)


def dress_emit(mat, col, strength):
    t = Tree(mat)
    b = bsdf_out(t)
    b.inputs["Base Color"].default_value = tuple(col) + (1.0,)
    b.inputs["Emission Color"].default_value = tuple(col) + (1.0,)
    b.inputs["Emission Strength"].default_value = strength
    b.inputs["Roughness"].default_value = 0.3
    mat.diffuse_color = tuple(col) + (1.0,)


def dress_decal(mat, img):
    t = Tree(mat)
    b = bsdf_out(t)
    tx = t.tex(img, t.uv("UVMap"), "decal atlas (rc_textures.py)")
    t.link(tx.outputs["Color"], b.inputs["Base Color"])
    t.link(tx.outputs["Alpha"], b.inputs["Alpha"])
    b.inputs["Roughness"].default_value = 0.40
    try:
        mat.blend_method = 'CLIP'
    except Exception:
        pass
    mat.diffuse_color = (0.9, 0.9, 0.9, 1.0)


def dress_screen(mat, img):
    t = Tree(mat)
    b = bsdf_out(t)
    tx = t.tex(img, t.uv("UVMap"), "status screen (rc_textures.py)")
    b.inputs["Base Color"].default_value = (0.004, 0.006, 0.007, 1)
    b.inputs["Roughness"].default_value = 0.12
    t.link(tx.outputs["Color"], b.inputs["Emission Color"])
    b.inputs["Emission Strength"].default_value = 1.4
    mat.diffuse_color = (0.05, 0.15, 0.25, 1)


def main(bake=True):
    rep = {}
    for gname, g in GROUPS.items():
        for n in g["objects"]:
            ob = bpy.data.objects[n]
            uv_tile_layer(ob.data)
            uv_atlas_layer(ob, g["size"], g.get("squares", {}).get(n, (0.0, 0.0, 1.0)))
            if "UV_Special" in ob.data.uv_layers:
                ob.data.uv_layers.remove(ob.data.uv_layers["UV_Special"])
    maps = {}
    for gname, g in GROUPS.items():
        if bake:
            maps[gname], rep[g["image"]] = bake_group(g)
        else:
            maps[gname] = bpy.data.images.get(g["image"]) or tex_image(os.path.join(TEX, g["image"] + ".png"), "Non-Color")
    # each group gets its own copy of a masked material ("<name><suffix>"), reading its own map
    done = set()
    for gname, g in GROUPS.items():
        for n in g["objects"]:
            me = bpy.data.objects[n].data
            for i, m in enumerate(me.materials):
                if m is None:
                    continue
                base = m.name
                for gg in GROUPS.values():
                    if gg["suffix"] and base.endswith(gg["suffix"]):
                        base = base[:-len(gg["suffix"])]
                if base not in SURF and base != "M_RC_Tread":
                    continue
                want = base + g["suffix"]
                cp = bpy.data.materials.get(want)
                if cp is None:
                    cp = bpy.data.materials[base].copy()
                    cp.name = want
                me.materials[i] = cp
                if want not in done:
                    if base == "M_RC_Tread":
                        dress_tread(cp, maps[gname])
                    else:
                        dress(cp, SURF[base], maps[gname])
                    done.add(want)
    for nm, (col, s) in EMIT.items():
        if nm in bpy.data.materials:
            dress_emit(bpy.data.materials[nm], col, s)
    for nm, (col, r) in PLASTIC.items():
        if nm in bpy.data.materials:
            dress_plastic(bpy.data.materials[nm], col, r)
    dress_glass(bpy.data.materials["M_RC_Glass"])
    dress_decal(bpy.data.materials["M_RC_Decal"], tex_image(os.path.join(TEX, "T_RC_Decals.png")))
    dress_screen(bpy.data.materials["M_RC_Screen"], tex_image(os.path.join(TEX, "T_RC_Screen_UI.png")))
    for m in list(bpy.data.materials):                      # materials left without users (the plain ones copied away)
        if m.users == 0:
            bpy.data.materials.remove(m)
    for im in bpy.data.images:
        if im.filepath and not im.filepath.startswith("//"):
            try:
                im.filepath = bpy.path.relpath(bpy.path.abspath(im.filepath))
            except Exception:
                pass
    rep["materials"] = sorted(m.name for m in bpy.data.materials)
    bpy.ops.wm.save_mainfile()
    print("SURFACES OK", json.dumps(rep))


if __name__ == "__main__":
    main(bake="--nobake" not in sys.argv)