"""Opens the rigged file fresh and checks it, then re-imports the skeletal FBX and compares it with the file.

    blender.exe -b --python-exit-code 1 --python verify_rig.py
"""
import bpy, os, math, json
from mathutils import Vector, Matrix

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.dirname(HERE)
ROT = Matrix.Rotation(-math.pi / 2, 4, 'Z')          # the FBX data is turned like this (and written in cm)
ok = True
def check(label, cond, detail=""):
    global ok
    ok &= bool(cond)
    print("RIGVERIFY", "ok  " if cond else "FAIL", label, detail)

bpy.ops.wm.open_mainfile(filepath=os.path.join(OUT, "DockArm_Carousel_Rigged.blend"))
vl = bpy.context.view_layer
rig = bpy.data.objects["DockArm_Rig"]
ad = rig.data
deform = [b.name for b in ad.bones if b.use_deform]
check("bones", len(ad.bones) == 98 and len(deform) == 86, "%d bones, %d deform (23 joints + 63 hose fans)" % (len(ad.bones), len(deform)))
drv = [fc.driver for fc in rig.animation_data.drivers]
check("drivers simple and valid (no auto-run needed)", all(d.is_simple_expression and d.is_valid for d in drv), "%d drivers" % len(drv))
check("no linked libraries, image packed", not bpy.data.libraries and all(i.packed_file for i in bpy.data.images if i.type == 'IMAGE'))
meshes = [o for o in bpy.data.objects if o.type == 'MESH' and o.name.startswith("SK_")]
bad = []
for o in meshes:
    names = [g.name for g in o.vertex_groups]
    hose = o.name.endswith("_OreHose")
    one = all(n in deform for n in names) and (len(names) == 25 if hose else len(names) == 1)      # a hose: 4 runs + 21 fans
    full = one and all(len(v.groups) == 1 and abs(v.groups[0].weight - 1.0) < 1e-6 for v in o.data.vertices)
    mod = [m for m in o.modifiers if m.type == 'ARMATURE' and m.object == rig]
    if not (one and full and len(mod) == 1 and o.parent == rig):
        bad.append(o.name)
check("26 skinned meshes (23 parts, 3 hoses), every vertex on one bone with weight 1", len(meshes) == 26 and not bad, str(bad))
check("every deform bone carries geometry", sorted({g.name for o in meshes for g in o.vertex_groups}) == sorted(deform))
check("rest pose saved (all channels zero)", all(pb.location.length < 1e-9 and Vector(pb.rotation_euler).length < 1e-9 for pb in rig.pose.bones))

# pose through the handles in the freshly opened file
P = rig.pose.bones
rest_head = {n: ad.bones[n].head_local.copy() for n in deform}
P["CTRL_Arm1_Hoist"].location.y = 5.3
P["CTRL_Arm3_Hoist"].location.y = 2.0
P["CTRL_Arm1_Claws"].rotation_euler.x = math.radians(25)
rig.update_tag(); vl.update()
dg = bpy.context.evaluated_depsgraph_get()
E = rig.evaluated_get(dg).pose.bones
w1 = E["Arm1_Head"].head - rest_head["Arm1_Head"]
w3 = E["Arm3_Head"].head - rest_head["Arm3_Head"]
check("handle 5.3 m: wrist goes straight down 5.3", (w1 - Vector((0, 0, -5.3))).length < 1e-4, str(tuple(round(c, 5) for c in w1)))
check("handle 2.0 m: wrist goes straight down 2.0", (w3 - Vector((0, 0, -2.0))).length < 1e-4, str(tuple(round(c, 5) for c in w3)))
lvl = (E["Arm1_Head"].matrix.to_3x3() - ad.bones["Arm1_Head"].matrix_local.to_3x3())
check("head stays level", max(abs(x) for r in lvl for x in r) < 1e-5)
d2 = (E["Arm2_Head"].head - rest_head["Arm2_Head"]).length
check("arm 2 untouched", d2 < 1e-5, "%.2e m" % d2)
c = E["Arm1_Claw1"].matrix.to_3x3().col[1].angle(ad.bones["Arm1_Claw1"].matrix_local.to_3x3().col[1])
check("claws handle 25 deg opens claw 25 deg", abs(math.degrees(c) - 25.0) < 1e-3, "%.4f" % math.degrees(c))
# joints by hand
P["CTRL_Arm1_Hoist"].location.y = 0.0
P["CTRL_Arm1_Hoist"]["auto"] = 0.0
P["Arm1_Upper"].rotation_euler.x = math.radians(40)
rig.update_tag(); vl.update()
E = rig.evaluated_get(bpy.context.evaluated_depsgraph_get()).pose.bones
elbow = E["Arm1_Lower"].head - rest_head["Arm1_Upper"]
check("auto 0: shoulder turned 40 deg by hand moves the elbow down", abs(elbow.z + 3.0 * math.sin(math.radians(40))) < 1e-4, "elbow dz %.4f" % elbow.z)
# the ore hose follows: its last ring is on the head, so the handle at 5.3 m takes it 5.3 m down
P["Arm1_Upper"].rotation_euler.x = 0.0
P["CTRL_Arm1_Hoist"]["auto"] = 1.0
P["CTRL_Arm1_Hoist"].location.y = 5.3
rig.update_tag(); vl.update()
hose = bpy.data.objects["SK_DockArm_Arm1_OreHose"]
hose.update_tag(); vl.update()
eh = hose.evaluated_get(bpy.context.evaluated_depsgraph_get())
last = max(i for i, v in enumerate(hose.data.vertices) if hose.vertex_groups[v.groups[0].group].name == "Arm1_Head")
drop = (eh.data.vertices[last].co - hose.data.vertices[last].co)
check("ore hose end goes down with the head", (drop - Vector((0, 0, -5.3))).length < 1e-4, str(tuple(round(c, 4) for c in drop)))
mid = [i for i, v in enumerate(hose.data.vertices) if hose.vertex_groups[v.groups[0].group].name == "Arm1_HoseE5"][0]
check("ore hose bends at the elbow (fan bone E5 moved)", (eh.data.vertices[mid].co - hose.data.vertices[mid].co).length > 1.0)
P["CTRL_Arm1_Hoist"].location.y = 0.0

want_bones = {b.name: ((ROT @ b.head_local), b.parent.name if b.parent and b.parent.use_deform else None) for b in ad.bones if b.use_deform}
def wb(o):
    vs = [ROT @ (o.matrix_world @ v.co) for v in o.data.vertices]
    return [round(min(v[i] for v in vs), 3) for i in range(3)] + [round(max(v[i] for v in vs), 3) for i in range(3)]
want_mesh = {o.name: (wb(o), sorted(g.name for g in o.vertex_groups), len(o.data.polygons)) for o in meshes}

bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.ops.import_scene.fbx(filepath=os.path.join(OUT, "fbx", "SK_DockArm_Carousel.fbx"))
bpy.context.view_layer.update()
arms = [o for o in bpy.data.objects if o.type == 'ARMATURE']
check("FBX: one armature", len(arms) == 1)
a = arms[0]
got = {b.name: (a.matrix_world @ b.head_local, b.parent.name if b.parent else None) for b in a.data.bones}
check("FBX: the 86 deform bones and nothing else", sorted(got) == sorted(want_bones), str(sorted(set(got) ^ set(want_bones))))
worst = max((got[n][0] - want_bones[n][0]).length for n in want_bones if n in got)
check("FBX: bone positions", worst < 2e-4, "worst %.6f m" % worst)
check("FBX: bone parents", all(got[n][1] == want_bones[n][1] for n in want_bones if n in got))
fm = {o.name: o for o in bpy.data.objects if o.type == 'MESH'}
check("FBX: 26 skinned meshes", sorted(fm) == sorted(want_mesh), str(sorted(set(fm) ^ set(want_mesh))))
badm = []
for n, (b, grp, npoly) in want_mesh.items():
    o = fm.get(n)
    if o is None:
        continue
    vs = [o.matrix_world @ v.co for v in o.data.vertices]
    g = [round(min(v[i] for v in vs), 3) for i in range(3)] + [round(max(v[i] for v in vs), 3) for i in range(3)]
    if max(abs(x - y) for x, y in zip(g, b)) > 2e-3 or sorted(x.name for x in o.vertex_groups) != grp or len(o.data.polygons) != npoly:
        badm.append((n, g, b, [x.name for x in o.vertex_groups]))
check("FBX: mesh bounds, weights, face counts", not badm, str(badm[:2]))
print("RIGVERIFY", "ALL OK" if ok else "FAILED", "| armature scale in the FBX re-import:", tuple(round(s, 4) for s in a.scale))
if not ok:
    raise SystemExit(1)