"""Audit non-neighbor triangle intersections on every frame of a skinned GLB."""

import argparse
from collections import Counter
import json
from pathlib import Path
import sys
import time

import bpy
from mathutils.bvhtree import BVHTree
import numpy as np

p = argparse.ArgumentParser()
p.add_argument("--input", required=True)
p.add_argument("--output", required=True)
p.add_argument("--frames", type=int, default=180)
p.add_argument("--corrective", type=int, default=0)
p.add_argument("--dual-quaternion", action="store_true")
p.add_argument("--weld", action="store_true")
a = p.parse_args(sys.argv[sys.argv.index("--") + 1 :])
out = Path(a.output)
out.mkdir(parents=True, exist_ok=True)
bpy.ops.object.select_all(action="SELECT")
bpy.ops.object.delete(use_global=False)
bpy.context.scene.render.fps = 30
if Path(a.input).suffix == ".blend":
    bpy.ops.wm.open_mainfile(filepath=str(Path(a.input).resolve()))
else:
    bpy.ops.import_scene.gltf(filepath=str(Path(a.input).resolve()))
objects = [
    o
    for o in bpy.context.scene.objects
    if o.type == "MESH" and any(m.type == "ARMATURE" for m in o.modifiers)
]
if a.weld:
    bpy.ops.object.select_all(action="DESELECT")
    for obj in objects:
        obj.select_set(True)
    bpy.context.view_layer.objects.active = objects[0]
    bpy.ops.object.join()
    bpy.ops.object.mode_set(mode="EDIT")
    bpy.ops.mesh.select_all(action="SELECT")
    bpy.ops.mesh.remove_doubles(threshold=1e-5)
    bpy.ops.object.mode_set(mode="OBJECT")
    objects = [bpy.context.object]
if a.dual_quaternion:
    for obj in objects:
        for modifier in obj.modifiers:
            if modifier.type == "ARMATURE":
                modifier.use_deform_preserve_volume = True
if a.corrective:
    for obj in objects:
        m = obj.modifiers.new("Collision deformation smoothing", "CORRECTIVE_SMOOTH")
        m.factor = 1.0
        m.iterations = a.corrective
        m.use_pin_boundary = False
        m.rest_source = "ORCO"
rig = next(o for o in bpy.context.scene.objects if o.type == "ARMATURE")
rig.data.pose_position = "REST"
for obj in objects:
    if obj.data.shape_keys:
        for key in obj.data.shape_keys.key_blocks:
            key.value = 0.0
bpy.context.view_layer.update()


def geometry():
    dg = bpy.context.evaluated_depsgraph_get()
    vv = []
    ff = []
    offset = 0
    for obj in objects:
        evaluated = obj.evaluated_get(dg)
        mesh = evaluated.to_mesh()
        mesh.calc_loop_triangles()
        vertices = np.array([tuple(evaluated.matrix_world @ v.co) for v in mesh.vertices])
        faces = np.array([tuple(t.vertices) for t in mesh.loop_triangles], dtype=np.int32)
        vv.append(vertices)
        ff.append(faces + offset)
        offset += len(vertices)
        evaluated.to_mesh_clear()
    return np.concatenate(vv), np.concatenate(ff)


rest, faces = geometry()
# Weld only for adjacency detection; retain all original export vertices and triangles.
_, weld = np.unique(np.round(rest / 1e-5).astype(np.int64), axis=0, return_inverse=True)
weld_faces = weld[faces]
regions = []
for obj in objects:
    lookup = {g.index: g.name for g in obj.vertex_groups}
    for v in obj.data.vertices:
        name = lookup[max(v.groups, key=lambda g: g.weight).group]
        if name in ["Head", "Neck1", "Neck2"]:
            region = "head"
        elif name in ["Hips", "Spine1", "Spine2", "Chest"]:
            region = "torso"
        else:
            region = name
        regions.append(region)
regions = np.array(regions)
face_regions = np.array([Counter(regions[f]).most_common(1)[0][0] for f in faces])


def intersections(vertices):
    tree = BVHTree.FromPolygons(vertices.tolist(), faces.tolist(), all_triangles=True, epsilon=0)
    pairs = np.array(tree.overlap(tree), dtype=np.int32)
    if not len(pairs):
        return set(), {}
    pairs = pairs[pairs[:, 0] < pairs[:, 1]]
    shared = (weld_faces[pairs[:, 0], :, None] == weld_faces[pairs[:, 1], None, :]).any(axis=(1, 2))
    pairs = pairs[~shared]
    # Strict triangle separating-axis test rejects AABB-only and coplanar near misses.
    if len(pairs):
        A = vertices[faces[pairs[:, 0]]]
        B = vertices[faces[pairs[:, 1]]]
        ea = np.roll(A, -1, axis=1) - A
        eb = np.roll(B, -1, axis=1) - B
        na = np.cross(ea[:, 0], ea[:, 1])
        nb = np.cross(eb[:, 0], eb[:, 1])
        axes = np.concatenate(
            [
                na[:, None],
                nb[:, None],
                np.cross(ea[:, :, None], eb[:, None, :]).reshape(-1, 9, 3),
                np.cross(na[:, None], ea),
                np.cross(nb[:, None], eb),
            ],
            axis=1,
        )
        lengths = np.linalg.norm(axes, axis=-1)
        axes = axes / np.maximum(lengths[:, :, None], 1e-20)
        pa = np.einsum("nti,nai->nat", A, axes)
        pb = np.einsum("nti,nai->nat", B, axes)
        separated = ((pa.max(2) < pb.min(2) - 1e-7) | (pb.max(2) < pa.min(2) - 1e-7)) & (lengths > 1e-10)
        # Count actual crossings, not triangles merely touching within float precision.
        una = na / np.maximum(np.linalg.norm(na, axis=1)[:, None], 1e-20)
        unb = nb / np.maximum(np.linalg.norm(nb, axis=1)[:, None], 1e-20)
        da = np.einsum("nti,ni->nt", A - B[:, 0, None, :], unb)
        db = np.einsum("nti,ni->nt", B - A[:, 0, None, :], una)
        crossing = (da.min(1) < -1e-6) & (da.max(1) > 1e-6) & (db.min(1) < -1e-6) & (db.max(1) > 1e-6)
        pairs = pairs[(~separated.any(1)) & crossing]
    groups = Counter("|".join(sorted((face_regions[i], face_regions[j]))) for i, j in pairs)
    return {int(i) * len(faces) + int(j) for i, j in pairs}, dict(groups)


def crossing_severity(vertices, encoded_pairs):
    if not encoded_pairs:
        return {
            "max_triangle_plane_depth": 0.0,
            "p95_triangle_plane_depth": 0.0,
            "affected_surface_area": 0.0,
            "affected_surface_fraction": 0.0,
        }
    pairs = np.array([(key // len(faces), key % len(faces)) for key in encoded_pairs], dtype=int)
    triangles = vertices[faces]
    normals = np.cross(triangles[:, 1] - triangles[:, 0], triangles[:, 2] - triangles[:, 0])
    lengths = np.linalg.norm(normals, axis=1)
    unit = normals / np.maximum(lengths[:, None], 1e-20)
    first, second = pairs.T
    da = np.einsum("nti,ni->nt", triangles[first] - triangles[second, 0, None], unit[second])
    db = np.einsum("nti,ni->nt", triangles[second] - triangles[first, 0, None], unit[first])
    depth = np.minimum.reduce([da.max(1), -da.min(1), db.max(1), -db.min(1)])
    affected = float(lengths[np.unique(pairs)].sum() * 0.5)
    return {
        "max_triangle_plane_depth": float(depth.max()),
        "p95_triangle_plane_depth": float(np.quantile(depth, 0.95)),
        "affected_surface_area": affected,
        "affected_surface_fraction": affected / max(float(lengths.sum() * 0.5), 1e-20),
    }


baseline, baseline_groups = intersections(rest)
print("REST", len(baseline), baseline_groups, flush=True)
rig.data.pose_position = "POSE"
rows = []
started = time.time()
worst = {}
saved = {}
joint_tracks = []
for frame in range(1, a.frames + 1):
    bpy.context.scene.frame_set(frame)
    vertices, _ = geometry()
    pairs, groups = intersections(vertices)
    new = pairs - baseline
    newgroups = Counter(
        "|".join(sorted((face_regions[k // len(faces)], face_regions[k % len(faces)]))) for k in new
    )
    feet = {
        side: float(
            vertices[
                (rest[:, 2] < rest[:, 2].min() + 0.09)
                & ((rest[:, 0] > 0) if side == "Left" else (rest[:, 0] < 0)),
                2,
            ].min()
        )
        for side in ["Left", "Right"]
    }
    row = {
        "frame": frame,
        "intersecting_triangle_pairs": len(pairs),
        "new_pairs_vs_rest": len(new),
        "groups": groups,
        "new_groups": dict(newgroups),
        "lowest_foot_z": feet,
        "lowest_vertex_z": float(vertices[:, 2].min()),
        "new_collision_severity": crossing_severity(vertices, new),
    }
    rows.append(row)
    joint_tracks.append(
        {
            n: list(rig.matrix_world @ rig.pose.bones[n].head)
            for n in ["LeftHand", "RightHand", "LeftFoot", "RightFoot", "Head"]
        }
    )
    if frame == 1 or len(new) > worst.get("new_pairs_vs_rest", -1):
        worst = row
        saved = {
            "vertices": vertices.copy(),
            "pairs": np.array([(k // len(faces), k % len(faces)) for k in sorted(new)], dtype=np.int32),
        }
    if frame % 15 == 0:
        print(
            frame,
            "pairs",
            len(pairs),
            "new",
            len(new),
            "elapsed",
            round(time.time() - started, 1),
            flush=True,
        )
    if frame % 30 == 0:
        (out / "progress.json").write_text(json.dumps(rows, indent=2))
np.savez_compressed(out / "worst.npz", faces=faces, **saved)
summary = {
    "input": str(Path(a.input).resolve()),
    "frames": a.frames,
    "triangles": len(faces),
    "rest_intersections": len(baseline),
    "rest_groups": baseline_groups,
    "frames_with_new_intersections": sum(r["new_pairs_vs_rest"] > 0 for r in rows),
    "peak_new_pairs": max(r["new_pairs_vs_rest"] for r in rows),
    "total_new_pair_frames": sum(r["new_pairs_vs_rest"] for r in rows),
    "worst_frame": worst,
    "ground_min_z": min(min(r["lowest_foot_z"].values()) for r in rows),
    "ground_min_all_vertices_z": min(r["lowest_vertex_z"] for r in rows),
    "peak_triangle_plane_depth": max(r["new_collision_severity"]["max_triangle_plane_depth"] for r in rows),
    "peak_affected_surface_fraction": max(
        r["new_collision_severity"]["affected_surface_fraction"] for r in rows
    ),
    "severity_method": "For each crossing triangle pair, minimum of both positive and negative vertex distances to the opposite triangle plane. This is a local crossing-depth estimate, not volumetric penetration. Affected area counts unique triangles once; dimensions use model units.",
    "method": "Deformed full-resolution triangles; BVH broad phase and triangle separating-axis narrow phase. Shared welded-vertex pairs excluded. Requires triangle-plane crossings deeper than 1e-6 model units; coplanar contact is excluded. Existing rest-pose intersections reported separately. Discrete 30 fps samples, not continuous collision detection.",
}
(out / "audit.json").write_text(
    json.dumps({"summary": summary, "frames": rows, "joints": joint_tracks}, indent=2)
)
print(json.dumps(summary), flush=True)