# Visual: paint masks as vertex colors and render the hip view — where are the lines
# relative to the garment mask (red), rim band (green), slit seams (blue)?
#   blender --background --python dbg_maskviz.py -- <blend> <masks.npz> <raw.glb> <out.png>
import bpy, sys, math
import numpy as np
from mathutils import Vector, Euler
from mathutils.kdtree import KDTree

argv = sys.argv[sys.argv.index("--") + 1:]
BLEND, MASKS, RAW, OUT = argv[0], argv[1], argv[2], argv[3]

bpy.ops.wm.open_mainfile(filepath=BLEND)
ob = max([o for o in bpy.data.objects if o.type == 'MESH'],
         key=lambda o: len(o.data.vertices))
me = ob.data
n_v = len(me.vertices)
co = np.empty(n_v * 3)
me.vertices.foreach_get("co", co)
co = co.reshape(-1, 3)
M = np.load(MASKS)
garment = M["garment"]

n_e = len(me.edges)
ev = np.empty(n_e * 2, dtype=np.int32)
me.edges.foreach_get("vertices", ev)
ev = ev.reshape(-1, 2)

before = set(bpy.data.objects)
bpy.ops.import_scene.gltf(filepath=RAW)
new = [o for o in bpy.data.objects if o not in before]
raw = max([o for o in new if o.type == 'MESH'], key=lambda o: len(o.data.vertices))
rme = raw.data
rn = len(rme.vertices)
rco = np.empty(rn * 3)
rme.vertices.foreach_get("co", rco)
rco = rco.reshape(-1, 3)
l_tot = np.empty(len(rme.polygons), dtype=np.int32)
rme.polygons.foreach_get("loop_total", l_tot)
l_start = np.empty(len(rme.polygons), dtype=np.int32)
rme.polygons.foreach_get("loop_start", l_start)
l_v = np.empty(len(rme.loops), dtype=np.int32)
rme.loops.foreach_get("vertex_index", l_v)
ecount = {}
for fs, ft in zip(l_start, l_tot):
    idxs = l_v[fs:fs + ft]
    for k in range(ft):
        a, b = idxs[k], idxs[(k + 1) % ft]
        kk = (a, b) if a < b else (b, a)
        ecount[kk] = ecount.get(kk, 0) + 1
rbnd = np.zeros(rn, dtype=bool)
for (a, b), c in ecount.items():
    if c == 1:
        rbnd[a] = rbnd[b] = True
for o in new:
    bpy.data.objects.remove(o, do_unlink=True)

kd = KDTree(n_v)
for i in range(n_v):
    kd.insert(Vector(co[i]), i)
kd.balance()
seam = np.zeros(n_v, dtype=bool)
for p in rco[rbnd]:
    hit = kd.find(Vector(p))
    if hit[0] is not None and hit[2] < 0.002:
        seam[hit[1]] = True


def grow(mask, rings):
    m = mask.copy()
    for _ in range(rings):
        h = m[ev[:, 0]] | m[ev[:, 1]]
        m2 = m.copy()
        m2[ev[:, 0]] |= h
        m2[ev[:, 1]] |= h
        m = m2
    return m


rim = grow(garment, 10) & ~(~grow(~garment, 10))
seam_g = grow(seam, 1)

col = np.ones((n_v, 4))
col[:, :3] = 0.75
col[garment, 0] = 0.9
col[garment, 1] = 0.45
col[garment, 2] = 0.45
col[rim, 0] = 0.3
col[rim, 1] = 0.8
col[rim, 2] = 0.3
col[seam_g, 0] = 0.15
col[seam_g, 1] = 0.25
col[seam_g, 2] = 0.95

vc = me.color_attributes.new("viz", 'FLOAT_COLOR', 'POINT')
vc.data.foreach_set("color", col.reshape(-1))

mat = bpy.data.materials.new("viz")
mat.use_nodes = True
nt = mat.node_tree
b = nt.nodes["Principled BSDF"]
attr = nt.nodes.new("ShaderNodeVertexColor")
attr.layer_name = "viz"
nt.links.new(attr.outputs["Color"], b.inputs["Base Color"])
b.inputs["Roughness"].default_value = 0.9
me.materials.clear()
me.materials.append(mat)

for o in list(bpy.data.objects):
    if o.type in ('LIGHT', 'CAMERA'):
        bpy.data.objects.remove(o, do_unlink=True)
sc = bpy.context.scene
sc.render.engine = 'BLENDER_EEVEE'
sc.render.resolution_x = sc.render.resolution_y = 1000
wd = bpy.data.worlds.new("w")
wd.use_nodes = True
wd.node_tree.nodes["Background"].inputs[0].default_value = (0.18, 0.18, 0.18, 1)
sc.world = wd
ld = bpy.data.lights.new("s", 'SUN')
ld.energy = 3.0
ld.use_shadow = False
lo = bpy.data.objects.new("s", ld)
lo.rotation_euler = Euler((math.radians(60), 0, math.radians(-30)))
bpy.context.collection.objects.link(lo)
ld2 = bpy.data.lights.new("s2", 'SUN')
ld2.energy = 1.2
ld2.use_shadow = False
lo2 = bpy.data.objects.new("s2", ld2)
lo2.rotation_euler = Euler((math.radians(120), 0, math.radians(150)))
bpy.context.collection.objects.link(lo2)

cam = bpy.data.cameras.new("c")
cam.lens = 60
cob = bpy.data.objects.new("c", cam)
bpy.context.collection.objects.link(cob)
sc.camera = cob
cob.location = Vector((0.0, -0.95, 0.53))
cob.rotation_euler = Euler((math.radians(90), 0, 0))
sc.render.filepath = OUT
bpy.ops.render.render(write_still=True)
print("VIZ_DONE")