"""Offline collision-from-visuals cooking for the complete LeKiwi arm subtree.

Run in build/venvs/collision, NOT the training or LeRobot environment. CoACD is
never imported by the browser/bridge. Each convex part is a separate MJCF mesh.
"""

import argparse
import hashlib
import importlib.metadata
import json
import os
import xml.etree.ElementTree as ET
from pathlib import Path

# Fix the numerical execution environment before importing numpy/CoACD.
os.environ["OMP_NUM_THREADS"] = "4"
os.environ["OPENBLAS_NUM_THREADS"] = "1"

ROOT = Path(__file__).resolve().parents[2]
CONFIG = ROOT / "robot_profiles/lekiwi-collision-source.json"
VERSIONS = {"coacd": "1.0.14", "trimesh": "4.12.2", "numpy": "2.1.3", "scipy": "1.17.0"}
PARAMETERS = {
    "threshold": 0.003,
    "real_metric": True,
    "preprocess_mode": "auto",
    "preprocess_resolution": 60,
    "resolution": 1000,
    "mcts_nodes": 12,
    "mcts_iterations": 60,
    "mcts_max_depth": 3,
    "decimate": True,
    "max_ch_vertex": 96,
    "seed": 42,
}


def digest(data):
    return hashlib.sha256(data).hexdigest()


def visuals(source, config):
    """Select every visual in the subtree, including welded motors/accessories."""
    urdf = (source / "LeKiwi.urdf").read_bytes()
    if digest(urdf) != config["urdfSha256"]:
        raise ValueError("URDF 与固定版本不符")
    robot = ET.fromstring(urdf)
    parents = {
        j.find("child").get("link"): j.find("parent").get("link") for j in robot.findall("joint")
    }
    result = []
    for link in robot.findall("link"):
        ancestor = link.get("name")
        visited = set()
        while ancestor in parents and ancestor != config["rootLink"]:
            if ancestor in visited:
                raise ValueError("URDF 运动链存在环")
            visited.add(ancestor)
            ancestor = parents[ancestor]
        if ancestor != config["rootLink"]:
            continue
        for visual in link.findall("visual"):
            mesh = visual.find("geometry/mesh")
            if mesh is None:
                raise ValueError("当前固定资产只支持 STL 视觉网格，不能静默遗漏几何")
            filename = mesh.get("filename")
            path = (source / filename).resolve()
            if not path.is_relative_to(source.resolve()) or filename not in config["meshes"]:
                raise ValueError(f"未批准的源网格：{filename}")
            sha = digest(path.read_bytes())
            if sha != config["meshes"][filename]:
                raise ValueError(f"源网格 SHA-256 不匹配：{filename}")
            scale = [float(x) for x in mesh.get("scale", "1 1 1").split()]
            if len(scale) != 3 or not all(0 < x < float("inf") for x in scale):
                raise ValueError("无效的视觉网格缩放")
            result.append(
                {
                    "link": link.get("name"),
                    "visual": visual.get("name"),
                    "mesh": filename,
                    "sha256": sha,
                    "scale": scale,
                }
            )
    if not result or {v["mesh"] for v in result} != set(config["meshes"]):
        raise ValueError("整臂视觉覆盖清单不完整")
    return result


def swept_bounds(vertices, axis):
    """Conservative axial/radial intervals for a full revolution, not pose sampling."""
    import numpy as np
    from scipy.spatial import ConvexHull

    axial = vertices @ axis
    reference = np.eye(3)[int(np.argmin(np.abs(axis)))]
    x = np.cross(axis, reference)
    x /= np.linalg.norm(x)
    y = np.cross(axis, x)
    projected = vertices @ np.column_stack([x, y])
    hull = ConvexHull(projected)
    radial_min = 0.0
    if hull.equations[:, 2].max() > 0:
        a = projected[hull.vertices]
        b = np.roll(a, -1, axis=0)
        edge = b - a
        t = np.clip(-(a * edge).sum(axis=1) / (edge * edge).sum(axis=1), 0, 1)
        radial_min = float(np.linalg.norm(a + t[:, None] * edge, axis=1).min())
    return (
        float(axial.min()),
        float(axial.max()),
        radial_min,
        float(np.linalg.norm(projected, axis=1).max()),
    )


def intervals_overlap(a, b, margin=0.0001):
    # 0.1 mm conservative allowance for serialized/compiled frame precision.
    return (
        max(a[0], b[0]) <= min(a[1], b[1]) + margin and max(a[2], b[2]) <= min(a[3], b[3]) + margin
    )


def assembly_frames(source, config, coverage):
    """URDF joint/visual frames for both geometric partitioning and pair policy."""
    import numpy as np
    from scipy.spatial.transform import Rotation

    robot = ET.parse(source / "LeKiwi.urdf").getroot()
    joints = {j.get("name"): j for j in robot.findall("joint")}
    parents = {j.find("child").get("link"): j for j in joints.values()}
    links = {link.get("name"): link for link in robot.findall("link")}
    poses = {}

    def origin(node):
        result = np.eye(4)
        if node is not None:
            result[:3, 3] = [float(x) for x in node.get("xyz", "0 0 0").split()]
            result[:3, :3] = Rotation.from_euler(
                "xyz", [float(x) for x in node.get("rpy", "0 0 0").split()]
            ).as_matrix()
        return result

    def pose(link):
        if link not in poses:
            j = parents.get(link)
            poses[link] = (
                np.eye(4)
                if j is None
                else pose(j.find("parent").get("link")) @ origin(j.find("origin"))
            )
        return poses[link]

    def owner(link):
        if link == config["rootLink"]:
            return "base"
        j = parents[link]
        return owner(j.find("parent").get("link")) if j.get("type") == "fixed" else j.get("name")

    by_visual = {}
    for item in coverage:
        item["weldJoint"] = owner(item["link"])
        visual = next(
            v for v in links[item["link"]].findall("visual") if v.get("name") == item["visual"]
        )
        by_visual[item["visual"]] = (
            item["weldJoint"],
            pose(item["link"]) @ origin(visual.find("origin")),
        )
    result = []
    for name, core in config["assemblyCores"].items():
        if not (0 < core["radiusM"] <= 0.025 and 0 < core["axialHalfExtentM"] <= 0.06):
            raise ValueError("装配配合区域超出已审核尺寸预算")
        j = joints[name]
        parent = owner(j.find("parent").get("link"))
        transform = pose(j.find("child").get("link"))
        axis = transform[:3, :3] @ np.array([float(x) for x in j.find("axis").get("xyz").split()])
        axis /= np.linalg.norm(axis)
        result.append((name, parent, transform, axis, core))
    return by_visual, result


def joint_policy(source, config, parts, coverage):
    """Bound exceptions by joint-local geometry, never by a whole body pair."""
    import numpy as np

    by_visual, frames = assembly_frames(source, config, coverage)
    result = []
    for name, parent, transform, axis, core in frames:
        members = []
        groups = {parent: {}, name: {}}
        for part in parts:
            group, visual_pose = by_visual[part["visual"]]
            if group not in (parent, name):
                continue
            vertices = np.fromstring(part["vertices"], sep=" ").reshape(-1, 3)
            offsets = vertices @ visual_pose[:3, :3].T + visual_pose[:3, 3] - transform[:3, 3]
            bounds = swept_bounds(offsets, axis)
            groups[group][part["name"]] = bounds
            if (
                bounds[3] <= core["radiusM"]
                and max(abs(bounds[0]), abs(bounds[1])) <= core["axialHalfExtentM"]
            ):
                members.append(part["name"])
        core_set = set(members)
        pairs = []
        filtered = pruned = 0
        for a, bound_a in groups[parent].items():
            for b, bound_b in groups[name].items():
                if a in core_set or b in core_set:
                    filtered += 1
                elif not intervals_overlap(bound_a, bound_b):
                    pruned += 1
                else:
                    pairs.append([a, b])
        if not pairs:
            raise ValueError(f"不能关闭整对相邻连杆碰撞：{name}")
        result.append(
            {
                "joint": name,
                "parent": parent,
                **core,
                "coreParts": members,
                "pairs": pairs,
                "coreFilteredPairs": filtered,
                "sweptPrunedPairs": pruned,
            }
        )
    return result


def clip_convex(vertices, normal, offset):
    """Intersect a convex hull with a half-space, retaining all crossing edges."""
    import numpy as np
    from scipy.spatial import ConvexHull, QhullError

    signed = vertices @ normal - offset
    if signed.max() <= 1e-10:
        return vertices
    if signed.min() >= -1e-10:
        return None
    hull = ConvexHull(vertices)
    result = list(vertices[signed <= 0])
    for triangle in hull.simplices:
        for i in range(3):
            a, b = triangle[i], triangle[(i + 1) % 3]
            if signed[a] * signed[b] < 0:
                result.append(
                    vertices[a] + (vertices[b] - vertices[a]) * signed[a] / (signed[a] - signed[b])
                )
    points = np.unique(np.asarray(result), axis=0)
    if len(points) < 4:
        return None
    try:
        # Remove edge/triangulation interpolation points BEFORE rounding;
        # rounding collinear points first invents tiny zigzag facet vertices.
        points = np.unique(np.round(points[ConvexHull(points).vertices], 6), axis=0)
        if len(points) < 4:
            return None
        hull = ConvexHull(points)
    except QhullError:
        return None  # sub-micrometre cutting sliver collapsed by serialization
    if hull.volume <= 1e-15:
        return None
    return points[hull.vertices]


def split_core(vertices, planes):
    """Partition, not delete: core and every outside fragment retain external contact."""
    result = []
    pending = vertices
    for normal, offset in planes:
        outside = clip_convex(pending, -normal, -offset)
        if outside is not None:
            result.append(outside)
        pending = clip_convex(pending, normal, offset)
        if pending is None:
            return result
    result.append(pending)
    return result


def partition_assembly_cores(source, config, parts, coverage):
    """Split straddling hulls so a bearing contact cannot lock a structural hull.

    A 16-sided inscribed prism stays inside each reviewed cylinder. The 2um
    inset keeps rounding from reclassifying its core as structural. No region
    is removed; only whole resulting core parts receive local mating exceptions.
    """
    import numpy as np
    from scipy.spatial import ConvexHull

    by_visual, frames = assembly_frames(source, config, coverage)
    shapes = [
        (part["visual"], np.fromstring(part["vertices"], sep=" ").reshape(-1, 3)) for part in parts
    ]
    for name, parent, transform, axis, core in frames:
        x = np.cross(axis, np.eye(3)[int(np.argmin(np.abs(axis)))])
        x /= np.linalg.norm(x)
        y = np.cross(axis, x)
        normals = [np.cos(t) * x + np.sin(t) * y for t in np.arange(16) * (2 * np.pi / 16)] + [
            axis,
            -axis,
        ]
        distances = [(core["radiusM"] - 0.000002) * np.cos(np.pi / 16)] * 16 + [
            core["axialHalfExtentM"] - 0.000002
        ] * 2
        divided = []
        for visual, vertices in shapes:
            group, pose = by_visual[visual]
            if group not in (parent, name):
                divided.append((visual, vertices))
                continue
            origin = pose[:3, 3] - transform[:3, 3]
            bounds = swept_bounds(vertices @ pose[:3, :3].T + origin, axis)
            if (
                (
                    bounds[3] <= core["radiusM"]
                    and max(abs(bounds[0]), abs(bounds[1])) <= core["axialHalfExtentM"]
                )
                or bounds[2] > core["radiusM"]
                or bounds[0] > core["axialHalfExtentM"]
                or bounds[1] < -core["axialHalfExtentM"]
            ):
                divided.append((visual, vertices))
                continue
            planes = [
                (pose[:3, :3].T @ n, distance - n @ origin)
                for n, distance in zip(normals, distances, strict=True)
            ]
            pieces = split_core(vertices, planes)
            original = ConvexHull(vertices)
            volume = sum(ConvexHull(piece).volume for piece in pieces)
            if abs(volume - original.volume) > max(original.volume * 0.01, original.area * 0.00001):
                raise ValueError(f"配合区切分体积误差超出微米量化预算：{visual}")
            divided.extend((visual, piece) for piece in pieces)
        shapes = divided
    result = []
    for item in coverage:
        pieces = [v for visual, v in shapes if visual == item["visual"]]
        pieces.sort(key=lambda v: tuple(v.mean(axis=0)))
        item["coacdParts"] = item["parts"]
        item["parts"] = len(pieces)
        if not 1 <= len(pieces) <= 256:
            raise ValueError("装配区域分割超出预算")
        for i, vertices in enumerate(pieces):
            if len(vertices) > 256:
                raise ValueError(f"装配分割凸包顶点超出预算：{item['visual']} / {len(vertices)}")
            result.append(
                {
                    "name": f"{item['link']}__h{i:03d}",
                    "visual": item["visual"],
                    "vertices": " ".join(
                        f"{value:.6f}" for point in sorted(map(tuple, vertices)) for value in point
                    ),
                }
            )
    return result


def cook(source, config, cache):
    import coacd
    import numpy as np
    import trimesh
    from scipy.spatial import ConvexHull

    for name, version in VERSIONS.items():
        if importlib.metadata.version(name) != version:
            raise ValueError(f"请使用离线固定依赖：{name}=={version}")
    coacd.set_log_level("info")
    parts = []
    coverage = []
    for visual in visuals(source, config):
        key = digest(
            json.dumps(
                {
                    "source": visual["sha256"],
                    "scale": visual["scale"],
                    "parameters": PARAMETERS,
                    "versions": VERSIONS,
                    "threads": 4,
                },
                sort_keys=True,
            ).encode()
        )
        cached = cache / f"{key}.json"
        mesh = trimesh.load(source / visual["mesh"], force="mesh")
        mesh.apply_scale(visual["scale"])
        if cached.exists():
            hulls = json.loads(cached.read_text())
        else:
            print(f"生成 {visual['mesh']}", flush=True)
            raw = coacd.run_coacd(coacd.Mesh(mesh.vertices, mesh.faces), **PARAMETERS)
            hulls = [sorted(np.round(v, 6).tolist()) for v, _ in raw]
            hulls.sort(key=lambda vertices: tuple(np.mean(vertices, axis=0)))
            cache.mkdir(parents=True, exist_ok=True)
            temporary = cached.with_suffix(".tmp")
            temporary.write_text(json.dumps(hulls))
            temporary.replace(cached)
        if not 1 <= len(hulls) <= 128:
            raise ValueError(f"凸包数量超出预算：{visual['mesh']}")
        # An explicit sampled coverage diagnostic, not a certified Hausdorff bound.
        samples = np.concatenate([mesh.vertices, mesh.triangles_center])[::3]
        outside = np.full(len(samples), np.inf)
        for index, vertices in enumerate(hulls):
            v = np.asarray(vertices)
            if (
                v.ndim != 2
                or v.shape[1] != 3
                or not 4 <= len(v) <= PARAMETERS["max_ch_vertex"]
                or not np.isfinite(v).all()
            ):
                raise ValueError("凸包顶点损坏或超出预算")
            if (v.min(axis=0) < mesh.bounds[0] - 0.002).any() or (
                v.max(axis=0) > mesh.bounds[1] + 0.002
            ).any():
                raise ValueError("凸包超出源网格边界预算")
            hull = ConvexHull(v)
            if hull.volume <= 1e-15:
                raise ValueError("退化凸包")
            signed_planes = (samples @ hull.equations[:, :3].T + hull.equations[:, 3]).max(axis=1)
            outside = np.minimum(outside, signed_planes)
            parts.append(
                {
                    "name": f"{visual['link']}__h{index:03d}",
                    "visual": visual["visual"],
                    "vertices": " ".join(f"{x:.6f}" for x in v.ravel()),
                }
            )
        if outside.max() > 0.0015:
            raise ValueError(f"源表面采样覆盖不合格：{visual['mesh']}")
        coverage.append(
            {
                **visual,
                "parts": len(hulls),
                "samples": len(samples),
                "coacdSampleOutsidePlanesMaxM": float(max(0, outside.max())),
                "sourceWatertight": bool(mesh.is_watertight),
                "cacheKey": key,
            }
        )
        print(
            f"完成 {visual['link']}: {len(hulls)} hulls; sampled plane gap={outside.max():.6g} m",
            flush=True,
        )
    parts = partition_assembly_cores(source, config, parts, coverage)
    policy = joint_policy(source, config, parts, coverage)
    geometry_sha = digest(json.dumps(parts, sort_keys=True, separators=(",", ":")).encode())
    physics = {"timestep": 0.001, "solref": [0.002, 1], "solimp": [0.95, 0.99, 0.001]}
    recipe_sha = digest(
        json.dumps(
            {"geometry": geometry_sha, "policy": policy, "physics": physics},
            sort_keys=True,
            separators=(",", ":"),
        ).encode()
    )
    return {
        "revision": 4,
        "geometrySha256": geometry_sha,
        "recipeSha256": recipe_sha,
        "physics": physics,
        "jointPolicy": policy,
        "source": config,
        "generator": {
            "versions": VERSIONS,
            "parameters": PARAMETERS,
            "threads": 4,
            "corePartition": {"sides": 16, "insetM": 0.000002},
        },
        "coverage": coverage,
        "parts": parts,
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", type=Path, default=ROOT / "build/lekiwi/URDF")
    parser.add_argument("--cache", type=Path, default=ROOT / "build/collision-cache")
    parser.add_argument(
        "--output", type=Path, default=ROOT / "robot_profiles/lekiwi-full-collision.json"
    )
    parser.add_argument("--check", action="store_true")
    args = parser.parse_args()
    result = cook(args.source, json.loads(CONFIG.read_text()), args.cache)
    if args.check:
        if json.loads(args.output.read_text()) != result:
            raise ValueError("已生成碰撞数据与源文件/算法参数不一致")
    else:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(result, indent=2) + "\n")
    print(f"整臂碰撞覆盖 {len(result['coverage'])} 个视觉网格 / {len(result['parts'])} 个凸包")


if __name__ == "__main__":
    main()