# Re-UVs the QuatSkin Lena with a Quaternius-style semantic atlas and transfers
# her skin texture into it (v2).
# v1 used planar +-Z for head/torso: side/down-facing surfaces (cheeks by the
# ears, under the jaw, armpit side strips) were edge-on to the projection ->
# texel starvation -> smear streaks. v2 unwraps head and torso as CYLINDERS
# around the vertical axis (u = angle x local radius, v = height) with small
# planar caps (crown, under-jaw, shoulder tops), and smooths the per-triangle
# region assignment so the arm/torso border stops jittering in the armpit.
#   arms/legs  cylinder around the bone chain (unchanged from v1)
#   torso/head cylinder around +Y, seam at center back, + planar caps
#   hands/feet planar +-Y (unchanged)
# Run headless:
#   tinqs.console.exe --headless --path tattoo-test -s res://reuv_lena.gd
extends SceneTree

const GLB_IN := "C:/Users/CAN/tinqs-ltd/ariki-game/assets/quaternius/derived-bodies/Ariki_Female_QuatSkin.glb"
const GLB_OUT := "C:/Users/CAN/tinqs-ltd/tattoo-test/lena_clean_uv.glb"
const TEX_OUT := "C:/Users/CAN/tinqs-ltd/tattoo-test/lena_clean_uv_base.png"
const ATLAS := 4096
const PAD_PX := 24
const JB := 0.03

var skel: Skeleton3D
var verts: PackedVector3Array
var norms: PackedVector3Array
var uvs_old: PackedVector2Array
var vbones: PackedInt32Array
var vweights: PackedFloat32Array
var indices: PackedInt32Array
var nverts: int
var influences: int
var region_w := {}

const ISLAND_ORDER := [
	"head_cyl", "head_top", "head_bot",
	"torso_cyl", "torso_top", "torso_bot",
	"arm_r", "arm_l", "leg_r", "leg_l",
	"hand_r_back", "hand_r_palm", "hand_l_back", "hand_l_palm",
	"foot_r_top", "foot_r_sole", "foot_l_top", "foot_l_sole",
]

func _init() -> void:
	var t_start := Time.get_ticks_msec()
	var doc := GLTFDocument.new()
	var state := GLTFState.new()
	if doc.append_from_file(GLB_IN, state) != OK:
		push_error("GLB load failed")
		quit(1)
		return
	var scene := doc.generate_scene(state)

	var body: MeshInstance3D = null
	var stack: Array = [scene]
	while not stack.is_empty():
		var n: Node = stack.pop_back()
		if n is Skeleton3D and skel == null:
			skel = n
		if n is MeshInstance3D and String(n.name).to_lower().contains("lena"):
			body = n
		for c in n.get_children():
			stack.push_back(c)

	var mesh: ArrayMesh = body.mesh
	var arrays: Array = mesh.surface_get_arrays(0)
	verts = arrays[Mesh.ARRAY_VERTEX]
	norms = arrays[Mesh.ARRAY_NORMAL]
	uvs_old = arrays[Mesh.ARRAY_TEX_UV]
	vbones = arrays[Mesh.ARRAY_BONES]
	vweights = arrays[Mesh.ARRAY_WEIGHTS]
	indices = arrays[Mesh.ARRAY_INDEX]
	nverts = verts.size()
	influences = int(float(vbones.size()) / float(nverts))

	var skin: Skin = body.skin
	var bind_bone := PackedInt32Array()
	for i in skin.get_bind_count():
		var b := skin.get_bind_bone(i)
		if b < 0:
			b = skel.find_bone(skin.get_bind_name(i))
		bind_bone.append(b)

	# ── region weights ──────────────────────────────────────────────────────
	var fingers_l := []
	var fingers_r := []
	for f in ["index", "middle", "pinky", "ring", "thumb"]:
		for seg in ["01", "02", "03", "04_leaf"]:
			fingers_l.append("%s_%s_l" % [f, seg])
			fingers_r.append("%s_%s_r" % [f, seg])
	var regions := {
		"head": ["Head", "neck_01"],
		"torso": ["pelvis", "spine_01", "spine_02", "spine_03", "clavicle_l", "clavicle_r"],
		"arm_l": ["upperarm_l", "lowerarm_l"],
		"arm_r": ["upperarm_r", "lowerarm_r"],
		"hand_l": ["hand_l"] + fingers_l,
		"hand_r": ["hand_r"] + fingers_r,
		"leg_l": ["thigh_l", "calf_l"],
		"leg_r": ["thigh_r", "calf_r"],
		"foot_l": ["foot_l", "ball_l", "ball_leaf_l"],
		"foot_r": ["foot_r", "ball_r", "ball_leaf_r"],
	}
	for key in regions:
		var binds := _bind_set(bind_bone, regions[key])
		var w := PackedFloat32Array()
		w.resize(nverts)
		for v in nverts:
			for k in influences:
				if binds.has(vbones[v * influences + k]):
					w[v] += vweights[v * influences + k]
		region_w[key] = w

	# ── per-triangle region (argmax) + adjacency smoothing ──────────────────
	var ntris := indices.size() / 3
	var tri_region := PackedStringArray()
	tri_region.resize(ntris)
	for t in ntris:
		var best := "torso"
		var best_s := -1.0
		for key in region_w:
			var w: PackedFloat32Array = region_w[key]
			var s: float = w[indices[t * 3]] + w[indices[t * 3 + 1]] + w[indices[t * 3 + 2]]
			if s > best_s:
				best_s = s
				best = key
		tri_region[t] = best

	# edge adjacency
	var edge_tri := {}
	var tri_nbrs := []
	tri_nbrs.resize(ntris)
	for t in ntris:
		tri_nbrs[t] = []
	for t in ntris:
		for j in 3:
			var a := indices[t * 3 + j]
			var b := indices[t * 3 + ((j + 1) % 3)]
			var key := "%d_%d" % [mini(a, b), maxi(a, b)]
			if edge_tri.has(key):
				var o: int = edge_tri[key]
				tri_nbrs[t].append(o)
				tri_nbrs[o].append(t)
			else:
				edge_tri[key] = t
	# majority smoothing: kills 1-triangle islands at region borders (armpit!)
	for pass_i in 3:
		var changed := 0
		var next := tri_region.duplicate()
		for t in ntris:
			var counts := {}
			counts[tri_region[t]] = 1
			for o in tri_nbrs[t]:
				var r: String = tri_region[o]
				counts[r] = counts.get(r, 0) + 1
			var best := tri_region[t]
			var best_c: int = counts[best]
			for r in counts:
				if counts[r] > best_c:
					best_c = counts[r]
					best = r
			if best != tri_region[t]:
				next[t] = best
				changed += 1
		tri_region = next
		print("region smoothing pass %d: %d flips" % [pass_i, changed])

	# ── axes for the vertical cylinders ──────────────────────────────────────
	var head_pos := _bone_pos("Head")
	var head_ax := Vector2(head_pos.x, head_pos.z)
	var tz := 0.0
	for bn in ["pelvis", "spine_01", "spine_02", "spine_03"]:
		tz += _bone_pos(bn).z
	var torso_ax := Vector2(0.0, tz / 4.0)

	# ── limb chains (unchanged) ──────────────────────────────────────────────
	var chains := {
		"arm_l": ["upperarm_l", "lowerarm_l", "hand_l", "spine_03"],
		"arm_r": ["upperarm_r", "lowerarm_r", "hand_r", "spine_03"],
		"leg_l": ["thigh_l", "calf_l", "foot_l", "pelvis"],
		"leg_r": ["thigh_r", "calf_r", "foot_r", "pelvis"],
	}
	var limb := {}
	for key in chains:
		var c: Array = chains[key]
		var p0 := _bone_pos(c[0])
		var p1 := _bone_pos(c[1])
		var p2 := _bone_pos(c[2])
		var center := _bone_pos(c[3])
		var axis := (p2 - p0).normalized()
		var outward := (p0 - center).normalized()
		var uref := (outward - axis * outward.dot(axis)).normalized()
		limb[key] = {"p0": p0, "p1": p1, "p2": p2,
			"len1": p0.distance_to(p1), "len2": p1.distance_to(p2),
			"uref": uref, "vref": axis.cross(uref)}

	# ── per-triangle island id (with normal-based caps) + smoothing ─────────
	var tri_island := PackedStringArray()
	tri_island.resize(ntris)
	for t in ntris:
		var reg := tri_region[t]
		var i0 := indices[t * 3]
		var an := (norms[i0] + norms[indices[t * 3 + 1]] + norms[indices[t * 3 + 2]]).normalized()
		var island := reg
		if reg == "head":
			island = "head_top" if an.y > 0.55 else ("head_bot" if an.y < -0.60 else "head_cyl")
		elif reg == "torso":
			island = "torso_top" if an.y > 0.65 else ("torso_bot" if an.y < -0.75 else "torso_cyl")
		elif reg.begins_with("hand"):
			island = reg + ("_back" if an.y >= 0.0 else "_palm")
		elif reg.begins_with("foot"):
			island = reg + ("_top" if an.y >= 0.0 else "_sole")
		tri_island[t] = island
	for pass_i in 2:
		var next := tri_island.duplicate()
		for t in ntris:
			var counts := {}
			counts[tri_island[t]] = 1
			for o in tri_nbrs[t]:
				if tri_region[o] == tri_region[t]:
					counts[tri_island[o]] = counts.get(tri_island[o], 0) + 1
			var best := tri_island[t]
			var best_c: int = counts[best]
			for r in counts:
				if counts[r] > best_c:
					best_c = counts[r]
					best = r
			next[t] = best
		tri_island = next

	# ── per-corner meter coords ──────────────────────────────────────────────
	var corner_um := PackedFloat32Array()
	corner_um.resize(indices.size())
	var corner_vm := PackedFloat32Array()
	corner_vm.resize(indices.size())
	for t in ntris:
		var island := tri_island[t]
		for j in 3:
			var ci := t * 3 + j
			var p := verts[indices[ci]]
			match island:
				"head_top", "torso_top":
					corner_um[ci] = p.x
					corner_vm[ci] = p.z
				"head_bot", "torso_bot":
					corner_um[ci] = p.x
					corner_vm[ci] = -p.z
				"head_cyl":
					corner_um[ci] = atan2(p.x - head_ax.x, p.z - head_ax.y)
					corner_vm[ci] = -p.y
				"torso_cyl":
					corner_um[ci] = atan2(p.x - torso_ax.x, p.z - torso_ax.y)
					corner_vm[ci] = -p.y
				"hand_l_back", "hand_r_back":
					corner_um[ci] = p.z
					corner_vm[ci] = absf(p.x)
				"hand_l_palm", "hand_r_palm":
					corner_um[ci] = -p.z
					corner_vm[ci] = absf(p.x)
				"foot_l_top", "foot_r_top":
					corner_um[ci] = p.x
					corner_vm[ci] = p.z
				"foot_l_sole", "foot_r_sole":
					corner_um[ci] = -p.x
					corner_vm[ci] = p.z
				_:
					var L: Dictionary = limb[island]
					var res := _limb_coords(p, L)
					corner_um[ci] = res.x
					corner_vm[ci] = res.y
		# cylinder wrap: one angular branch per triangle, then angle -> arc meters
		var is_cyl := island.begins_with("arm") or island.begins_with("leg") \
			or island == "head_cyl" or island == "torso_cyl"
		if is_cyl:
			var a0 := corner_um[t * 3]
			var a1 := corner_um[t * 3 + 1]
			var a2 := corner_um[t * 3 + 2]
			if maxf(a0, maxf(a1, a2)) - minf(a0, minf(a1, a2)) > PI:
				for j in 3:
					if corner_um[t * 3 + j] < 0.0:
						corner_um[t * 3 + j] += TAU
			for j in 3:
				var ci2 := t * 3 + j
				var p2 := verts[indices[ci2]]
				var r := 0.015
				if island == "head_cyl":
					r = maxf(Vector2(p2.x - head_ax.x, p2.z - head_ax.y).length(), 0.015)
				elif island == "torso_cyl":
					r = maxf(Vector2(p2.x - torso_ax.x, p2.z - torso_ax.y).length(), 0.015)
				else:
					r = _limb_radius(p2, limb[island])
				corner_um[ci2] = corner_um[ci2] * r

	# ── island bboxes, pack, weld, transfer, export ──────────────────────────
	var isl := {}
	for name in ISLAND_ORDER:
		isl[name] = {"mn": Vector2(1e9, 1e9), "mx": Vector2(-1e9, -1e9)}
	for t in ntris:
		for j in 3:
			var ci := t * 3 + j
			var m: Dictionary = isl[tri_island[t]]
			var q := Vector2(corner_um[ci], corner_vm[ci])
			m["mn"] = (m["mn"] as Vector2).min(q)
			m["mx"] = (m["mx"] as Vector2).max(q)

	var lo := 100.0
	var hi := 4000.0
	for it in 24:
		var mid := (lo + hi) * 0.5
		if _try_pack(isl, mid, false):
			lo = mid
		else:
			hi = mid
	var scale := lo
	_try_pack(isl, scale, true)
	print("atlas scale: %.0f px/m" % scale)
	for name in ISLAND_ORDER:
		var m: Dictionary = isl[name]
		print("  %s -> (%d,%d) %dx%d px" % [name, m["x"], m["y"], m["w"], m["h"]])

	var new_verts := PackedVector3Array()
	var new_norms := PackedVector3Array()
	var new_uvs := PackedVector2Array()
	var new_uvs_old := PackedVector2Array()
	var new_bones := PackedInt32Array()
	var new_weights := PackedFloat32Array()
	var new_indices := PackedInt32Array()
	new_indices.resize(indices.size())
	var weld := {}
	for t in ntris:
		for j in 3:
			var ci := t * 3 + j
			var v := indices[ci]
			var m: Dictionary = isl[tri_island[t]]
			var px: float = m["x"] + (corner_um[ci] - (m["mn"] as Vector2).x) * scale
			var py: float = m["y"] + (corner_vm[ci] - (m["mn"] as Vector2).y) * scale
			var uv := Vector2(px / ATLAS, py / ATLAS)
			var key := "%d_%d_%d" % [v, int(round(px * 4.0)), int(round(py * 4.0))]
			if not weld.has(key):
				var nv := new_verts.size()
				new_verts.append(verts[v])
				new_norms.append(norms[v])
				new_uvs.append(uv)
				new_uvs_old.append(uvs_old[v])
				for k in influences:
					new_bones.append(vbones[v * influences + k])
					new_weights.append(vweights[v * influences + k])
				weld[key] = nv
			new_indices[ci] = weld[key]
	print("verts %d -> %d" % [nverts, new_verts.size()])

	var mat: BaseMaterial3D = mesh.surface_get_material(0)
	var src: Image = mat.albedo_texture.get_image()
	src.clear_mipmaps()
	src.convert(Image.FORMAT_RGBA8)
	var avg := Color(0, 0, 0)
	var n_avg := 0
	for sy in range(0, src.get_height(), 61):
		for sx in range(0, src.get_width(), 61):
			var c := src.get_pixel(sx, sy)
			if (c.r + c.g + c.b) / 3.0 > 0.2:
				avg += c
				n_avg += 1
	avg = Color(avg.r / n_avg, avg.g / n_avg, avg.b / n_avg)
	var dst := Image.create(ATLAS, ATLAS, false, Image.FORMAT_RGBA8)
	dst.fill(avg)

	var painted := 0
	for t in ntris:
		painted += _transfer_tri(dst, src,
			new_indices[t * 3], new_indices[t * 3 + 1], new_indices[t * 3 + 2],
			new_uvs, new_uvs_old)
	print("transferred px: ", painted)
	dst.save_png(TEX_OUT)

	var new_arrays := []
	new_arrays.resize(Mesh.ARRAY_MAX)
	new_arrays[Mesh.ARRAY_VERTEX] = new_verts
	new_arrays[Mesh.ARRAY_NORMAL] = new_norms
	new_arrays[Mesh.ARRAY_TEX_UV] = new_uvs
	new_arrays[Mesh.ARRAY_BONES] = new_bones
	new_arrays[Mesh.ARRAY_WEIGHTS] = new_weights
	new_arrays[Mesh.ARRAY_INDEX] = new_indices
	var new_mesh := ArrayMesh.new()
	new_mesh.add_surface_from_arrays(Mesh.PRIMITIVE_TRIANGLES, new_arrays)
	mat.albedo_texture = ImageTexture.create_from_image(dst)
	new_mesh.surface_set_material(0, mat)
	body.mesh = new_mesh

	var out_doc := GLTFDocument.new()
	var out_state := GLTFState.new()
	if out_doc.append_from_scene(scene, out_state) != OK:
		push_error("export append failed")
		quit(1)
		return
	if out_doc.write_to_filesystem(out_state, GLB_OUT) != OK:
		push_error("export write failed")
		quit(1)
		return
	print("saved: %s + %s" % [GLB_OUT, TEX_OUT])
	print("took %d ms" % (Time.get_ticks_msec() - t_start))
	quit()

func _limb_coords(pos: Vector3, L: Dictionary) -> Vector2:
	var p0: Vector3 = L["p0"]
	var p1: Vector3 = L["p1"]
	var p2: Vector3 = L["p2"]
	var len1: float = L["len1"]
	var len2: float = L["len2"]
	var ax1 := (p1 - p0) / len1
	var ax2 := (p2 - p1) / len2
	var s1 := clampf((pos - p0).dot(ax1), 0.0, len1)
	var cp1 := p0 + ax1 * s1
	var s2 := clampf((pos - p1).dot(ax2), 0.0, len2)
	var cp2 := p1 + ax2 * s2
	var wb := smoothstep(-JB, JB, pos.distance_to(cp1) - pos.distance_to(cp2))
	var t_m := lerpf(s1, len1 + s2, wb)
	var dv: Vector3 = pos - cp1.lerp(cp2, wb)
	var ang := atan2(dv.dot(L["vref"]), dv.dot(L["uref"]))
	return Vector2(ang, t_m)

func _limb_radius(pos: Vector3, L: Dictionary) -> float:
	var p0: Vector3 = L["p0"]
	var p1: Vector3 = L["p1"]
	var p2: Vector3 = L["p2"]
	var len1: float = L["len1"]
	var len2: float = L["len2"]
	var ax1 := (p1 - p0) / len1
	var ax2 := (p2 - p1) / len2
	var s1 := clampf((pos - p0).dot(ax1), 0.0, len1)
	var cp1 := p0 + ax1 * s1
	var s2 := clampf((pos - p1).dot(ax2), 0.0, len2)
	var cp2 := p1 + ax2 * s2
	var wb := smoothstep(-JB, JB, pos.distance_to(cp1) - pos.distance_to(cp2))
	return maxf(pos.distance_to(cp1.lerp(cp2, wb)), 0.015)

func _try_pack(isl: Dictionary, scale: float, commit: bool) -> bool:
	var cx := PAD_PX
	var cy := PAD_PX
	var row_h := 0
	for name in ISLAND_ORDER:
		var m: Dictionary = isl[name]
		var w := int(((m["mx"] as Vector2).x - (m["mn"] as Vector2).x) * scale) + 1
		var h := int(((m["mx"] as Vector2).y - (m["mn"] as Vector2).y) * scale) + 1
		if w > ATLAS - 2 * PAD_PX:
			return false
		if cx + w >= ATLAS - PAD_PX:
			cx = PAD_PX
			cy += row_h + PAD_PX
			row_h = 0
		if cy + h >= ATLAS - PAD_PX:
			return false
		if commit:
			m["x"] = cx
			m["y"] = cy
			m["w"] = w
			m["h"] = h
		cx += w + PAD_PX
		row_h = maxi(row_h, h)
	return true

func _transfer_tri(dst: Image, src: Image, i0: int, i1: int, i2: int,
		nuv: PackedVector2Array, ouv: PackedVector2Array) -> int:
	var p0 := nuv[i0] * ATLAS
	var p1 := nuv[i1] * ATLAS
	var p2 := nuv[i2] * ATLAS
	var denom := (p1.y - p2.y) * (p0.x - p2.x) + (p2.x - p1.x) * (p0.y - p2.y)
	if absf(denom) < 1e-9:
		return 0
	var minx := clampi(int(floor(min(p0.x, min(p1.x, p2.x)))) - 1, 0, ATLAS - 1)
	var maxx := clampi(int(ceil(max(p0.x, max(p1.x, p2.x)))) + 1, 0, ATLAS - 1)
	var miny := clampi(int(floor(min(p0.y, min(p1.y, p2.y)))) - 1, 0, ATLAS - 1)
	var maxy := clampi(int(ceil(max(p0.y, max(p1.y, p2.y)))) + 1, 0, ATLAS - 1)
	var painted := 0
	for py in range(miny, maxy + 1):
		for px in range(minx, maxx + 1):
			var fx := px + 0.5
			var fy := py + 0.5
			var b0 := ((p1.y - p2.y) * (fx - p2.x) + (p2.x - p1.x) * (fy - p2.y)) / denom
			var b1 := ((p2.y - p0.y) * (fx - p2.x) + (p0.x - p2.x) * (fy - p2.y)) / denom
			var b2 := 1.0 - b0 - b1
			if b0 < -0.06 or b1 < -0.06 or b2 < -0.06:
				continue
			b0 = clampf(b0, 0.0, 1.0)
			b1 = clampf(b1, 0.0, 1.0)
			b2 = clampf(b2, 0.0, 1.0)
			var s := b0 + b1 + b2
			var ou: Vector2 = (ouv[i0] * b0 + ouv[i1] * b1 + ouv[i2] * b2) / s
			dst.set_pixel(px, py, _sample_bilinear(src, ou.x, ou.y))
			painted += 1
	return painted

func _bind_set(bind_bone: PackedInt32Array, names: Array) -> Dictionary:
	var bone_ids := {}
	for bn in names:
		var idx := skel.find_bone(bn)
		if idx < 0:
			push_error("bone not found: " + str(bn))
		bone_ids[idx] = true
	var out := {}
	for i in bind_bone.size():
		if bone_ids.has(bind_bone[i]):
			out[i] = true
	return out

func _bone_pos(bone: String) -> Vector3:
	return skel.get_bone_global_rest(skel.find_bone(bone)).origin

func _sample_bilinear(img: Image, u: float, v: float) -> Color:
	var w := img.get_width()
	var h := img.get_height()
	var x := clampf(u * w - 0.5, 0.0, w - 1.001)
	var y := clampf(v * h - 0.5, 0.0, h - 1.001)
	var x0 := int(x)
	var y0 := int(y)
	var x1 := mini(x0 + 1, w - 1)
	var y1 := mini(y0 + 1, h - 1)
	var fx := x - x0
	var fy := y - y0
	return img.get_pixel(x0, y0).lerp(img.get_pixel(x1, y0), fx) \
		.lerp(img.get_pixel(x0, y1).lerp(img.get_pixel(x1, y1), fx), fy)