#!/usr/bin/env python3
"""
hyworld_to_pybullet.py — Load a HY-World-2.0 point cloud into PyBullet as a
static, collidable physics environment, and (optionally) drop the RX1
humanoid into it to verify the reconstructed geometry holds up under gravity.

Pipeline:
  points.ply (raw point cloud)
    -> Open3D: normal estimation + Poisson surface reconstruction
    -> density-based trimming (removes hallucinated low-density blobs)
    -> quadric decimation (collision-friendly triangle count)
    -> world_mesh.obj
    -> PyBullet: static (mass=0) concave collision + visual body
    -> [optional] RX1 URDF spawned above the mesh, simulated under gravity

Usage:
  python scripts/hyworld_to_pybullet.py --points points.ply --out <dir>
  python scripts/hyworld_to_pybullet.py --points points.ply --out <dir> --no-robot
  python scripts/hyworld_to_pybullet.py --points points.ply --out <dir> --collision-triangles 10000
"""
from __future__ import annotations

import argparse
import json
import logging
import pathlib
import sys
import time

log = logging.getLogger("hyworld_to_pybullet")

# RX1 URDF is bundled alongside the MuJoCo assets — mesh paths inside it are
# relative to the URDF's own directory, so no separate mesh copy is needed.
_RX1_URDF = pathlib.Path(__file__).resolve().parent.parent / "rx1_mujoco" / "rx1.urdf"


def _build_collision_mesh(points_path: pathlib.Path, out_dir: pathlib.Path, target_triangles: int) -> dict:
    """Point cloud -> mesh pipeline. Returns metadata about the produced mesh."""
    try:
        import open3d as o3d
        import numpy as np
    except ImportError as e:
        return {"error": f"open3d not installed: {e}"}

    log.info("Loading point cloud: %s", points_path)
    pcd = o3d.io.read_point_cloud(str(points_path))
    if len(pcd.points) == 0:
        return {"error": f"no points found in {points_path}"}
    log.info("Loaded %d points", len(pcd.points))

    log.info("Estimating normals...")
    pcd.estimate_normals()
    pcd.orient_normals_consistent_tangent_plane(k=100)

    log.info("Poisson surface reconstruction...")
    mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pcd, depth=9)
    densities = np.asarray(densities)

    log.info("Trimming low-density (hallucinated) regions...")
    threshold = np.quantile(densities, 0.05)
    mesh.remove_vertices_by_mask(densities < threshold)
    log.info("Mesh after trimming: %d vertices, %d triangles", len(mesh.vertices), len(mesh.triangles))

    if len(mesh.triangles) == 0:
        return {"error": "mesh has 0 triangles after trimming — reconstruction failed"}

    log.info("Decimating to ~%d triangles for collision performance...", target_triangles)
    if len(mesh.triangles) > target_triangles:
        decimated = mesh.simplify_quadric_decimation(target_number_of_triangles=target_triangles)
        # Known Open3D issue on non-watertight meshes: quadric decimation can
        # fail to converge near the target. Fall back to vertex clustering.
        if len(decimated.triangles) == 0 or len(decimated.triangles) > target_triangles * 3:
            log.warning("Quadric decimation didn't converge — falling back to vertex clustering")
            bbox = mesh.get_axis_aligned_bounding_box()
            diag = np.linalg.norm(bbox.get_extent())
            voxel_size = diag / (target_triangles ** 0.5)
            decimated = mesh.simplify_vertex_clustering(voxel_size)
        mesh = decimated
    log.info("Final collision mesh: %d vertices, %d triangles", len(mesh.vertices), len(mesh.triangles))

    # Clean up small disconnected components (common photogrammetry noise)
    try:
        import trimesh
        tm = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.triangles)
        components = tm.split(only_watertight=False)
        if len(components) > 1:
            largest = max(components, key=lambda c: len(c.faces))
            log.info("Dropped %d small disconnected components", len(components) - 1)
            mesh = o3d.geometry.TriangleMesh(
                o3d.utility.Vector3dVector(largest.vertices),
                o3d.utility.Vector3iVector(largest.faces),
            )
    except ImportError:
        log.info("trimesh not available — skipping disconnected-component cleanup")
    except Exception as e:
        log.warning("Component cleanup failed, using undedup mesh: %s", e)

    mesh.compute_vertex_normals()
    out_dir.mkdir(parents=True, exist_ok=True)
    mesh_path = out_dir / "world_mesh.obj"
    o3d.io.write_triangle_mesh(str(mesh_path), mesh)
    log.info("Wrote collision mesh: %s", mesh_path)

    bbox = mesh.get_axis_aligned_bounding_box()
    return {
        "mesh_path": str(mesh_path),
        "triangles": len(mesh.triangles),
        "vertices": len(mesh.vertices),
        "bbox_min": bbox.min_bound.tolist(),
        "bbox_max": bbox.max_bound.tolist(),
    }


def _run_pybullet(mesh_info: dict, include_robot: bool, spawn_height: float) -> dict:
    try:
        import pybullet as p
        import pybullet_data
    except ImportError as e:
        return {"error": f"pybullet not installed: {e}"}

    mesh_path = mesh_info["mesh_path"]
    bbox_min, bbox_max = mesh_info["bbox_min"], mesh_info["bbox_max"]

    log.info("Connecting to PyBullet (GUI)...")
    p.connect(p.GUI)
    p.setAdditionalSearchPath(pybullet_data.getDataPath())
    p.setGravity(0, 0, -9.81)

    log.info("Loading world mesh as static collision + visual body...")
    col_id = p.createCollisionShape(p.GEOM_MESH, fileName=mesh_path, flags=p.GEOM_FORCE_CONCAVE_TRIMESH)
    vis_id = p.createVisualShape(p.GEOM_MESH, fileName=mesh_path)
    world_id = p.createMultiBody(baseMass=0, baseCollisionShapeIndex=col_id, baseVisualShapeIndex=vis_id)

    robot_id = None
    if include_robot:
        if not _RX1_URDF.exists():
            log.warning("RX1 URDF not found at %s — skipping robot spawn", _RX1_URDF)
        else:
            top_z = bbox_max[2] if bbox_max else 1.0
            spawn_pos = [
                (bbox_min[0] + bbox_max[0]) / 2 if bbox_min else 0,
                (bbox_min[1] + bbox_max[1]) / 2 if bbox_min else 0,
                top_z + spawn_height,
            ]
            log.info("Spawning RX1 at %s", spawn_pos)
            robot_id = p.loadURDF(str(_RX1_URDF), basePosition=spawn_pos, useFixedBase=False)

    log.info("Simulating settle-down (5s)...")
    for _ in range(int(5 / (1 / 240))):
        p.stepSimulation()
        time.sleep(1 / 240)

    robot_settled = None
    if robot_id is not None:
        pos, _ = p.getBasePositionAndOrientation(robot_id)
        robot_settled = pos[2] > (bbox_min[2] if bbox_min else -100) - 1.0  # didn't fall through
        log.info("Robot final position: %s (settled=%s)", pos, robot_settled)

    log.info("Simulation running — close the PyBullet window to exit.")
    try:
        while p.isConnected():
            p.stepSimulation()
            time.sleep(1 / 240)
    except KeyboardInterrupt:
        pass

    return {
        "world_id": world_id,
        "robot_id": robot_id,
        "robot_settled": robot_settled,
    }


def main():
    ap = argparse.ArgumentParser(description="Load a HY-World point cloud into PyBullet")
    ap.add_argument("--points", required=True, help="path to points.ply")
    ap.add_argument("--out", required=True, help="output directory for the collision mesh")
    ap.add_argument("--collision-triangles", type=int, default=20000)
    ap.add_argument("--robot", dest="robot", action="store_true", default=True)
    ap.add_argument("--no-robot", dest="robot", action="store_false")
    ap.add_argument("--spawn-height", type=float, default=1.0)
    args = ap.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s  %(name)s  %(levelname)s  %(message)s",
        datefmt="%H:%M:%S",
    )

    points_path = pathlib.Path(args.points)
    out_dir = pathlib.Path(args.out)
    if not points_path.exists():
        result = {"error": f"points file not found: {points_path}"}
        print("PYBULLET_RESULT " + json.dumps(result))
        sys.exit(1)

    mesh_info = _build_collision_mesh(points_path, out_dir, args.collision_triangles)
    if "error" in mesh_info:
        log.error("Mesh build failed: %s", mesh_info["error"])
        print("PYBULLET_RESULT " + json.dumps(mesh_info))
        sys.exit(1)

    sim_info = _run_pybullet(mesh_info, args.robot, args.spawn_height)
    if "error" in sim_info:
        log.error("PyBullet failed: %s", sim_info["error"])
        print("PYBULLET_RESULT " + json.dumps({**mesh_info, **sim_info}))
        sys.exit(1)

    result = {**mesh_info, **sim_info, "out_dir": str(out_dir)}
    print("PYBULLET_RESULT " + json.dumps(result))


if __name__ == "__main__":
    main()