"""Surfaces for the Mako, the Spearfish's and the Manta's method (models/tender-spearfish/build/surfaces_spearfish.py), run on
Tender_Mako.blend after the build: blender -b Tender_Mako.blend --python build/surfaces_mako.py
(generated by build/make_surfaces.py from this head + the Spearfish's helper functions and run, with this head's screen_uvs)

  * hard surfaces take tiling sets from the user's own Season 2 packs (already in the Unreal project), read in place from
    models/vessel-interior/textures: ScratchedMetal under the paints, StainedSteel for bare metal, FloorPanel_A on the cockpit
    deck, Glass_DirtRough as a smudge mask, T_Scratch_Norm for fine scratches. The seats take the CC0 leather scan the captain's
    quarters brought in (models/vessel-interior/source_textures/leather_white). Nothing is downloaded.
  * two maps baked from the model in Cycles per atlas: Masks (R occlusion, G edges, B crevices and seams) and BevelNormal (hard
    edges rounded 12 mm, the shell's panel seams as grooves). Exterior atlas 4096 px, cabin atlas 2048 px.
  * the colours follow the user's sheet: an orange paint under a clear coat on the shell, hood, collar, jaw and fins, a dark
    slate lower hull and side pods, mid grey big fins and ventral fin, light blue-tinted glass, graphite fittings, yellow marks
    on the docking clips. Worn as a working machine is: chips at exposed edges down to dark primer, grime in crevices, smudges
    in the clear coat; no rust.

Two UV layers on every mesh (the order matters for Unreal, which builds its tangents from the first):
    UVMap    (UV0)  a unique unwrap in the part's square of its atlas (the baked maps); the screens map onto T_Mako_Screens.
    UV_Tile  (UV1)  metres, box-projected (the tiling sets).
"""
import bpy, bmesh, math, os, time, json
import numpy as np
from mathutils import Vector

OUT = os.path.dirname(bpy.data.filepath)
HERE = os.path.join(OUT, "build")
TEX = os.path.join(OUT, "textures")
os.makedirs(TEX, exist_ok=True)
VI = os.path.normpath(os.path.join(OUT, "..", "vessel-interior"))
PACK = os.path.join(VI, "textures", "ModSci_Engineer", "Textures")
SCR = ("Base/ScratchedMetal/T_ScratchedMetal_Diff.png", "Base/ScratchedMetal/T_ScratchedMetal_Rough.png", "Base/ScratchedMetal/T_ScratchedMetal_Normal.png")
STN = ("Base/StainedSteel/T_StainSteeled_Diff.png", "Base/StainedSteel/T_StainSteeled_Rough.png", None)
FLOOR = ("FloorPanel_A/T_FloorPanel_A_diff.png", "FloorPanel_A/T_FloorPanel_A_Rough.png", "FloorPanel_A/T_FloorPanel_A_normals.png")
LEATHER = tuple(os.path.join(VI, "source_textures", "leather_white", f) for f in
                ("leather_white_diff_2k.jpg", "leather_white_rough_2k.jpg", "leather_white_nor_gl_2k.jpg"))
DIRT = "Base/Glass/T_Glass_DirtRough.png"
SCRATCH_N = "Base/Glass/T_Scratch_Norm.png"
REPORT = {}
t0 = time.time()
# a re-run replaces the maps of the run before: drop their images first, or the new ones come out named ".001"
for _im in list(bpy.data.images):
    if _im.name.startswith("T_Mako_"):
        bpy.data.images.remove(_im)
for _f in os.listdir(TEX):
    if _f.startswith("T_Mako_") and _f.count(".") > 1:
        os.remove(os.path.join(TEX, _f))

ATLASES = {
    "Exterior": dict(size=4096, rects={"SM_Mako_Hull": (0.0, 0.0, 0.86), "SM_Mako_Glass": (0.86, 0.0, 0.14),
                                        "SM_Mako_Hatch": (0.86, 0.14, 0.14), "SM_Mako_Arm_L": (0.86, 0.28, 0.07),
                                        "SM_Mako_Arm_R": (0.93, 0.28, 0.07), "SM_Mako_Rotor": (0.86, 0.35, 0.07)}),
    "Cabin": dict(size=2048, rects={"SM_Mako_Cabin": (0.0, 0.0, 1.0)}),
}
BEVEL_R, EDGE_R = 0.012, 0.020
SEAM_W0, SEAM_W1, SEAM_DEPTH = 0.0120, 0.0040, 0.0050      # v04 (the Critique): bolder, they must read at the sheet's scale

SURF = {
    "MAKO | Hull paint":      dict(set=SCR, tint=(0.97, 0.27, 0.008), metal=0.0, tile=1.20, shade=(0.96, 1.03), rough=(0.34, 0.26), normal=0.04,
                                   dirt=0.02, dirt_tile=2.4, grime=0.16, wear=0.14, wear_col=(0.15, 0.155, 0.16), wear_rough=0.45,
                                   coat=1.0, coat_rough=0.05, coat_dirt=0.04),
    "MAKO | Jaw paint":       dict(set=SCR, tint=(0.80, 0.075, 0.008), metal=0.0, tile=1.20, shade=(0.96, 1.03), rough=(0.36, 0.28), normal=0.04,
                                   dirt=0.02, dirt_tile=2.4, grime=0.08, wear=0.06, wear_col=(0.15, 0.155, 0.16), wear_rough=0.45,
                                   coat=0.45, coat_rough=0.06, coat_dirt=0.04),
    "MAKO | Lower hull":      dict(set=(STN[0], SCR[1], SCR[2]), tint=(0.085, 0.105, 0.13), metal=0.0, tile=0.90, shade=(0.92, 1.08), rough=(0.50, 0.38), normal=0.12,
                                   dirt=0.05, dirt_tile=1.9, grime=0.15, wear=0.30, wear_col=(0.26, 0.27, 0.28), wear_rough=0.36,
                                   coat=0.35, coat_rough=0.10, coat_dirt=0.06),
    "MAKO | Graphite":        dict(set=(STN[0], SCR[1], SCR[2]), tint=(0.028, 0.031, 0.036), metal=0.0, tile=0.70, shade=(0.92, 1.08), rough=(0.55, 0.42), normal=0.15,
                                   dirt=0.06, dirt_tile=1.9, grime=0.15, wear=0.40, wear_col=(0.30, 0.31, 0.32), wear_rough=0.36),
    "MAKO | Fin grey":        dict(set=SCR, tint=(0.38, 0.46, 0.54), metal=0.0, tile=1.00, shade=(0.94, 1.05), rough=(0.40, 0.30), normal=0.06,
                                   dirt=0.03, dirt_tile=2.0, grime=0.18, wear=0.22, wear_col=(0.42, 0.43, 0.44), wear_rough=0.36,
                                   coat=0.6, coat_rough=0.08, coat_dirt=0.05),
    "MAKO | Side body":       dict(set=SCR, tint=(0.19, 0.23, 0.28), metal=0.0, tile=1.00, shade=(0.94, 1.05), rough=(0.42, 0.32), normal=0.06,
                                   dirt=0.03, dirt_tile=2.0, grime=0.20, wear=0.22, wear_col=(0.40, 0.41, 0.42), wear_rough=0.36,
                                   coat=0.5, coat_rough=0.08, coat_dirt=0.05),
    "MAKO | Recess grey":     dict(set=(STN[0], SCR[1], SCR[2]), tint=(0.25, 0.29, 0.34), metal=0.0, tile=0.90, shade=(0.92, 1.08), rough=(0.48, 0.38), normal=0.10,
                                   dirt=0.05, dirt_tile=1.9, grime=0.25, wear=0.20, wear_col=(0.30, 0.31, 0.32), wear_rough=0.36),
    "MAKO | Frame light":     dict(set=SCR, tint=(0.55, 0.60, 0.64), metal=0.0, tile=0.60, shade=(0.96, 1.03), rough=(0.36, 0.30), normal=0.05,
                                   dirt=0.02, dirt_tile=1.5, grime=0.10, wear=0.10, wear_col=(0.60, 0.61, 0.62), wear_rough=0.36),
    "MAKO | Black gloss":     dict(set=(None, SCR[1], None), tint=(0.006, 0.007, 0.009), metal=0.0, tile=0.30, shade=(1.0, 1.0), rough=(0.12, 0.05), normal=0.0,
                                   dirt=0.05, dirt_tile=0.5, grime=0.0, wear=0.0, coat=1.0, coat_rough=0.03, coat_dirt=0.05),
    "MAKO | Steel":           dict(set=STN, tint=(0.60, 0.61, 0.63), metal=1.0, tile=0.90, shade=(0.93, 1.05), rough=(0.24, 0.32), normal=0.50, normal_map=SCRATCH_N,
                                   dirt=0.05, dirt_tile=0.9, grime=0.20, wear=0.0),
    "MAKO | Rubber":          dict(set=(None, SCR[1], None), tint=(0.017, 0.017, 0.018), metal=0.0, tile=0.40, shade=(1.0, 1.0), rough=(0.86, 0.62), normal=0.0,
                                   dirt=0.08, dirt_tile=1.2, grime=0.10, wear=0.0, dust=0.03),
    "MAKO | Yellow marking":  dict(set=SCR, tint=(0.80, 0.50, 0.025), metal=0.0, tile=0.60, shade=(0.94, 1.04), rough=(0.40, 0.32), normal=0.06,
                                   dirt=0.06, dirt_tile=1.4, grime=0.22, wear=0.50, wear_col=(0.50, 0.51, 0.52), wear_rough=0.36, coat=0.6, coat_rough=0.06, coat_dirt=0.1),
    "MAKO | Cabin panel":     dict(set=SCR, tint=(0.10, 0.14, 0.21), metal=0.0, tile=0.90, shade=(0.90, 1.05), rough=(0.62, 0.42), normal=0.30,
                                   dirt=0.04, dirt_tile=2.0, grime=0.25, wear=0.15, wear_col=(0.45, 0.46, 0.47), wear_rough=0.40),
    "MAKO | Cabin floor":     dict(set=FLOOR, tint=(0.18, 0.18, 0.19), metal=0.6, tile=1.00, shade=(0.80, 1.10), rough=(0.55, 0.35), normal=0.80,
                                   dirt=0.07, dirt_tile=1.5, grime=0.30, wear=0.0),
    "MAKO | Seat":            dict(set=LEATHER, tint=(0.022, 0.045, 0.10), metal=0.0, tile=0.55, shade=(0.80, 1.15), rough=(0.62, 0.40), normal=0.6,
                                   dirt=0.04, dirt_tile=1.2, grime=0.25, wear=0.0, abs_paths=True),
}
GLOW = {"MAKO | Lamp light": ((1.0, 0.92, 0.80), 30.0), "MAKO | Cabin light": ((0.72, 0.86, 1.0), 18.0)}
GLASSES = {"MAKO | Glass upper": (0.42, 0.86, 1.0), "MAKO | Bay glass": (0.05, 0.09, 0.15), "MAKO | Glass": (0.34, 0.82, 1.0), "MAKO | Lamp lens": (0.98, 0.99, 1.0)}   # light blue, as the sheet (passed twice)
SCREEN_INK = (0.35, 0.80, 1.0)        # blue-white instruments, amber for alerts
SCREEN_ALERT = (1.0, 0.55, 0.10)


def screen_uvs(o):
    """every Screen plate: its display face (the biggest face of the plate, on the side toward the cockpit's middle) gets a
    third of the sheet (the console's three in order across, the small ones by turn); the plate's other faces a dark corner"""
    me = o.data
    names = [m.name for m in me.materials]
    si = names.index("MAKO | Screen")
    uvl = me.uv_layers["UVMap"].data
    bm = bmesh.new(); bm.from_mesh(me); bm.faces.ensure_lookup_table()
    todo = {f.index for f in bm.faces if f.material_index == si}
    islands = []
    while todo:
        seed = todo.pop(); isl = [seed]; stack = [seed]
        while stack:
            f = bm.faces[stack.pop()]
            for e in f.edges:
                for g in e.link_faces:
                    if g.index in todo:
                        todo.discard(g.index); isl.append(g.index); stack.append(g.index)
        islands.append(isl)
    bm.free()
    def centre(ids):
        return sum((me.polygons[i].center for i in ids), Vector()) / len(ids)
    islands.sort(key=lambda isl: (round(centre(isl).y, 1), centre(isl).x))
    n_disp = 0
    for k, isl in enumerate(islands):
        c = centre(isl)
        # which way the display looks: the console's three toward the seat, the rear bay's rack lights outboard, the rest up
        D = Vector((0, 0.616, 0.788)) if c.y < -0.6 else (Vector((math.copysign(1, c.x), 0, 0)) if c.y > 0.5 else Vector((0, 0, 1)))
        amax = max(me.polygons[i].area for i in isl)
        best = max(isl, key=lambda i: me.polygons[i].normal.dot(D) + (0.0 if me.polygons[i].area > 0.5 * amax else -2.0))
        third = k % 3
        for i in isl:
            p = me.polygons[i]
            if i == best:
                nrm = p.normal
                right = Vector((1, 0, 0)) if abs(nrm.x) < 0.8 else Vector((0, 1, 0))
                right = (right - right.dot(nrm) * nrm).normalized()
                up = nrm.cross(right).normalized()
                cos_ = [me.vertices[me.loops[li].vertex_index].co for li in range(p.loop_start, p.loop_start + p.loop_total)]
                us = [q.dot(right) for q in cos_]; vs = [q.dot(up) for q in cos_]
                for li, q in zip(range(p.loop_start, p.loop_start + p.loop_total), cos_):
                    uu = (third + 0.04 + 0.92 * (q.dot(right) - min(us)) / max(1e-6, max(us) - min(us))) / 3.0
                    vv = 0.04 + 0.92 * (q.dot(up) - min(vs)) / max(1e-6, max(vs) - min(vs))
                    uvl[li].uv = (uu, vv)
                n_disp += 1
            else:
                for li in range(p.loop_start, p.loop_start + p.loop_total):
                    uvl[li].uv = (0.004, 0.004)
    return n_disp

# ------------------------------------------------------------------ panel seams (v03), drawn per pixel from their measured lines
# The sheet draws the shell as plates with thick dark seams: TOP the transverse band between the front and the rear plates
# (two curved lines a side) and seams across the rear shell; SIDE the rear hull's recessed side panel. They are computed in the
# bake's shader from the surface position (distance to each line), so they are crisp at the atlas's resolution and do not
# depend on where the grid's vertices are. The collar, the dorsal spine and the shell's edge are real geometry.
_MS = {}
exec(open(os.path.join(HERE, "mk_measure.py"), encoding="utf-8").read(), _MS)
TOP_SEGS = []
for _line in _MS["PLATE_LINES"]:
    _pts = [(lat * _MS["S"], _MS["Yi"](x)) for x, lat in _line]
    for _sg in (1, -1):
        _P = [(_sg * a, b) for a, b in _pts]
        TOP_SEGS += list(zip(_P, _P[1:]))
for _x in (352.0, 730.0):                                    # v04: behind the bezel (SIDE x 345-352), across the rear shell
    TOP_SEGS.append(((-0.75, _MS["Yi"](_x)), (0.75, _MS["Yi"](_x))))
SIDE_SEGS = [((_MS["Yi"](480.0), _MS["Zi"](830.0)), (_MS["Yi"](640.0), _MS["Zi"](763.0)))]   # v04: SIDE's diagonal stripe
for _line in _MS.get("SIDE_SEAMS4", []):                     # v04 round 2: the rear capsule's rounded front
    _P = [(_MS["Yi"](x), _MS["Zi"](v)) for x, v in _line]
    SIDE_SEGS += list(zip(_P, _P[1:]))


def seam_value(nt, loc_x=0):
    """0..1: how much of a panel seam this point is (the shell's plate seams seen from above, the rear side panel's outline)"""
    N = nt.nodes
    L = nt.links
    def mth(op, a, b=None, clamp=False):
        n = N.new('ShaderNodeMath'); n.operation = op; n.use_clamp = clamp
        for i, v in enumerate((a, b)):
            if v is None:
                continue
            if isinstance(v, (int, float)):
                n.inputs[i].default_value = v
            else:
                L.new(v, n.inputs[i])
        return n.outputs[0]
    def vm(op, a, b=None):
        n = N.new('ShaderNodeVectorMath'); n.operation = op
        for i, v in enumerate((a, b)):
            if v is None:
                continue
            if isinstance(v, (tuple, list)):
                n.inputs[i].default_value = v
            else:
                L.new(v, n.inputs[i])
        return n.outputs["Value"] if op in ('LENGTH', 'DOT_PRODUCT') else n.outputs["Vector"]
    geo = N.new('ShaderNodeNewGeometry')
    sp = N.new('ShaderNodeSeparateXYZ'); L.new(geo.outputs["Position"], sp.inputs[0])
    sn = N.new('ShaderNodeSeparateXYZ'); L.new(geo.outputs["Normal"], sn.inputs[0])
    cxy = N.new('ShaderNodeCombineXYZ'); L.new(sp.outputs["X"], cxy.inputs[0]); L.new(sp.outputs["Y"], cxy.inputs[1])
    P2 = cxy.outputs[0]
    dmin = None
    for (a, b) in TOP_SEGS:
        A = (a[0], a[1], 0.0); AB = (b[0] - a[0], b[1] - a[1], 0.0)
        L2 = max(1e-9, AB[0] ** 2 + AB[1] ** 2)
        pa = vm('SUBTRACT', P2, A)
        t = mth('DIVIDE', vm('DOT_PRODUCT', pa, AB), L2, clamp=True)
        sc = N.new('ShaderNodeVectorMath'); sc.operation = 'SCALE'; sc.inputs[0].default_value = AB; L.new(t, sc.inputs["Scale"])
        d = vm('LENGTH', vm('SUBTRACT', pa, sc.outputs["Vector"]))
        dmin = d if dmin is None else mth('MINIMUM', dmin, d)
    gate_top = mth('GREATER_THAN', sn.outputs["Z"], 0.35)
    def band(d):
        mr = N.new('ShaderNodeMapRange'); mr.interpolation_type = 'SMOOTHSTEP'
        mr.inputs["From Min"].default_value = SEAM_W0; mr.inputs["From Max"].default_value = SEAM_W1
        mr.inputs["To Min"].default_value = 0.0; mr.inputs["To Max"].default_value = 1.0
        L.new(d, mr.inputs["Value"])
        return mr.outputs["Result"]
    top = mth('MULTIPLY', band(dmin), gate_top)
    # v04: SIDE's diagonal panel edge on both flanks, from the cheek panel up and back to the glass bay
    cyz = N.new('ShaderNodeCombineXYZ'); L.new(sp.outputs["Y"], cyz.inputs[0]); L.new(sp.outputs["Z"], cyz.inputs[1])
    Q2 = cyz.outputs[0]
    ds = None
    for (a, b) in SIDE_SEGS:
        A = (a[0], a[1], 0.0); AB = (b[0] - a[0], b[1] - a[1], 0.0)
        L2 = max(1e-9, AB[0] ** 2 + AB[1] ** 2)
        pa = vm('SUBTRACT', Q2, A)
        t = mth('DIVIDE', vm('DOT_PRODUCT', pa, AB), L2, clamp=True)
        sc = N.new('ShaderNodeVectorMath'); sc.operation = 'SCALE'; sc.inputs[0].default_value = AB; L.new(t, sc.inputs["Scale"])
        d = vm('LENGTH', vm('SUBTRACT', pa, sc.outputs["Vector"]))
        ds = d if ds is None else mth('MINIMUM', ds, d)
    gate_side = mth('GREATER_THAN', mth('ABSOLUTE', sn.outputs["X"]), 0.55)
    gate_lat = mth('GREATER_THAN', mth('ABSOLUTE', sp.outputs["X"]), 0.15)
    side = mth('MULTIPLY', mth('MULTIPLY', band(ds), gate_side), gate_lat)
    return mth('MAXIMUM', top, side)


def dress_glass(mat, tint):
    """the canopy reads light blue, as the sheet draws it: mostly clear, a tenth of it a light-blue body colour, a clear coat
    for the sheen, fingerprints and dust in the roughness"""
    t = Tree(mat)
    b = t.node('ShaderNodeBsdfPrincipled')
    out = t.node('ShaderNodeOutputMaterial')
    t.link(b.outputs[0], out.inputs["Surface"])
    dirt = t.tex(tex_image(os.path.join(PACK, DIRT), "Non-Color"), t.mapped(t.uv("UV_Tile"), 1.0 / 0.7), "smudges and dust").outputs["Color"]
    b.inputs["Base Color"].default_value = tuple(tint) + (1,)
    b.inputs["Transmission Weight"].default_value = {"MAKO | Glass": 0.90, "MAKO | Glass upper": 0.70,
                                                     "MAKO | Bay glass": 0.55}.get(mat.name, 1.0)   # v04 round 2: sky cyan, pale above, dark bay
    b.inputs["IOR"].default_value = 1.12                               # v04: a thin canopy, little distortion or glare (the cockpit shows)
    b.inputs["Coat Weight"].default_value = 0.15
    b.inputs["Coat Roughness"].default_value = 0.02
    t.put(b.inputs["Roughness"], t.ranged(dirt, 0.0, 0.03, "smudged", 0.0, 1.0))
    mat.diffuse_color = tuple(tint) + (0.35,)
    try:
        mat.surface_render_method = 'DITHERED'
    except Exception:
        pass