"""Inspect one connected skinned-mesh component at bind pose and one action sample.

Usage:
  blender --background --python diagnose-animated-component.py -- \
    input.glb output.json mesh_name component_index "Action - name" middle
"""

import json
import math
import os
import sys

import bpy
from mathutils import Vector


def arguments():
    values = sys.argv[sys.argv.index("--") + 1 :]
    if len(values) != 6:
        raise RuntimeError(
            "Expected input, output, mesh name, component index, action name, frame label."
        )
    return values[0], values[1], values[2], int(values[3]), values[4], values[5]


class UnionFind:
    def __init__(self, count):
        self.parent = list(range(count))
        self.rank = [0] * count

    def find(self, value):
        root = value
        while self.parent[root] != root:
            root = self.parent[root]
        while self.parent[value] != value:
            next_value = self.parent[value]
            self.parent[value] = root
            value = next_value
        return root

    def union(self, left, right):
        left_root = self.find(left)
        right_root = self.find(right)
        if left_root == right_root:
            return
        if self.rank[left_root] < self.rank[right_root]:
            left_root, right_root = right_root, left_root
        self.parent[right_root] = left_root
        if self.rank[left_root] == self.rank[right_root]:
            self.rank[left_root] += 1


def mesh_components(mesh_object):
    mesh = mesh_object.data
    union_find = UnionFind(len(mesh.vertices))
    for edge in mesh.edges:
        union_find.union(edge.vertices[0], edge.vertices[1])
    grouped = {}
    for index in range(len(mesh.vertices)):
        grouped.setdefault(union_find.find(index), []).append(index)
    return sorted(grouped.values(), key=lambda indices: (indices[0], len(indices)))


def rounded_vector(value):
    return [round(float(axis), 6) for axis in value]


def evaluated_positions(mesh_object):
    graph = bpy.context.evaluated_depsgraph_get()
    evaluated = mesh_object.evaluated_get(graph)
    mesh = evaluated.to_mesh()
    try:
        return [evaluated.matrix_world @ vertex.co for vertex in mesh.vertices]
    finally:
        evaluated.to_mesh_clear()


def component_metrics(points):
    centroid = sum(points, Vector()) / len(points)
    minimum = Vector(tuple(min(point[axis] for point in points) for axis in range(3)))
    maximum = Vector(tuple(max(point[axis] for point in points) for axis in range(3)))
    extents = maximum - minimum
    return {
        "centroid": rounded_vector(centroid),
        "minimum": rounded_vector(minimum),
        "maximum": rounded_vector(maximum),
        "diagonal": round(float(extents.length), 6),
    }


def set_bind(scene, armatures):
    for armature in armatures:
        if armature.animation_data is not None:
            armature.animation_data.action = None
        armature.data.pose_position = "REST"
    scene.frame_set(0)
    bpy.context.view_layer.update()


def set_action(scene, armatures, action, frame):
    for index, armature in enumerate(armatures):
        armature.data.pose_position = "POSE"
        if armature.animation_data is None:
            armature.animation_data_create()
        armature.animation_data.action = action if index == 0 else None
    integer = math.floor(frame)
    scene.frame_set(integer, subframe=frame - integer)
    bpy.context.view_layer.update()


input_file, output_file, mesh_name, component_index, action_name, frame_label = arguments()
bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.ops.import_scene.gltf(filepath=os.path.abspath(input_file))
scene = bpy.context.scene
mesh_objects = [obj for obj in scene.objects if obj.type == "MESH"]
armatures = [obj for obj in scene.objects if obj.type == "ARMATURE"]
mesh_object = next((obj for obj in mesh_objects if obj.name == mesh_name), None)
if mesh_object is None:
    raise RuntimeError("Mesh object not found: " + mesh_name)
components = mesh_components(mesh_object)
if component_index < 0 or component_index >= len(components):
    raise RuntimeError("Component index is outside topology.")
indices = components[component_index]
index_set = set(indices)
faces = [
    list(polygon.vertices)
    for polygon in mesh_object.data.polygons
    if all(vertex in index_set for vertex in polygon.vertices)
]

set_bind(scene, armatures)
bind_positions = evaluated_positions(mesh_object)
bind_points = [bind_positions[index] for index in indices]

action = next((candidate for candidate in bpy.data.actions if candidate.name == action_name), None)
if action is None:
    raise RuntimeError("Action not found: " + action_name)
first, last = (float(value) for value in action.frame_range)
frame_by_label = {
    "start": first,
    "middle": (first + last) * 0.5,
    "end": last,
}
if frame_label not in frame_by_label:
    raise RuntimeError("Frame label must be start, middle, or end.")
frame = frame_by_label[frame_label]
set_action(scene, armatures, action, frame)
posed_positions = evaluated_positions(mesh_object)
posed_points = [posed_positions[index] for index in indices]

vertex_rows = []
bone_totals = {}
for vertex_index in indices:
    vertex = mesh_object.data.vertices[vertex_index]
    influences = []
    for membership in sorted(vertex.groups, key=lambda entry: -entry.weight):
        group_name = mesh_object.vertex_groups[membership.group].name
        weight = float(membership.weight)
        influences.append({"bone": group_name, "weight": round(weight, 8)})
        row = bone_totals.setdefault(
            group_name,
            {"bone": group_name, "vertices": 0, "weightSum": 0.0, "maximumWeight": 0.0},
        )
        row["vertices"] += 1
        row["weightSum"] += weight
        row["maximumWeight"] = max(row["maximumWeight"], weight)
    vertex_rows.append(
        {
            "index": vertex_index,
            "bind": rounded_vector(bind_positions[vertex_index]),
            "sample": rounded_vector(posed_positions[vertex_index]),
            "influences": influences,
        }
    )
for row in bone_totals.values():
    row["weightSum"] = round(row["weightSum"], 8)
    row["maximumWeight"] = round(row["maximumWeight"], 8)

bind_metrics = component_metrics(bind_points)
sample_metrics = component_metrics(posed_points)
bind_centroid = Vector(bind_metrics["centroid"])
sample_centroid = Vector(sample_metrics["centroid"])
sorted_bone_totals = sorted(
    bone_totals.values(),
    key=lambda row: (-row["weightSum"], row["bone"].casefold()),
)
primary_bone = sorted_bone_totals[0]["bone"]
primary_group = mesh_object.vertex_groups.get(primary_bone)
components_sharing_primary_bone = []
if primary_group is not None:
    for candidate_index, candidate_vertices in enumerate(components):
        weight_sum = 0.0
        maximum_weight = 0.0
        weighted_vertices = 0
        for vertex_index in candidate_vertices:
            for membership in mesh_object.data.vertices[vertex_index].groups:
                if membership.group != primary_group.index:
                    continue
                weight_sum += float(membership.weight)
                maximum_weight = max(maximum_weight, float(membership.weight))
                weighted_vertices += 1
        if weight_sum <= 0:
            continue
        candidate_bind = component_metrics(
            [bind_positions[index] for index in candidate_vertices]
        )
        candidate_posed = component_metrics(
            [posed_positions[index] for index in candidate_vertices]
        )
        components_sharing_primary_bone.append(
            {
                "componentIndex": candidate_index,
                "vertices": len(candidate_vertices),
                "weightedVertices": weighted_vertices,
                "weightSum": round(weight_sum, 8),
                "maximumWeight": round(maximum_weight, 8),
                "bindCentroid": candidate_bind["centroid"],
                "sampleCentroid": candidate_posed["centroid"],
                "centroidShift": round(
                    float(
                        (
                            Vector(candidate_posed["centroid"])
                            - Vector(candidate_bind["centroid"])
                        ).length
                    ),
                    6,
                ),
                "diagonalRatio": round(
                    candidate_posed["diagonal"] / candidate_bind["diagonal"],
                    6,
                ) if candidate_bind["diagonal"] else None,
            }
        )
result = {
    "schemaVersion": 1,
    "input": os.path.abspath(input_file),
    "mesh": mesh_name,
    "componentIndex": component_index,
    "sourceTopology": {
        "meshVertices": len(mesh_object.data.vertices),
        "meshEdges": len(mesh_object.data.edges),
        "meshFaces": len(mesh_object.data.polygons),
        "connectedComponents": len(components),
        "componentVertices": len(indices),
        "componentFaces": len(faces),
        "vertexIndices": indices,
        "faces": faces,
    },
    "sample": {
        "action": action_name,
        "frameLabel": frame_label,
        "frame": round(frame, 6),
        "range": [round(first, 6), round(last, 6)],
    },
    "bind": bind_metrics,
    "posed": sample_metrics,
    "shape": {
        "diagonalRatio": round(sample_metrics["diagonal"] / bind_metrics["diagonal"], 6),
        "centroidShift": round(float((sample_centroid - bind_centroid).length), 6),
    },
    "boneInfluences": sorted_bone_totals,
    "primaryBone": primary_bone,
    "componentsSharingPrimaryBone": components_sharing_primary_bone,
    "vertices": vertex_rows,
}
os.makedirs(os.path.dirname(os.path.abspath(output_file)), exist_ok=True)
with open(output_file, "w", encoding="utf-8") as stream:
    json.dump(result, stream, indent=2)
    stream.write("\n")
print(json.dumps({
    "output": os.path.abspath(output_file),
    "componentVertices": len(indices),
    "componentFaces": len(faces),
    "boneInfluences": result["boneInfluences"],
    "diagonalRatio": result["shape"]["diagonalRatio"],
    "centroidShift": result["shape"]["centroidShift"],
}))