"""Extract rig-agnostic evaluated-mesh evidence from a Blender scene.

Run only through Blender:
  blender -b --factory-startup -P evaluator_extract.py -- \
    scene.blend evidence.json evaluator_config.json reference|candidate

Object and bone names are retained only as diagnostics. Scoring correspondence
uses connected-component topology and normalized geometry descriptors.
"""

import bpy
import json
import math
import sys
from pathlib import Path
from mathutils import Vector
from bpy_extras.object_utils import world_to_camera_view


def args_after_dash():
    return sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []


def rounded(v, n=6):
    if isinstance(v, Vector):
        return [round(x, n) for x in v]
    return round(float(v), n)


def quantile(values, q):
    if not values:
        return 0.0
    s = sorted(values)
    return s[min(len(s) - 1, max(0, int(round(q * (len(s) - 1)))))]


def mesh_components(mesh):
    """Return geometric components after welding coincident split normals.

    The glTF uses flat-shaded duplicate vertices, so raw index connectivity
    would incorrectly call nearly every triangle a separate component.
    """
    def key(index):
        co = mesh.vertices[index].co
        return tuple(round(float(x), 6) for x in co)

    key_to_indices = {}
    for vertex in mesh.vertices:
        key_to_indices.setdefault(key(vertex.index), []).append(vertex.index)
    adjacency = {k: set() for k in key_to_indices}
    for poly in mesh.polygons:
        keys = [key(i) for i in poly.vertices]
        for i, a in enumerate(keys):
            b = keys[(i + 1) % len(keys)]
            adjacency[a].add(b)
            adjacency[b].add(a)
    seen = set()
    comps = []
    for seed in adjacency:
        if seed in seen:
            continue
        stack = [seed]
        seen.add(seed)
        coord_keys = []
        while stack:
            cur = stack.pop()
            coord_keys.append(cur)
            for nxt in adjacency.get(cur, ()):
                if nxt not in seen:
                    seen.add(nxt)
                    stack.append(nxt)
        verts = sorted(i for k in coord_keys for i in key_to_indices[k])
        vset = set(verts)
        poly_count = sum(
            1 for p in mesh.polygons if all(i in vset for i in p.vertices)
        )
        edge_indices = [
            e.index for e in mesh.edges
            if e.vertices[0] in vset and e.vertices[1] in vset
        ]
        comps.append((sorted(verts), edge_indices, poly_count))
    return comps


def component_descriptor(mesh, verts, edge_indices):
    pts = [mesh.vertices[i].co.copy() for i in verts]
    lo = Vector((min(p[k] for p in pts) for k in range(3)))
    hi = Vector((max(p[k] for p in pts) for k in range(3)))
    diag = max((hi - lo).length, 1e-9)
    dims = sorted((hi[k] - lo[k]) / diag for k in range(3))
    lengths = [
        (mesh.vertices[mesh.edges[i].vertices[0]].co -
         mesh.vertices[mesh.edges[i].vertices[1]].co).length / diag
        for i in edge_indices
    ]
    vset = set(verts)
    signed_volume = 0.0
    for poly in mesh.polygons:
        if len(poly.vertices) < 3 or not all(i in vset for i in poly.vertices):
            continue
        a = mesh.vertices[poly.vertices[0]].co
        for j in range(1, len(poly.vertices) - 1):
            b = mesh.vertices[poly.vertices[j]].co
            c = mesh.vertices[poly.vertices[j + 1]].co
            signed_volume += a.dot(b.cross(c)) / 6.0
    return {
        "vertex_count": len(verts),
        "edge_count": len(edge_indices),
        "bbox_ratios_sorted": [round(x, 6) for x in dims],
        "edge_length_quantiles": [
            round(quantile(lengths, q), 6) for q in (0.1, 0.25, 0.5, 0.75, 0.9)
        ],
        # Rotation/translation/scale invariant after normalization, but changes
        # sign under reflection. This preserves anatomical identity for mirrored
        # symmetric parts without consulting an object or bone name.
        "signed_volume_normalized": round(signed_volume / (diag ** 3), 8),
    }


def landmark_groups(mesh, verts):
    """Topology-stable regions selected from undeformed local coordinates."""
    pts = {i: mesh.vertices[i].co.copy() for i in verts}
    count = max(1, int(math.ceil(len(verts) * 0.10)))
    groups = {"center": list(verts)}
    for axis, label in enumerate(("x", "y", "z")):
        ordered = sorted(verts, key=lambda i: pts[i][axis])
        groups[f"{label}_low"] = ordered[:count]
        groups[f"{label}_high"] = ordered[-count:]
    return groups


def mean_points(points):
    if not points:
        return Vector((0, 0, 0))
    total = Vector((0, 0, 0))
    for p in points:
        total += p
    return total / len(points)


scene_path, output_path, config_path, role = args_after_dash()
config = json.loads(Path(config_path).read_text())
bpy.ops.wm.open_mainfile(filepath=str(Path(scene_path).resolve()))
scene = bpy.context.scene

source_to_id = {
    info["private_source_name"]: pid
    for pid, info in config["canonical_parts"].items()
}

# Detect a stage floor by scale-relative planar geometry, never by a name or an
# absolute dimension (uniformly scaled valid characters remain admissible).
all_meshes = [
    obj for obj in bpy.data.objects
    if obj.type == "MESH" and len(obj.data.vertices) and not obj.hide_render
]
sizes = [max(o.dimensions.x, o.dimensions.y, o.dimensions.z)
         for o in all_meshes]
median_size = sorted(sizes)[len(sizes) // 2] if sizes else 1.0
floor_candidates = [
    o for o in all_meshes
    if o.dimensions.z <= max(1e-4, 0.02 * max(o.dimensions.x, o.dimensions.y))
    and max(o.dimensions.x, o.dimensions.y) >= 2.0 * median_size
]
floor_object = max(
    floor_candidates, key=lambda o: o.dimensions.x * o.dimensions.y,
    default=None)
meshes = [o for o in all_meshes if o != floor_object]
meshes.sort(key=lambda obj: obj.name)

component_defs = []
for obj in meshes:
    base = obj.data
    for ordinal, (verts, edges, poly_count) in enumerate(mesh_components(base)):
        desc = component_descriptor(base, verts, edges)
        desc["polygon_count"] = poly_count
        cid = f"c{len(component_defs):04d}"
        mats = []
        for mat in obj.data.materials:
            if mat:
                color = list(mat.diffuse_color[:3])
                mats.append({"diagnostic_name": mat.name,
                             "rgb": [round(x, 6) for x in color]})
        component_defs.append({
            "id": cid,
            "diagnostic_object_name": obj.name,
            "canonical_part": source_to_id.get(obj.name) if role == "reference" else None,
            "object": obj,
            "vertex_indices": verts,
            "landmark_groups": landmark_groups(base, verts),
            "descriptor": desc,
            "materials": mats,
        })

frames = {}
camera_visibility = {}
topology_warnings = []
start = int(scene.frame_start)
end = int(scene.frame_end)
for frame in range(start, end + 1):
    scene.frame_set(frame)
    bpy.context.view_layer.update()
    depsgraph = bpy.context.evaluated_depsgraph_get()
    per_frame = {}
    all_world_points = []
    eval_cache = {}
    try:
        for comp in component_defs:
            obj = comp["object"]
            if obj.name not in eval_cache:
                ev = obj.evaluated_get(depsgraph)
                em = ev.to_mesh()
                eval_cache[obj.name] = (ev, em)
            ev, em = eval_cache[obj.name]
            inds = comp["vertex_indices"]
            if not inds or max(inds) >= len(em.vertices):
                topology_warnings.append(
                    f"frame {frame}: evaluated topology changed for {obj.name}"
                )
                continue
            world = [ev.matrix_world @ em.vertices[i].co for i in inds]
            all_world_points.extend(world)
            lo = [min(p[k] for p in world) for k in range(3)]
            hi = [max(p[k] for p in world) for k in range(3)]
            landmarks = {}
            for name, group in comp["landmark_groups"].items():
                landmarks[name] = rounded(mean_points(
                    [ev.matrix_world @ em.vertices[i].co for i in group]
                ))
            z_band = max((hi[2] - lo[2]) * 0.02, 1e-5)
            landmarks["dynamic_sole"] = rounded(mean_points(
                [p for p in world if p.z <= lo[2] + z_band]
            ))
            per_frame[comp["id"]] = {
                "centroid": rounded(mean_points(world)),
                "bbox_min": rounded(Vector(lo)),
                "bbox_max": rounded(Vector(hi)),
                "landmarks": landmarks,
            }
            if scene.camera:
                projected_component = [
                    world_to_camera_view(scene, scene.camera, p) for p in world
                ]
                inside_component = [
                    p for p in projected_component
                    if p.z > 0 and 0 <= p.x <= 1 and 0 <= p.y <= 1
                ]
                per_frame[comp["id"]]["camera_visible_fraction"] = round(
                    len(inside_component) / len(projected_component), 6)
            else:
                per_frame[comp["id"]]["camera_visible_fraction"] = 0.0
    finally:
        for ev, _em in eval_cache.values():
            ev.to_mesh_clear()
    frames[str(frame)] = per_frame
    if scene.camera and all_world_points:
        projected = [world_to_camera_view(scene, scene.camera, p)
                     for p in all_world_points]
        inside = [p for p in projected
                  if p.z > 0 and 0 <= p.x <= 1 and 0 <= p.y <= 1]
        camera_visibility[str(frame)] = {
            "visible_fraction": round(len(inside) / len(projected), 6),
            "projected_bounds": [
                round(min(p.x for p in projected), 6),
                round(min(p.y for p in projected), 6),
                round(max(p.x for p in projected), 6),
                round(max(p.y for p in projected), 6),
            ],
        }
    else:
        camera_visibility[str(frame)] = {
            "visible_fraction": 0.0, "projected_bounds": None
        }

components_json = []
for comp in component_defs:
    components_json.append({
        k: v for k, v in comp.items()
        if k not in ("object", "vertex_indices", "landmark_groups")
    })

armatures = [o for o in bpy.data.objects if o.type == "ARMATURE"]
animated_armatures = []
for armature in armatures:
    animation = armature.animation_data
    if animation and (
        animation.action
        or len(animation.nla_tracks)
        or len(animation.drivers)
    ):
        animated_armatures.append(armature)

armature_driven_objects = set()
for obj in meshes:
    modifier_driven = any(
        modifier.type == "ARMATURE"
        and modifier.object in armatures
        for modifier in obj.modifiers
    )
    bone_parented = (
        obj.parent in armatures and obj.parent_type == "BONE"
    )
    if modifier_driven or bone_parented:
        armature_driven_objects.add(obj)

armature_driven_component_ids = [
    comp["id"] for comp in component_defs
    if comp["object"] in armature_driven_objects
]
evidence = {
    "schema_version": 1,
    "scene": str(Path(scene_path).resolve()),
    "role": role,
    "fresh_process_open": True,
    "blender_version": bpy.app.version_string,
    "frame_start": start,
    "frame_end": end,
    "fps": scene.render.fps / max(scene.render.fps_base, 1e-9),
    "resolution": [scene.render.resolution_x, scene.render.resolution_y],
    "mesh_object_count": len(meshes),
    "detected_floor": floor_object.name if floor_object else None,
    "detected_floor_z": (
        round(floor_object.matrix_world.translation.z, 6)
        if floor_object else None),
    "armature_count": len(armatures),
    "animated_armature_count": len(animated_armatures),
    "armature_driven_mesh_object_count": len(armature_driven_objects),
    "armature_driven_component_ids": armature_driven_component_ids,
    "diagnostic_armature_names": [o.name for o in armatures],
    "component_count": len(component_defs),
    "components": components_json,
    "topology_warnings": sorted(set(topology_warnings)),
    "active_camera": scene.camera.name if scene.camera else None,
    "camera_visibility": camera_visibility,
    "frames": frames,
}
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
Path(output_path).write_text(json.dumps(evidence, indent=1))
print(f"[extract] {len(meshes)} meshes, {len(component_defs)} components, "
      f"{len(frames)} frames -> {output_path}")