import os
os.environ["OMNI_KIT_ACCEPT_EULA"] = "YES" 
from isaacsim import SimulationApp

simulation_app = SimulationApp({"headless": False}) 

from omni.isaac.core.utils.extensions import enable_extension
enable_extension("omni.isaac.ros2_bridge")
simulation_app.update() 

import numpy as np
from PIL import Image 
import omni.kit.commands
from pathlib import Path

# 【底层基石 USD API】
import omni.usd                  
from pxr import Gf, Sdf, UsdShade, UsdGeom, Vt, PhysxSchema

from omni.isaac.core import World
from omni.isaac.core.objects import FixedCuboid
# 【新增引入】DynamicCuboid 用于生成受物理世界重力影响的动态刚体车辆
from omni.isaac.core.objects import DynamicCuboid 

from omni.isaac.core.utils.prims import create_prim
from omni.isaac.core.utils.viewports import set_camera_view
from omni.isaac.core.utils.rotations import euler_angles_to_quat

from omni.isaac.sensor import Camera
import omni.replicator.core as rep
import omni.graph.core as og
from omni.isaac.core.objects import VisualCuboid

# 【引入 URDF 导入器和机器人核心类】
import omni.kit.commands
from omni.importer.urdf import _urdf
from omni.isaac.core.robots import Robot
from omni.isaac.core.utils.stage import add_reference_to_stage


def create_checkerboard_image(filepath="checkerboard.png", rows=6, cols=9, square_size_px=500):
    width = cols * square_size_px
    height = rows * square_size_px
    img = np.ones((height, width, 3), dtype=np.uint8) * 255
    
    for r in range(rows):
        for c in range(cols):
            if (r + c) % 2 == 1:
                img[r*square_size_px:(r+1)*square_size_px, c*square_size_px:(c+1)*square_size_px] = 0
                
    border = square_size_px
    img_with_border = np.pad(img, pad_width=((border, border), (border, border), (0, 0)), mode='constant', constant_values=255)
    
    pil_img = Image.fromarray(img_with_border)
    abs_filepath = Path(filepath).resolve()
    pil_img.save(abs_filepath)
    usd_filepath = str(abs_filepath).replace("\\", "/")
    print(f"[*] 棋盘格纹理已自动生成: {usd_filepath}")
    return usd_filepath

def add_corner_rotary_lidars(room_length=10.0, room_width=6.0, height=3.5, 
                             lidar_config="Example_Rotary", topic_prefix="/workshop/lidar"):
    offset = 0.3
    x_pos = (room_length / 2.0) - offset
    y_pos = (room_width / 2.0) - offset

    lidar_configs = [
        {"name": "FL", "pos": [ x_pos,  y_pos, height], "yaw": np.degrees(np.arctan2(-y_pos, -x_pos))},
        {"name": "FR", "pos": [ x_pos, -y_pos, height], "yaw": np.degrees(np.arctan2( y_pos, -x_pos))},
        {"name": "BL", "pos": [-x_pos,  y_pos, height], "yaw": np.degrees(np.arctan2(-y_pos,  x_pos))},
        {"name": "BR", "pos": [-x_pos, -y_pos, height], "yaw": np.degrees(np.arctan2( y_pos,  x_pos))}
    ]

    keys = og.Controller.Keys
    graph_path = "/World/ROS2_Lidar_Graph"
    nodes = [("OnTick", "omni.graph.action.OnTick"), ("ReadSimTime", "omni.isaac.core_nodes.IsaacReadSimulationTime"), ("PublishTF", "omni.isaac.ros2_bridge.ROS2PublishTransformTree")]
    connections = [("OnTick.outputs:tick", "PublishTF.inputs:execIn"), ("ReadSimTime.outputs:simulationTime", "PublishTF.inputs:timeStamp")]
    set_values = []
    lidar_paths = [] 

    for cfg in lidar_configs:
        lidar_path = f"/World/Sensors/Lidar_{cfg['name']}"
        lidar_paths.append(lidar_path)
        quat = euler_angles_to_quat(np.array([0, 15.0, cfg['yaw']]), degrees=True)
        orientation = Gf.Quatd(quat[0], quat[1], quat[2], quat[3])
        omni.kit.commands.execute("IsaacSensorCreateRtxLidar", path=lidar_path, parent=None, config=lidar_config, translation=Gf.Vec3d(*cfg["pos"]), orientation=orientation)
        render_product = rep.create.render_product(lidar_path, [1, 1])
        helper_name = f"ROS2LidarHelper_{cfg['name']}"
        nodes.append((helper_name, "omni.isaac.ros2_bridge.ROS2RtxLidarHelper"))
        connections.append(("OnTick.outputs:tick", f"{helper_name}.inputs:execIn"))
        set_values.extend([
            (f"{helper_name}.inputs:renderProductPath", str(render_product.path)), (f"{helper_name}.inputs:topicName", f"{topic_prefix}/{cfg['name'].lower()}/pointcloud"),
            (f"{helper_name}.inputs:frameId", f"Lidar_{cfg['name']}"), (f"{helper_name}.inputs:type", "point_cloud"), (f"{helper_name}.inputs:fullScan", True) 
        ])

    set_values.append(("PublishTF.inputs:targetPrims", lidar_paths))
    og.Controller.edit({"graph_path": graph_path, "evaluator_name": "execution"}, {keys.CREATE_NODES: nodes, keys.CONNECT: connections, keys.SET_VALUES: set_values})

def create_raw_usd_material(stage, mat_path, tex_path):
    material = UsdShade.Material.Define(stage, mat_path)
    pbr_shader = UsdShade.Shader.Define(stage, f"{mat_path}/PBRShader")
    pbr_shader.CreateIdAttr("UsdPreviewSurface")
    pbr_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(1.0) 
    pbr_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)  
    
    tex_sampler = UsdShade.Shader.Define(stage, f"{mat_path}/diffuseTexture")
    tex_sampler.CreateIdAttr("UsdUVTexture")
    tex_sampler.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(Sdf.AssetPath(tex_path))
    tex_sampler.CreateInput("magFilter", Sdf.ValueTypeNames.Token).Set("nearest")
    tex_sampler.CreateInput("minFilter", Sdf.ValueTypeNames.Token).Set("nearest")
    
    st_reader = UsdShade.Shader.Define(stage, f"{mat_path}/stReader")
    st_reader.CreateIdAttr("UsdPrimvarReader_float2")
    st_reader.CreateInput("varname", Sdf.ValueTypeNames.Token).Set("st")
    
    tex_sampler.CreateInput("st", Sdf.ValueTypeNames.Float2).ConnectToSource(st_reader.ConnectableAPI(), "result")
    pbr_shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).ConnectToSource(tex_sampler.ConnectableAPI(), "rgb")
    material.CreateSurfaceOutput().ConnectToSource(pbr_shader.ConnectableAPI(), "surface")
    return material

def create_textured_board(stage, prim_path, width, height, center, euler_rot_deg, usd_material):
    mesh = UsdGeom.Mesh.Define(stage, prim_path)
    w, h = width / 2.0, height / 2.0
    mesh.GetPointsAttr().Set(Vt.Vec3fArray([Gf.Vec3f(-w, -h, 0), Gf.Vec3f( w, -h, 0), Gf.Vec3f( w,  h, 0), Gf.Vec3f(-w,  h, 0)]))
    mesh.GetFaceVertexCountsAttr().Set([4])
    mesh.GetFaceVertexIndicesAttr().Set([0, 1, 2, 3])
    mesh.GetNormalsAttr().Set([Gf.Vec3f(0, 0, 1)] * 4)
    mesh.SetNormalsInterpolation(UsdGeom.Tokens.vertex)
    
    primvars_api = UsdGeom.PrimvarsAPI(mesh)
    st_primvar = primvars_api.CreatePrimvar("st", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.vertex)
    st_primvar.Set([Gf.Vec2f(0, 0), Gf.Vec2f(1, 0), Gf.Vec2f(1, 1), Gf.Vec2f(0, 1)])
    mesh.GetExtentAttr().Set([Gf.Vec3f(-w, -h, -0.01), Gf.Vec3f(w, h, 0.01)])
    
    xform = UsdGeom.Xformable(mesh)
    xform.AddTranslateOp().Set(Gf.Vec3d(*center))
    xform.AddRotateXYZOp().Set(Gf.Vec3f(*euler_rot_deg)) 
    UsdShade.MaterialBindingAPI.Apply(mesh.GetPrim()).Bind(usd_material)
    return mesh


def build_workshop():
    world = World(stage_units_in_meters=1.0)
    L, W, H, T = 10.0, 6.0, 3.5, 0.2
    floor_color = np.array([0.2, 0.2, 0.2])
    wall_color = np.array([0.8, 0.8, 0.8])

    world.scene.add(FixedCuboid(prim_path="/World/Workshop/Floor", name="floor", position=np.array([0, 0, -T/2]), scale=np.array([L + 2*T, W + 2*T, T]), color=floor_color))
    world.scene.add(FixedCuboid(prim_path="/World/Workshop/Ceiling", name="ceiling", position=np.array([0, 0, H + T/2]), scale=np.array([L + 2*T, W + 2*T, T]), color=wall_color))
    world.scene.add(FixedCuboid(prim_path="/World/Workshop/Wall_Front", name="wall_front", position=np.array([L/2 + T/2, 0, H/2]), scale=np.array([T, W, H]), color=wall_color))
    world.scene.add(FixedCuboid(prim_path="/World/Workshop/Wall_Back", name="wall_back", position=np.array([-L/2 - T/2, 0, H/2]), scale=np.array([T, W, H]), color=wall_color))
    world.scene.add(FixedCuboid(prim_path="/World/Workshop/Wall_Left", name="wall_left", position=np.array([0, W/2 + T/2, H/2]), scale=np.array([L + 2*T, T, H]), color=wall_color))
    world.scene.add(FixedCuboid(prim_path="/World/Workshop/Wall_Right", name="wall_right", position=np.array([0, -W/2 - T/2, H/2]), scale=np.array([L + 2*T, T, H]), color=wall_color))

    light_positions = [(L/4, W/4, H - 0.5), (L/4, -W/4, H - 0.5), (-L/4, W/4, H - 0.5), (-L/4, -W/4, H - 0.5)]
    for i, pos in enumerate(light_positions):
        create_prim(prim_path=f"/World/Workshop/Lights/Light_{i}", prim_type="SphereLight", position=np.array(pos), attributes={"inputs:radius": 0.3, "inputs:intensity": 30000.0, "inputs:color": (1.0, 1.0, 0.95)})

    cb_rows, cb_cols = 6, 9
    cb_square_size = 0.20 
    tex_path = create_checkerboard_image("checkerboard.png", rows=cb_rows, cols=cb_cols, square_size_px=500)
    stage = omni.usd.get_context().get_stage()

    mat_path = "/World/Workshop/Materials/CheckerboardMat"
    usd_material = create_raw_usd_material(stage, mat_path, tex_path)

    board_w = (cb_cols + 2) * cb_square_size 
    board_h = (cb_rows + 2) * cb_square_size 
    z_height = H / 2.0
    offset = 0.05 

    board_configs = [
        ("/World/Workshop/CalibrationBoards/Front", [ L/2 - offset, 0, z_height], [90, 0,  90]),
        ("/World/Workshop/CalibrationBoards/Back",  [-L/2 + offset, 0, z_height], [90, 0, -90]),
        ("/World/Workshop/CalibrationBoards/Left",  [0,  W/2 - offset, z_height], [90, 0,   0]),
        ("/World/Workshop/CalibrationBoards/Right", [0, -W/2 + offset, z_height], [90, 0, 180])
    ]
    for path, pos, euler_rot in board_configs:
        create_textured_board(stage, path, board_w, board_h, pos, euler_rot, usd_material)
    
    camera_z = H - 0.1 
    camera = Camera(prim_path="/World/Workshop/CalibrationCamera", position=np.array([0.0, 0.0, camera_z]), frequency=20, resolution=(1280, 720))
    camera.set_world_pose(orientation=np.array([1.0, 0.0, 0.0, 0.0]))
    camera.initialize()
    camera.set_focal_length(5.0)

    keys = og.Controller.Keys
    og.Controller.edit({"graph_path": "/World/ROS2_Camera_Graph", "evaluator_name": "execution"},
        {keys.CREATE_NODES: [("OnTick", "omni.graph.action.OnTick"), ("ROS2Camera", "omni.isaac.ros2_bridge.ROS2CameraHelper")],
         keys.CONNECT: [("OnTick.outputs:tick", "ROS2Camera.inputs:execIn")],
         keys.SET_VALUES: [("ROS2Camera.inputs:renderProductPath", camera.get_render_product_path()), ("ROS2Camera.inputs:topicName", "/AutoCalib_Workshop/camera/image_raw"), ("ROS2Camera.inputs:type", "rgb")]})

    add_corner_rotary_lidars(room_length=L, room_width=W, height=H - 0.2, lidar_config="Example_Rotary", topic_prefix="/AutoCalib_Workshop/lidar")


    # ================= 【🚗核心新增 1：构建物理层 AGV 底盘】 =================
    # 主车体：带质量的物理刚体
    # world.scene.add(
    #     DynamicCuboid(
    #         prim_path="/World/Workshop/Vehicle", name="agv_vehicle",
    #         position=np.array([0.0, 0.0, 0.2]), # 中心高度 20cm，完美贴地防穿模
    #         scale=np.array([0.8, 0.5, 0.3]),    # 车辆尺寸：长0.8m x 宽0.5m x 高0.3m
    #         color=np.array([0.2, 0.6, 1.0]),    # 亮蓝色车身
    #         mass=50.0                           # 赋予 50kg 的真实物理质量
    #     )
    # )
    
    # 车头指示器：红色方块，挂载在车头正前方，明确指示 +X 前进方向
    # VisualCuboid(
    #     prim_path="/World/Workshop/Vehicle/DirectionMarker",
    #     name="direction_marker",
    #     position=np.array([0.4, 0.0, 0.0]), # 相对车身局部前移
    #     scale=np.array([0.1, 0.51, 0.31]), 
    #     color=np.array([1.0, 0.0, 0.0])     # 直接通过内置的 color 参数设置
    # )

   # ================= 【🚗核心新增 1：导入真实 URDF 替换基础方块】 =================
    urdf_file_path = "/home/nvidia/study/AutoCalib-Workshop/models/ack_m.urdf"
    # 🚨 新增：指定转换后的 USD 文件保存在哪里（必须带 .usd 后缀）
    dest_usd_path = "/home/nvidia/study/AutoCalib-Workshop/models/ack_m.usd" 
    
    dest_prim_path = "/World/Workshop/Vehicle"

    # 1. 配置 URDF 导入参数
    import_config = _urdf.ImportConfig()
    import_config.merge_fixed_joints = True  # 🚨 关键修改：改为 True！将雷达、相机和空节点合并进主车身
    import_config.convex_decomp = False
    import_config.fix_base = False 
    import_config.make_default_prim = True

    # 2. 将 URDF 解析并保存为本地的 USD 文件
    omni.kit.commands.execute(
        "URDFParseAndImportFile",
        urdf_path=urdf_file_path,
        import_config=import_config,
        dest_path=dest_usd_path  # 传入的是硬盘文件路径
    )

    # 3. 🔥 将硬盘上的 USD 文件作为引用（Reference）挂载到场景树中
    add_reference_to_stage(usd_path=dest_usd_path, prim_path=dest_prim_path)

    # 4. 包装为 Robot 对象
    world.scene.add(
        Robot(
            prim_path=dest_prim_path, 
            name="agv_vehicle",
            position=np.array([0.0, 0.0, 0.05]) 
        )
    )
    # ========================================================================================

    # 强制关闭物理引擎对车辆的“休眠优化(Sleep)”，确保它随时能被指令叫醒移动
    physx_rb = PhysxSchema.PhysxRigidBodyAPI.Get(stage, "/World/Workshop/Vehicle")
    if physx_rb:
        physx_rb.GetSleepThresholdAttr().Set(0.0)

    # ================= 【🚗核心新增 2：无缝底层 Twist 订阅图】 =================
    og.Controller.edit({"graph_path": "/World/ROS2_Twist_Graph", "evaluator_name": "execution"},
        {
            keys.CREATE_NODES: [
                ("OnTick", "omni.graph.action.OnTick"),
                ("TwistSub", "omni.isaac.ros2_bridge.ROS2SubscribeTwist"), # <--- 已修正
            ],
            keys.CONNECT: [
                ("OnTick.outputs:tick", "TwistSub.inputs:execIn"),
            ],
            keys.SET_VALUES: [
                ("TwistSub.inputs:topicName", "/cmd_vel"),
            ]
        })
    # ========================================================================================

    return world


def main():
    world = build_workshop()

    world.reset()

    # 上帝视角的俯视监控
    set_camera_view(eye=np.array([0.001, 0.0, 3.4]), target=np.array([0.0, 0.0, 0.0]))

    print("======================================================")
    print(" 🎯 标定车间完美运行！全向 AGV (蓝身红头) 已就绪！")
    print(" ---------------------------------------------------")
    print(" 🎮 车辆控制指南：请打开您的**原生终端** (不激活 conda)，输入:")
    print("   ros2 topic pub /cmd_vel geometry_msgs/msg/Twist \"{linear: {x: 0.5}, angular: {z: 0.8}}\"")
    print("======================================================")

    agv = world.scene.get_object("agv_vehicle")

    while simulation_app.is_running():
        # ================= 【🚗核心闭环：实时提取 Twist，转换物理运动学】 =================
        try:
            # 1. 每一帧从底层 ActionGraph 中拉取解包出来的 ROS2 速度指令
            lin_vel = og.Controller.get(og.Controller.attribute("/World/ROS2_Twist_Graph/TwistSub.outputs:linearVelocity"))
            ang_vel = og.Controller.get(og.Controller.attribute("/World/ROS2_Twist_Graph/TwistSub.outputs:angularVelocity"))
            
            if lin_vel is not None and ang_vel is not None and len(lin_vel) == 3:
                # 2. 获取车辆在物理世界中实时的绝对姿态四元数 [w, x, y, z] 和 自然下落速度
                pos, quat = agv.get_world_pose()
                curr_lin_vel = agv.get_linear_velocity()
                
                # 3. 构造 3D 四元数和旋转矩阵
                q = Gf.Quatd(float(quat[0]), float(quat[1]), float(quat[2]), float(quat[3]))
                rot_mat = Gf.Matrix3d(Gf.Rotation(q))
                
                # 🚨 核心修复：USD (pxr.Gf) 中，向量与矩阵相乘直接使用 * 运算符（行向量右乘矩阵）
                world_lin_vel = Gf.Vec3d(*lin_vel) * rot_mat
                world_ang_vel = Gf.Vec3d(*ang_vel) * rot_mat
                
                # 4. 物理防翻车约束
                target_lin_vel = np.array([world_lin_vel[0], world_lin_vel[1], curr_lin_vel[2]])
                target_ang_vel = np.array([0.0, 0.0, world_ang_vel[2]])
                
                # 5. 直接对车辆物理质心施加强制推演覆盖
                agv.set_linear_velocity(target_lin_vel)
                agv.set_angular_velocity(target_ang_vel)
                
                # Debug 日志：确认 ROS2 话题是否真正连通
                if abs(lin_vel[0]) > 0.01 or abs(ang_vel[2]) > 0.01:
                    print(f"\r[ROS2 Debug] 车辆移动中: 线速度 {lin_vel[0]:.2f}, 角速度 {ang_vel[2]:.2f}", end="")

        except Exception as e:
            # 🚨 永远不要用 pass 吞掉这里的报错
            print(f"\n[ERROR] 运动学控制循环异常: {e}")
        # ============================================================================================

        world.step(render=True)

    simulation_app.close()

if __name__ == "__main__":
    main()