import trimesh
import importlib
import numpy as np
from pathlib import Path
from copy import deepcopy
import transforms3d as t3d
from threading import Thread
import readline

import sys

import trimesh.bounds

sys.path.append(".")
from envs.utils import *
import sapien.core as sapien
from sapien.utils.viewer import Viewer
from tqdm import tqdm
from PIL import Image

import re
import time
from typing import List, Literal

from sapien import Pose

# obj square
CAMERA_POSE = Pose([0, 0.134123, 0.96], [0.684988, 0.174248, 0.173926, -0.685696])
# main graph
# CAMERA_POSE = Pose([0.0293144, -1.12261, 1.52599], [0.665553, 0.233024, 0.231257, -0.670268])

class Helper:
    POINTS = [
        ("target_pose", "target"),
        ("contact_points_pose", "contact"),
        ("functional_matrix", "functional"),
        ("orientation_point", "orientation"),
    ]

    def create_scene(self, viewer=True, **kwargs):
        """
        Set the scene
            - Set up the basic scene: light source, viewer.
        """
        self.engine = sapien.Engine()
        # declare sapien renderer
        from sapien.render import set_global_config

        set_global_config(max_num_materials=50000, max_num_textures=50000)
        self.renderer = sapien.SapienRenderer()
        # give renderer to sapien sim
        self.engine.set_renderer(self.renderer)

        sapien.render.set_camera_shader_dir("rt")
        sapien.render.set_ray_tracing_samples_per_pixel(32)
        sapien.render.set_ray_tracing_path_depth(8)
        sapien.render.set_ray_tracing_denoiser("oidn")

        # declare sapien scene
        scene_config = sapien.SceneConfig()
        self.scene = self.engine.create_scene(scene_config)
        # set simulation timestep
        self.scene.set_timestep(kwargs.get("timestep", 1 / 250))

        # initialize viewer with camera position and orientation
        if viewer:
            self.viewer = Viewer(self.renderer)
            self.viewer.set_scene(self.scene)
            self.viewer.set_camera_xyz(
                x=kwargs.get("camera_xyz_x", 0.4),
                y=kwargs.get("camera_xyz_y", 0.22),
                z=kwargs.get("camera_xyz_z", 1.5),
            )
            self.viewer.set_camera_rpy(
                r=kwargs.get("camera_rpy_r", 0),
                p=kwargs.get("camera_rpy_p", -0.8),
                y=kwargs.get("camera_rpy_y", 2.45),
            )
        else:
            self.viewer = None
        self.camera = self.scene.add_camera("camera", 2390, 1000, 1.57, 0.1, 1000)
        self.camera.set_pose(CAMERA_POSE)
        # scale = 1
        # self.camera = self.scene.add_camera(name="", width=2560*scale, height=1600*scale, fovy=1.57, near=0.1, far=1e+03)
        # self.camera.set_local_pose(sapien.Pose([-0.893507, -0.358009, 0.983116], [0.869079, 0.128192, 0.298895, -0.372735]))

    def create_table_and_wall(self):
        # add ground to scene
        # self.scene.add_ground(0)
        # set default physical material
        self.scene.default_physical_material = self.scene.create_physical_material(0.5, 0.5, 0)
        # give some white ambient light of moderate intensity
        self.scene.set_ambient_light([0.5, 0.5, 0.5])
        # default enable shadow unless specified otherwise
        shadow = False
        # default spotlight angle and intensity
        direction_lights = [[[0, 0.5, -1], [0.5, 0.5, 0.5]]]
        self.direction_light_lst = []
        for direction_light in direction_lights:
            self.direction_light_lst.append(self.scene.add_directional_light(direction_light[0], direction_light[1], shadow=shadow))
        # default point lights position and intensity
        point_lights = [
            [[1, 0, 1.8], [1, 1, 1]],
            [[-1, 0, 1.8], [1, 1, 1]],
            [[2.6, -1.7, 0.76], [1, 1, 1]],
            [[-2.6, -1.7, 0.76], [1, 1, 1]],
            [[-1.2, -4.4, 0.76], [1, 1, 1]],
            [[1.2, -4.4, 0.76], [1, 1, 1]],
        ]
        self.point_light_lst = []
        for point_light in point_lights:
            self.point_light_lst.append(self.scene.add_point_light(point_light[0], point_light[1], shadow=shadow))

        # creat wall
        wall_texture, table_texture = None, None
        # self.wall_texture, self.table_texture = 0, 0
        # self.wall = create_box(
        #     self.scene,
        #     sapien.Pose(p=[0, 1, 1.5]),
        #     half_size=[3, 0.6, 1.5],
        #     color=(1, 0.9, 0.9),
        #     name='wall',
        #     texture_id=wall_texture
        # )

        # self.table_z_bias = np.random.random()*0.3 - 0.15
        # print('bias:', self.table_z_bias)
        # self.table_z_bias = 0
        # table_height = self.table_z_bias + 0.74

        # creat table
        # self.table = create_table(
        #     self.scene,
        #     sapien.Pose(p = [0, 0, table_height]),
        #     length=2,
        #     width=4,
        #     height=table_height,
        #     thickness=0.05,
        #     is_static=True,
        #     texture_id=table_texture
        # )

    def init_messy(self):
        with open("./assets/objects/objaverse/list.json", "r") as file:
            self.messy_item_info = json.load(file)
        self.obj_names = self.messy_item_info["item_names"]
        self.size_dict = []
        self.obj_list = []
        self.max_obj_num = 1

    def add_messy_obj(self, name, idx, xlim=[-0.3, 0.3], ylim=[-0.2, 0.2], zlim=[0.741]):
        tyrs, max_try = 0, 100
        success_count, messy_obj = 0, None
        while tyrs < max_try:
            obj_str = f"{name}_{idx}"
            obj_radius = self.messy_item_info["radius"][obj_str]
            obj_offset = self.messy_item_info["z_offset"][obj_str]
            obj_maxz = self.messy_item_info["z_max"][obj_str]

            success, messy_obj = rand_create_cluttered_actor(
                self.scene,
                xlim=xlim,
                ylim=ylim,
                zlim=np.array(zlim),
                modelname=obj_str,
                rotate_rand=True,
                rotate_lim=[0, 0, np.pi],
                size_dict=self.size_dict,
                obj_radius=obj_radius,
                z_offset=obj_offset,
                z_max=obj_maxz,
                prohibited_area=[],
            )
            if not success:
                continue
            # self.viewer.paused = True
            # while self.viewer.paused:
            #     self.scene.update_render()
            #     self.viewer.render()
            messy_obj: sapien.Entity = messy_obj[0]
            messy_obj.set_name(obj_str)
            messy_obj.find_component_by_type(sapien.physx.PhysxRigidDynamicComponent).mass = 0.01
            success_count += 1

            pose = sapien.pysapien.Entity.get_pose(messy_obj).p.tolist()
            pose.append(obj_radius)
            self.size_dict.append(pose)
            self.obj_list.append(messy_obj)

            if len(self.obj_list) > self.max_obj_num:
                obj = self.obj_list.pop(0)
                self.size_dict.pop(0)
                self.scene.remove_actor(obj)
                self.scene.update_render()
                self.viewer.render()

            break
        return success_count == 1, messy_obj

    def check_urdf(self, name, idx, d_range=50, pose=None):
        if pose is None:
            success, obj = self.add_messy_obj(name, idx)
            if not success:
                return False

            def to_array(pose: sapien.Pose) -> np.ndarray:
                return np.array(pose.p.tolist() + pose.q.tolist())

            is_step, max_step = 0, 200
            pose_list = [to_array(obj.get_pose())]
            while is_step < max_step:
                self.scene.step()
                self.scene.update_render()
                self.viewer.render()

                new_pose = obj.get_pose()
                pose_list.append(to_array(new_pose))

                if len(pose_list) > d_range:
                    check_succ = True
                    for i in range(-d_range, 0):
                        if not np.allclose(pose_list[i], pose_list[-d_range], 1e-4):
                            check_succ = False
                            break
                    if check_succ:
                        break
                is_step += 1

            if is_step > 0 and is_step < max_step:
                success = True
            elif is_step >= max_step:
                success = False

            # self.viewer.paused = True
            # while self.viewer.paused:
            #     self.scene.update_render()
            #     self.viewer.render()
            return success
        else:
            modeldir = f"./assets/objects/objaverse/{name}/{idx}/"
            loader: sapien.URDFLoader = self.scene.create_urdf_loader()

            loader.fix_root_link = True
            loader.load_multiple_collisions_from_file = False
            object = loader.load_multiple(modeldir + "model.urdf")[1][0]

            object.set_pose(sapien.Pose(pose, [1, 0, 0, 0]))
            object.set_name(name)

            return True

    def test_messy(self):
        self.create_scene()
        self.create_table_and_wall()
        self.init_messy()

        self.result = []

        test_list = []
        for name in self.obj_names:
            for idx in self.messy_item_info["list_of_items"][name]:
                test_list.append((name, idx))

        # test_list = test_list[19:]
        for cnt, (name, idx) in enumerate(tqdm(test_list)):
            if name != "ramen_package":
                continue
            # if cnt > 0 and cnt % 100 == 0:
            #     time.sleep(3)
            #     self.scene.clear()
            #     self.create_table_and_wall()
            success = self.check_urdf(name, idx)
            self.result.append({"name": name, "idx": idx, "success": success})
            with open("result.jsonl", "a", encoding="utf-8") as f:
                f.write(json.dumps(self.result[-1]) + "\n")

        while len(self.obj_list) > 0:
            obj = self.obj_list.pop(0)
            self.size_dict.pop(0)
            self.scene.remove_actor(obj)
            self.scene.update_render()
            self.viewer.render()

    @staticmethod
    def trans_mat(to_mat: np.ndarray, from_mat: np.ndarray):
        to_rot = to_mat[:3, :3]
        from_rot = from_mat[:3, :3]
        rot_mat = to_rot @ from_rot.T

        trans_mat = to_mat[:3, 3] - from_mat[:3, 3]

        result = np.eye(4)
        result[:3, :3] = rot_mat
        result[:3, 3] = trans_mat
        result = np.where(np.abs(result) < 1e-5, 0, result)
        return result

    @staticmethod
    def trans_base(
            init_pose_mat: np.ndarray,
            now_base_mat: np.ndarray,
            init_base_mat: np.ndarray = np.eye(4),
    ):
        now_pose_mat = np.eye(4)
        base_trans_mat = Helper.trans_mat(now_base_mat, init_base_mat)
        now_pose_mat[:3, :3] = (base_trans_mat[:3, :3] @ init_pose_mat[:3, :3] @ base_trans_mat[:3, :3].T)
        now_pose_mat[:3, 3] = base_trans_mat[:3, :3] @ init_pose_mat[:3, 3]

        # 转化为世界坐标
        p = now_pose_mat[:3, 3] + now_base_mat[:3, 3]
        q_mat = now_pose_mat[:3, :3] @ now_base_mat[:3, :3]
        return sapien.Pose(p, t3d.quaternions.mat2quat(q_mat))

    def add_visual_box(self, pose: sapien.Pose, name: str = "box"):
        box, _ = create_obj(
            scene=self.scene,
            pose=pose,
            modelname="vis_box",
            # modelname="cube",
            is_static=True,
            scale=[0.025, 0.025, 0.025],
            no_collision=True,
        )
        box.set_name(name)

    def check_obj(self, name, idx, mid, d_range=50, pose=None, anno=None):
        if pose is None:
            obj, config = rand_create_actor(
                self.scene,
                xlim=[0, 0],
                ylim=[0, 0],
                zlim=[0.743],
                modelname=f"{idx}_{name}",
                model_id=mid,
                convex=True,
                qpos=[0, 0, 0.707107, 0.707107],
                scale=(0.1, 0.1, 0.1),
            )
        else:
            obj = create_actor(
                self.scene,
                pose=sapien.Pose(pose[:3], [0, 0, 0.707107, 0.707107]),
                modelname=f"{idx}_{name}",
                model_id=mid,
                convex=True,
                is_static=True,
            )
            if obj is None:
                print(f"create obj[{idx}_{name}/{mid}] failed")
                return False
            if anno is not None and (anno is True or (anno[0] <= pose[0] <= anno[1] and anno[2] <= pose[1] <= anno[3])):
                try:
                    scale = config["scale"]
                    base_mat = obj.get_pose().to_transformation_matrix()
                    for key, name in self.POINTS:
                        if key == "orientation_point":
                            if len(config.get(key, [])) <= 1:
                                continue
                            points = [config.get(key, [])]
                        else:
                            points = config.get(key, [])

                        for idx, mat in enumerate(points):
                            mat = np.array(mat)
                            mat[:3, 3] *= scale
                            pose = self.trans_base(mat, base_mat)
                            self.add_visual_box(pose, name=f"{name}_{idx}")
                except:
                    return False

        # def to_array(pose:sapien.Pose) -> np.ndarray:
        #     return np.array(pose.p.tolist()+pose.q.tolist())

        # is_step, max_step = 0, 200
        # pose_list = [to_array(obj.get_pose())]
        # while is_step < max_step:
        #     self.scene.step()
        #     self.scene.update_render()
        #     self.viewer.render()

        #     new_pose = obj.get_pose()
        #     pose_list.append(to_array(new_pose))

        #     if len(pose_list) > d_range:
        #         check_succ = True
        #         for i in range(-d_range, 0):
        #             if not np.allclose(pose_list[i], pose_list[-d_range], 1e-4):
        #                 check_succ = False
        #                 break
        #         if check_succ:
        #             break
        #     is_step += 1

        # if is_step > 0 and is_step < max_step:
        #     success = True
        # elif is_step >= max_step:
        #     success = False

        # self.scene.remove_actor(obj)
        # self.scene.update_render()
        # self.viewer.render()
        # return success

    def add_robot(self):

        def init_joints(entity: sapien.physx.PhysxArticulation, config):
            # set joints
            active_joints = entity.get_active_joints()
            arm_joints = [entity.find_joint_by_name(i) for i in config["arm_joints_name"][0]]

            def get_gripper_joints(find, gripper_name: str):
                gripper = [(find(gripper_name["base"]), 1.0, 0.0)]
                for g in gripper_name["mimic"]:
                    gripper.append((find(g[0]), g[1], g[2]))
                return gripper

            gripper = get_gripper_joints(entity.find_joint_by_name, config["gripper_name"][0])

            for i, joint in enumerate(active_joints):
                joint.set_drive_property(
                    stiffness=config.get("joint_stiffness", 1000),
                    damping=config.get("joint_damping", 200),
                )
            for joint in gripper:
                joint[0].set_drive_property(
                    stiffness=config.get("gripper_stiffness", 1000),
                    damping=config.get("gripper_damping", 200),
                )
            for i, joint in enumerate(active_joints):
                joint.set_drive_target(config["joints"][0][i])
            for i, joint in enumerate(gripper):
                real_joint: sapien.physx.PhysxArticulationJoint = joint[0]
                drive_target = config["gripper_scale"][1] * joint[1] + joint[2]
                drive_velocity_target = (np.clip(drive_target - real_joint.drive_target, -1.0, 1.0) * 0.05)
                real_joint.set_drive_target(drive_target)
                real_joint.set_drive_velocity_target(drive_velocity_target)

        radius = 2.5
        count, max_count = 0, 13
        emb = Path("./assets/embodiments")

        joint_dict = {
            "ARX-X5": [-6.155617, 1.1425792, 1.4179262, -0.97225964, -1.4429708e-05, -3.082031e-06, 0.044, 0.044],
            # "ARX-X5": [
            #     -6.155634,
            #     0.816421,
            #     1.0468683,
            #     -0.9384637,
            #     -3.4565306e-05,
            #     -8.612996e-06,
            #     0.044,
            #     0.044,
            # ],
            "piper": [
                -0.34990656,
                1.2450953,
                -1.5324507,
                0.10282991,
                1.22,
                0.00065908127,
                0.039999943,
                0.03997663,
            ],
            "franka-panda": [
                -0.00021794076,
                0.041278794,
                -0.0013123713,
                -1.8957008,
                0.009215873,
                2.0166128,
                0.8549956,
                0.04,
                0.04,
            ],
            "aloha-agilex": [
                0.0,
                0.0,
                -2.5302018e-14,
                -2.5302018e-14,
                -2.5302018e-14,
                -2.5302018e-14,
                1.1234251e-05,
                1.0832736e-05,
                -0.00048545605,
                1.5486969e-05,
                -2.5418809e-17,
                -2.5418809e-17,
                -2.5418809e-17,
                -2.5418809e-17,
                0.002626635,
                0.002626792,
                0.0027120241,
                0.0021979488,
                -0.0399116,
                -0.03991316,
                -0.031362604,
                -0.031362318,
                -0.0021148901,
                -0.002130989,
                -0.0031363545,
                -0.0031357573,
                -0.00090792944,
                -0.0009686581,
                -1.6246497e-06,
                -1.6584742e-06,
                -6.803319e-05,
                -6.932296e-05,
                1.0387723e-06,
                1.125215e-06,
                0.044976402,
                0.044976484,
                0.04762502,
                0.047625143,
            ],
            "ur5-wsg": [-1.5452573, -1.7434453, -1.3246999, -1.75, 1.5422482, -3.1415927, -0.055, 0.055],
            "z1": [
                0.2046731,
                1.5261446,
                -1.7666384,
                1.1484289,
                8.120951e-06,
                -7.348934e-05,
                -7.787227e-08,
                0.040000536,
                0.040000137,
            ],
            "ufactory_lite6": [
                -0.16042127,
                0.53086734,
                2.0658371,
                0.006172284,
                0.92715985,
                1.5044298,
                3.7193262e-05,
                0.040008515,
                0.03999608,
            ],
            # 'ARX-X5': [-6.1558957, 0.81342375, 1.0558599, -0.937343, -3.2896776e-05, -7.4935256e-06, 0.044, 0.044],
            # 'ufactory_lite6': [-0.25563017, 0.35529876, 2.0722473, 0.005538411, 0.9270778, 1.5045198, 3.138106e-05, 0.040005907, 0.03998957],
            # 'franka-panda': [-0.00016283647, 0.0074461037, -0.0010076275, -1.8719686, 0.008220577, 2.018346, 0.85500133, 0.04, 0.04],
            # 'aloha-agilex-1': [0.0, 0.0, -2.2630008e-14, -2.2630008e-14, -2.2630008e-14, -2.2630008e-14, 7.525569e-06, 7.171039e-06, -0.00035828358, 1.0665836e-06, -1.6544881e-17, -1.6544881e-17, -1.6544881e-17, -1.6544881e-17, 0.0017437901, 0.0017439453, 0.0018454427, 0.0014031429, -0.026461456, -0.02646318, -0.023825448, -0.023819776, -0.0014078408, -0.001425338, -0.0022293383, -0.0022303618, -0.0006110415, -0.00067730586, -9.73302e-07, -9.84728e-07, -4.5029174e-05, -4.632743e-05, 2.28171e-07, 2.6904584e-07, 0.044995338, 0.044996887, 0.04765, 0.04765],
            # 'ur5-wsg-gripper': [-1.5494769, -1.5602797, -1.3733442, -1.7500004, 1.5424018, -0.000120613506, -0.055, 0.055],
            # 'z1': [0.20469421, 1.5193493, -1.7742655, 1.1475929, 7.234784e-06, -5.6515753e-05, -1.4886399e-07, 0.039999034, 0.039998993],
            # 'piper': [-0.23096707, 1.2409755, -1.4549325, 0.10388685, 1.2199999, 0.000500452, 0.039999936, 0.03997564],
            # 'ur5-wsg-gripper': [-1.57, -0.78, -1.33, -1.70, 1.56, 3.14, 0, 0],
            # 'franka-panda': [-2.89, 1.03, 2.89, -1.93, -0.21, 1.27, 0.78, 0, 0],
        }

        pose_dict = {
            "ARX-X5": sapien.Pose([-0.821443, -1.6714, 0.781873], [0.999601, 5.86649e-08, -7.04128e-07, -0.0282362]),
            "piper": sapien.Pose([0.846021, -1.70083, 0.731933], [-0.0329616, 4.47035e-08, -1.49012e-08, 0.999457]),
            "franka-panda": sapien.Pose(
                [0.880834, -2.30439, 0.75],
                [0.4564, 4.47035e-08, -1.16415e-10, 0.889775],
            ),
            "aloha-agilex": sapien.Pose(
                [1.75423e-08, -2.39183, 0.465],
                [0.709881, 8.9407e-08, -4.74683e-08, 0.704321],
            ),
            "ur5-wsg": sapien.Pose(
                [-0.907954, -2.31459, 0.77098],
                [0.956935, -2.55658e-08, -1.0741e-07, -0.290302],
            ),
            # "z1": sapien.Pose(
            #     [-1.08728, -1.74981, 0.743286],
            #     [0.999384, -5.50994e-07, 4.56203e-09, 0.035102],
            # ),
            # "ufactory_lite6": sapien.Pose([-1.14082, -1.26895, 0.654833], [1, 0, 0, 8.66251e-07]),
            # 'ARX-X5': sapien.Pose([0.97715, -1.28326, 0.783988], [-9.09963e-07, 6.85768e-07, 9.68444e-08, 1]),
            # 'ufactory_lite6': sapien.Pose([1.11059, -1.72772, 0.654833], [-9.09963e-07, 0, 0, 1]),
            # 'franka-panda': sapien.Pose([0.880834, -2.30439, 0.75], [0.4564, 4.47035e-08, -1.16415e-10, 0.889775]),
            # 'aloha-agilex-1': sapien.Pose([8.77117e-09, -2.39183, 0.315], [0.709881, 8.9407e-08, -4.74683e-08, 0.704321]),
            # 'ur5-wsg-gripper': sapien.Pose([-0.907954, -2.31459, 0.77098], [0.956935, -2.55658e-08, -1.0741e-07, -0.290302]),
            # 'z1': sapien.Pose([-1.08728, -1.74981, 0.743286], [0.999384, -5.50994e-07, 8.28732e-09, 0.035102]),
            # 'piper': sapien.Pose([-1.02042, -1.22711, 0.731933], [1, 0, 0, 0]),
            # 'piper': sapien.Pose([-1.08728, -1.28692, 0.743286], [-0.99999, 5.36442e-07, 4.47035e-08, -0.00445333]),
            # 'ur5-wsg-gripper': sapien.Pose([1.11433, -2.06679, 0.75], [0.935868, 3.53431e-08, -1.04601e-07, 0.35235]),
            # 'z1': sapien.Pose([-1.08376, -0.879804, 0.75], [0.999983, 8.50705e-08, -2.12434e-07, 0.00582502]),
            # 'ufactory_lite6': sapien.Pose([1.10847, -0.967647, 0.6], [0.0313626, 0, 0, 0.999508]),
            # 'ARX-X5': sapien.Pose([1.072, -1.29, 0.783988], [0.0159038, 6.87222e-07, 8.59252e-08, 0.999874]),
            # 'franka-wsg': sapien.Pose([0.622937, -2.38999, 0.750243], [0.710587, 9.08971e-07, 9.36911e-07, 0.703609]),
            # 'ur5-robotic85-gripper': sapien.Pose([-1.38769, -1.8014, 0.75], [-0.879306, 5.93055e-08, 8.07258e-08, 0.476258]),
            # 'RM65B-EG24C2': sapien.Pose([1.02471, -1.10745, 0.85], [0.999792, 3.72529e-08, 2.26079e-07, -0.0204]),
            # 'ufactory_xarm7': sapien.Pose([-1.09916, -1.35288, 0.72], [0.998108, -4.47035e-08, -2.98023e-08, 0.0614921]),
            # 'piper': sapien.Pose([-1.08728, -0.786483, 0.743286], [0.99999, -2.98023e-08, 7.45058e-08, 0.00445369]),
            # 'z1': sapien.Pose([-1.08376, -0.536377, 0.75], [0.999983, 8.58563e-08, -2.12394e-07, 0.00582501]),
            # 'ufactory_lite6': sapien.Pose([1.10847, -0.601613, 0.6], [0.0313626, 0, 0, 0.999508]),
            # 'ARX-X5': sapien.Pose([1.072, -0.884706, 0.783988], [0.0159033, 6.87222e-07, 8.59257e-08, 0.999874]),
            # 'kinova': sapien.Pose([-1.29571, -1.1639, 0.618725], [0.999138, -2.0683e-06, 0.000454868, 0.0415031]),
            # 'rethink_robotics_sawyer': sapien.Pose([1.40007, -1.41166, 0.7], [0.0748435, -7.82311e-08, -4.47035e-08, 0.997195]),
        }
        name_list = list(pose_dict.keys())

        for robot in emb.iterdir():
            if not robot.is_dir():
                continue
            cfg_path = robot / "config.yml"
            if robot.name not in name_list:
                continue
            if not cfg_path.exists():
                continue

            cfg = yaml.load(open(cfg_path, "r", encoding="utf-8"), Loader=yaml.FullLoader)
            urdf_path = robot / cfg["urdf_path"]
            loader: sapien.URDFLoader = self.scene.create_urdf_loader()
            loader.fix_root_link = True
            entity: sapien.physx.PhysxArticulation = loader.load(str(urdf_path))
            entity.set_name(robot.name)
            print(f"load {robot.name} from {urdf_path}")
            # x = 0.1 - radius * np.cos(np.pi/12+(np.pi*5/6)/max_count*count)
            # entity.set_pose(sapien.Pose([
            #     x,
            #     -radius * np.sin(np.pi/12+(np.pi*5/6)/max_count*count),
            #     cfg['robot_pose'][0][2]
            #     ], t3d.quaternions.axangle2quat([0, 0, 1], np.pi/max_count*count)
            # ))
            cfg["joints"] = [joint_dict[robot.name]]
            entity.set_pose(pose_dict[robot.name])
            init_joints(entity, cfg)
            count += 1

    def block(self):
        if self.viewer is None:
            return
        while True:
            self.scene.step()
            self.scene.update_render()
            self.viewer.render()

    def run(self, step=200, no_step=False):
        if no_step:
            self.scene.update_render()
            if self.viewer is not None:
                self.viewer.render()
            return
        for _ in tqdm(range(step), desc="running"):
            if not no_step:
                self.scene.step()
            self.scene.update_render()
            if self.viewer is not None:
                self.viewer.render()

    def take_picture(self, name="camera.png"):
        print("start taking picture")
        self.camera.take_picture()
        camera_rgba = self.camera.get_picture("Color")
        position = self.camera.get_picture("Position")
        depth = -position[..., 2]

        camera_rgba_img = (camera_rgba * 255).clip(0, 255).astype("uint8")[:, :, :3]
        camera_rgba_img = Image.fromarray(camera_rgba_img)
        camera_rgba_img.save(name)
        np.save("depth_data.npy", depth)
        print("picture saved:", name)

    def generate_in(
        self,
        obj_list,
        x_min,
        x_max,
        y_min,
        y_max,
        z=0.74,
        padding=0.05,
        anno=None,
        table=True,
        max_z_stack=1,
        logo=False,
    ):
        max_z_count = 0

        def create_table():
            nonlocal max_z_count, table
            max_z_count += 1
            if max_z_count > max_z_stack:
                return False
            if not table:
                return True
            builder = self.scene.create_actor_builder()
            builder.set_physx_body_type("static")

            length, width, thickness = x_max - x_min, y_max - y_min, 0.02
            tabletop_pose = sapien.Pose([0.0, 0.0, -thickness / 2])  # Center the tabletop at z=0
            tabletop_half_size = [length / 2, width / 2, thickness / 2]
            builder.add_box_collision(
                pose=tabletop_pose,
                half_size=tabletop_half_size,
                material=self.scene.default_physical_material,
            )

            builder.add_box_visual(
                pose=tabletop_pose,
                half_size=tabletop_half_size,
                material=(1, 1, 1),
            )
            table = builder.build("table")
            table.set_pose(sapien.Pose(p=[(x_min + x_max) / 2, (y_min + y_max) / 2, z], q=[0, 0, 0, 1]))
            return True

        def load_logo():
            name = "rbt.glb"
            scale = (0.6, ) * 3
            builder = self.scene.create_actor_builder()
            builder.set_physx_body_type("static")
            builder.add_multiple_convex_collisions_from_file(filename=name, scale=scale)
            builder.add_visual_from_file(filename=name, scale=scale)
            mesh = builder.build(name="logo")
            mesh.set_pose(sapien.Pose([0, -1.37182, 0.991556], [-4.88642e-06, 4.02623e-06, 0.348843, 0.937181]))

        create_table()

        if logo:
            load_logo()
            y_max -= 0.5

        sum_x, sum_y, max_y, max_z = x_min, y_max - padding, 0, 0
        batch = []
        for cnt, (name, idx, mid, tagged, t, height) in enumerate(tqdm(obj_list)):
            if t == "obj":
                cfg = Path(f"./assets/objects/{idx}_{name}/model_data{mid}.json")
                if not cfg.exists():
                    print(f"WARNING: {idx}_{name}/{mid} not found")
                    continue
                with open(cfg, "r", encoding="utf-8") as f:
                    cfg = json.load(f)
                w, h, tall = (
                    cfg["extents"][0] * cfg["scale"][0],
                    cfg["extents"][2] * cfg["scale"][2],
                    cfg["extents"][1] * cfg["scale"][1],
                )
            else:
                w, h, tall, z_off = (
                    self.messy_item_info["radius"][f"{name}_{idx}"] * 2,
                    self.messy_item_info["radius"][f"{name}_{idx}"] * 2,
                    self.messy_item_info["z_max"][f"{name}_{idx}"] - self.messy_item_info["z_offset"][f"{name}_{idx}"],
                    self.messy_item_info["z_offset"][f"{name}_{idx}"],
                )
            if sum_y - padding - h < y_min or sum_x + padding > x_max or cnt == len(obj_list) - 1:
                for x, y, zz, (n, i, m, tg, tp, h) in tqdm(batch):
                    if tp == "obj":
                        success = self.check_obj(n, i, m, pose=[x, sum_y - max_y / 2, zz], anno=anno)
                    else:
                        success = self.check_urdf(n, i, d_range=50, pose=[x, sum_y - max_y / 2, zz])
                batch = []
                sum_y -= max_y + padding
                max_y = 0
                
                if sum_y - padding - h < y_min:
                    sum_y = y_max - padding
                    # z += max_z + 0.01
                    # z += 0.3
                    # z -= 0.3
                    z -= 0.3
                    x_min -= 0.25
                    x_max += 0.25
                    if not create_table():
                        return obj_list[cnt:]
                    max_z = 0
                sum_x = x_min

            if t == "obj":
                batch.append((sum_x + padding + w / 2, h, z, (name, idx, mid, tagged, t, h)))
            else:
                batch.append((
                    sum_x + padding + w / 2,
                    h,
                    z - z_off,
                    (name, idx, mid, tagged, t, h),
                ))

            sum_x += w + padding
            max_y = max(max_y, h)
            max_z = max(tall, max_z)
        return []

    def test_obj(self):
        try:
            self.create_scene(viewer=True)
            self.rendered = True
        except:
            self.create_scene(viewer=False)
            self.rendered = False
        self.create_table_and_wall()

        if self.viewer is not None:
            self.viewer.set_camera_pose(self.camera.get_pose())

        self.result = []
        test_list_1, test_list_2 = [], []
        for root_path in Path("./assets/objects").iterdir():
            if not root_path.is_dir():
                continue
            if re.search(r"^(\d+)_(.*)$", root_path.name) is None:
                continue

            new_list = []
            try:
                idx, name = root_path.name.split("_", 1)
                if name in ["dustbin", "tabletrashbin"]:
                    continue
                collision = [i.name for i in (root_path / "collision").iterdir()]
                visual = [i.name for i in (root_path / "visual").iterdir()]
                config = [i.name for i in (root_path).iterdir() if i.name.endswith(".json")]
                models = set(collision) & set(visual)
                for model in models:
                    modelid = re.search(r"(\d+)", model)
                    minz = 999.9
                    if modelid is not None:
                        modelid = int(modelid.group(1))
                        cfg = Path(f"./assets/objects/{idx}_{name}/model_data{modelid}.json")
                        if not cfg.exists():
                            print(f"WARNING: {idx}_{name}/{modelid} not found")
                            continue
                        with open(cfg, "r", encoding="utf-8") as f:
                            cfg = json.load(f)
                        # cfg["scale"] = cfg.get("scale", [0.1, 0.1, 0.1])
                        size = np.array(cfg["extents"]) * np.array(cfg["scale"])
                        minz = np.exp(min(max(2*(size[0]-size[1])**2, (size[0]-size[2])**2, 2*(size[1]-size[2])**2), minz))+size[1]
                        new_list.append([
                            name,
                            idx,
                            modelid,
                            f"model_data{modelid}.json" in config,
                            "obj",
                            0.0
                        ])
                new_list.sort(key=lambda x: x[2])
                for i in new_list: i[5] = minz

                if name in ['sauce-can', 'french-fries', 'hamburg', 'stapler', 'tea-box', 'coffee-box', 'tissue-box', 'bread', 'toycar', 'playingcards', 'small-speaker', 'cup']:
                    test_list_1 += new_list
                else:
                    test_list_2 += new_list
                # test_list += new_list
            except Exception as e:
                print(f"WARNING: [{name}_{idx}] failed:", e)

        # self.init_messy()
        # for name in self.obj_names:
        #     for idx in self.messy_item_info["list_of_items"][name]:
        #         test_list_2.append((
        #             name,
        #             idx,
        #             "",
        #             "",
        #             "urdf",
        #             self.messy_item_info["z_max"][f"{name}_{idx}"],
        #         ))

        # test_list_1.sort(key=lambda x: x[5])
        # np.random.seed(42)
        # np.random.shuffle(test_list_2)
        # test_list_2.sort(key=lambda x: x[5][1])
        # test_list_1 = test_list_1
        # test_list_2 = test_list_2

        # self.add_robot()
        # self.run(500)
        # self.block()
        print(f"{len(test_list_1)=}, {len(test_list_2)=}")
        # test_list_1 = test_list_1 + test_list_2
        test_list_1 = test_list_1 + test_list_2
        test_list_1.sort(key=lambda x: x[5])
        # np.random.shuffle(test_list_1)
        # test_list_1_idx = np.random.choice(np.arange(len(test_list_1)), size=250, replace=False)
        # test_list_1 = [i for idx, i in enumerate(test_list_1) if idx in test_list_1_idx]
        # test_list_1, test_list_2 = [], []

        # self.generate_in(test_list_1, -0.5, 0.5, -2.2, -1.1, z=1.0, logo=True)
        if self.rendered:
            test_list_1 = test_list_1[:50]
            print(test_list_1)
        # res = self.generate_in(test_list_1, -1.25, 1.25, -1.8, -0.4, table=False, max_z_stack=5)
        res = self.generate_in(test_list_1, -1.2, 1.2, -1.8, -0.4, table=False, max_z_stack=5)
        print('rest', len(res))
        # self.block()
        # list_2 = self.generate_in(test_list_2, 1.7, 2.4, -3, -0.5, 0.2, max_z_stack=5)
        # list_2 = self.generate_in(list_2, -2.4, -1.7, -3, -0.5, 0.2, max_z_stack=5)
        # list_2 = self.generate_in(list_2, -2.4, 2.4, -4, -3.5, 0.4, max_z_stack=5)

        # list_2 = self.generate_in(test_list_2[:2], 1.7, 2.4, -3, -0.5, 0.2, max_z_stack=1)
        # list_2 = self.generate_in(test_list_2[:2], -2.4, -1.7, -3, -0.5, 0.2, max_z_stack=1)
        # list_2 = self.generate_in(test_list_2[:2], -2.4, 2.4, -4, -3.5, 0.4, max_z_stack=1)
        # self.block()
        # x_max = 1
        # sum_x = -x_max
        # sum_y, max_y = -0.2, 0
        # padding = 0.05
        # for cnt, (name, idx, mid, tagged) in enumerate(tqdm(test_list_1)):
        #     cfg = Path(f'./assets/objects/{idx}_{name}/model_data{mid}.json')
        #     if not cfg.exists():
        #         print(f'WARNING: {idx}_{name}/{mid} not found')
        #         continue
        #     with open(cfg, 'r', encoding='utf-8') as f:
        #         cfg = json.load(f)
        #     w, h = cfg['extents'][0] * cfg['scale'][0], cfg['extents'][2] * cfg['scale'][2]
        #     if sum_x + padding + w > x_max:
        #         sum_x = -x_max
        #         sum_y -= max_y + padding
        #         max_y = 0

        #     success = self.check_obj(name, idx, mid, pose=[sum_x+padding+w/2, sum_y-h/2, 0.743])
        #     sum_x += w + padding
        #     max_y = max(max_y, h)

        #     self.result.append({
        #         'name': f'{idx}_{name}',
        #         'id': mid,
        #         'tagged': tagged,
        #         'stable': success
        #     })
        # with open('result.jsonl', 'a', encoding='utf-8') as f:
        #     f.write(json.dumps(self.result[-1]) + '\n')
        # with open('success.txt', 'a', encoding='utf-8') as f:
        #     f.write(
        #         f'{idx}_{name:<15}/{mid:2d} 标定：{"是" if tagged else "否"}  稳定：{"是" if success else "否"}\n')

        # if cnt > 0 and cnt % 20 == 0:
        #     self.scene.clear()
        #     time.sleep(2)
        #     self.create_table_and_wall()
        # self.block()
        self.run(no_step=True)
        self.take_picture('./script/camera.png')
        self.block()


import os


def cpy():
    models = []
    with open("./success.txt", "r", encoding="utf-8") as f:
        lines = [i.strip() for i in f.readlines()]
        for i in lines:
            i_split = i.split("_")
            models.append(("_".join(i_split[:-1]), i_split[-1], i))

    for name, idx, original_name in tqdm(models):
        from_path = Path(f"./assets/messy_objects/{original_name}")
        to_path = Path(f"./assets/messy_objects_stable/{original_name}")
        os.system(f"cp -r {from_path} {to_path}")

    with open("./assets/messy_objects/list.json", "r", encoding="utf-8") as f:
        metadata = json.load(f)

    list_of_items = {}
    for name, idx, original_name in models:
        if name not in list_of_items:
            list_of_items[name] = []
        list_of_items[name].append(idx)

    new_metadata = {
        "item_names": list(set([n[0] for n in models])),
        "list_of_items": list_of_items,
        "radius": {
            n[2]: metadata["radius"][n[2]]
            for n in models
        },
        "z_offset": {
            n[2]: metadata["z_offset"][n[2]]
            for n in models
        },
        "z_max": {
            n[2]: metadata["z_max"][n[2]]
            for n in models
        },
    }

    with open("./assets/messy_objects_stable/list.json", "w", encoding="utf-8") as f:
        json.dump(new_metadata, f, ensure_ascii=False, indent=4)


def cfg():
    result = []
    with open("result.jsonl", "r", encoding="utf-8") as f:
        for i in f.readlines():
            result.append(json.loads(i.strip()))
    root_path = Path("./assets/objects")
    for res in result:
        res_cfg = root_path / res["name"] / f'model_data{res["id"]}.json'
        if res_cfg.exists():
            with open(res_cfg, "r", encoding="utf-8") as f:
                cfg = json.load(f)
            cfg["stable"] = res["stable"]
            with open(res_cfg, "w", encoding="utf-8") as f:
                json.dump(cfg, f, ensure_ascii=False, indent=4)
        else:
            print(f'WARNING: {res["name"]}/{res["id"]} not found')

if __name__ == "__main__":
    helper = Helper()
    helper.test_obj()
    # cpy()
    # cfg()
#     pass

# all_items = [
#     "bottle", "bowl", "brush", "can", "chip_can", "clock", "drinkbox", "hammer", "marker", "notebook", "pencil", "plate", "pot", "ramen_box", "remote", "slipper", "snack_box", "snack_package", "sneaker", "spoon", "steel_tape", "tape", "thermos", "tissue", "toothbrush", "toy_car", "wallet",
    
#     "001_bottle", "002_bowl", "003_plate", "004_fluted-block", "007_shoe_box", "019_coaster", "020_hammer", "021_cup", "022_cup-with-liquid", "027_table-tennis", "028_dustpan", "030_drill", "032_screwdriver", "033_fork", "034_knife", "035_apple", "036_cabinet", "037_box", "039_mug", "040_rack", "041_shoe", "042_wooden_box", "043_book", "045_sand-clock", "046_alarm-clock", "047_mouse", "048_stapler", "049_shampoo", "050_bell", "051_candlestick", "052_dumbbell", "053_teanet", "054_baguette", "055_small-speaker", "057_toycars", "058_markpen", "059_pencup", "061_battery", "062_plasticbox", "063_tabletrashbin", "068_boxdrink", "069_vagetables", "070_paymentsign", "071_cans", "072_electronicscale", "073_rubikscube", "074_displaystand", "075_bread", "076_breadbasket", "077_phone", "078_phonestand", "079_remotecontrol", "080_pillbottle", "081_playingcards", "082_smallshovel", "083_brush", "084_woodenmallet", "085_gong", "086_woodenblock", "087_waterer", "088_wineglass", "089_globe", "090_trophy", "091_kettle", "092_notebook", "093_brush-pen", "094_rest", "095_glue", "096_cleaner", "097_screen", "098_speaker", "099_fan", "100_seal", "101_milk-tea", "103_fruits", "104_board", "105_sauce-can", "106_skillet", "108_block", "110_basket", "111_callbell", "112_tea-box", "113_coffee-box", "109_hydrating-oil", "107_soap", "102_roller", "067_steamer", "066_vinegar", "065_soy-sauce", "064_msg", "060_kitchenpot", "056_switch", "044_microwave", "038_milk-box", "031_jam-jar", "029_olive-oil", "028_roll-paper", "026_pet-collar", "025_chips-tub", "024_scanner", "023_tissue-box", "018_microphone", "017_calculator", "016_oven", "015_laptop", "014_bookcase", "013_dumbbell-rack", "012_plant-pot", "011_dustbin", "010_pen", "009_kettle", "008_tray", "006_hamburg", "005_french-fries"
# ]

# # # Stricter groups based on high visual/shape similarity
# strict_groups = [
#     # Groups from the user's example
#     ["plate", "003_plate"],
#     ["toy_car", "057_toycars"],
#     ["remote", "079_remotecontrol"],
#     ["marker", "058_markpen", "pencil"],
#     ["can", "071_cans"],
#     ["mug", "039_mug"],

#     # Additional groups based on clear, direct similarities,
#     # trying to match the style and granularity of the examples.

#     # Direct L1 to L2 counterparts or very similar items
#     ["bottle", "001_bottle"],
#     ["bowl", "002_bowl"],
#     ["brush", "083_brush"],
#     ["clock", "046_alarm-clock"], # "alarm-clock" is a type of clock
#     ["hammer", "020_hammer"],
#     ["notebook", "092_notebook"], # L1 "notebook" matches L2 "notebook"
#     ["pot", "060_kitchenpot"], # "kitchenpot" is a type of pot
#     ["tissue", "023_tissue-box"], # "tissue" and "tissue-box" are directly related

#     # Small groups of highly similar items from L1 and/or L2
#     ["chip_can", "025_chips-tub"], # Both are tube-shaped snack containers
#     ["slipper", "sneaker", "041_shoe"], # All are types of footwear
#     ["spoon", "033_fork", "034_knife"], # Cutlery items
#     ["steel_tape", "tape"], # Types of tape
#     ["drinkbox", "068_boxdrink"], # Drinks in boxes
#     ["091_kettle", "009_kettle"], # Both are kettles
#     ["055_small-speaker", "098_speaker"], # Types of speakers
#     ["035_apple", "103_fruits"], # Fruits
#     ["054_baguette", "075_bread"], # Types of bread
#     ["050_bell", "111_callbell"], # Types of bells
#     ["004_fluted-block", "086_woodenblock", "108_block"], # Types of blocks

#     # Grouping L2 items that are similar concepts, or direct matches not yet covered
#     ["021_cup", "022_cup-with-liquid"], # Cups
#     ["112_tea-box", "113_coffee-box"], # Specific types of boxes
#     ["037_box", "042_wooden_box", "062_plasticbox", "007_shoe_box"], # General boxes

#     ["019_coaster", "008_tray"],
#     ["036_cabinet", "040_rack", "014_bookcase", "013_dumbbell-rack"], # Furniture/storage units
#     ["063_tabletrashbin", "011_dustbin"], # Trash receptacles
#     ["074_displaystand", "078_phonestand"],
#     ["110_basket", "076_breadbasket"],
# ]

# similar_items_dict = {item: set() for item in all_items}

# for group in strict_groups:
#     # Ensure all items in defined groups are known (they should be from all_items)
#     for item_in_group in group:
#         if item_in_group not in similar_items_dict:
#             # This case should ideally not happen if groups only contain items from all_items
#             print(f"Warning: Item '{item_in_group}' in a group is not in the master list of all_items.")
#             continue # Skip if item not in master list, or handle as error

#     for item_in_group in group:
#         if item_in_group in similar_items_dict: # Process only if item is in master list
#             # Add all other items from this specific group as similar
#             for other_item_in_group in group:
#                 if item_in_group != other_item_in_group:
#                     similar_items_dict[item_in_group].add(other_item_in_group)

# # Convert sets to sorted lists for consistent output
# final_strict_dict = {item: sorted(list(similar_set)) for item, similar_set in similar_items_dict.items()}

# # Example of how to print the resulting dictionary (optional)
# # import json
# # print(json.dumps(final_strict_dict, indent=2))

# # To display the dictionary (optional, for verification)
# import json
# json.dump(final_strict_dict, open('similar_items.json', 'w'), indent=4)
