import logging
import argparse
import readline
import importlib
import numpy as np
from pathlib import Path
from copy import deepcopy
import transforms3d as t3d
from threading import Thread, Lock
import trimesh
import trimesh.bounds

import sys

sys.path.append(".")
from envs.utils import *

import sapien
from sapien.render import set_global_config

render_pause = False


class BaseViewer:
    scene: sapien.Scene
    viewer: sapien.utils.Viewer

    actor: Actor
    modelid: str
    modelname: str
    config_path: Path
    EMPTY_CONFIG: dict
    POINTS: list[tuple[str, str]]

    def __init__(self):
        # create scene and viewer
        set_global_config(max_num_materials=50000, max_num_textures=50000)
        self.scene = sapien.Scene()
        self.scene.set_timestep(1 / 250)

        # initialize viewer with camera position and orientation
        self.viewer = None
        self.reset()
    
    def open_viewer(self):
        if self.viewer is not None and not self.viewer.closed:
            return 
        self.viewer = self.scene.create_viewer()
        self.viewer.set_scene(self.scene)
        self.viewer.set_camera_pose(pose=sapien.Pose(
            [-0.0096987, -0.19846, 0.0955636],
            [0.71241, -0.118063, 0.123576, 0.680634],
        ))

    def reset(self):
        self.scene.clear()
        self.open_viewer()

        # ground
        self.scene.add_ground(0)

        # lights
        self.scene.set_ambient_light([0.5, 0.5, 0.5])
        shadow = True
        # default spotlight angle and intensity
        direction_lights = [[[0, 0.5, -1], [0.5, 0.5, 0.5]]]
        for direction_light in direction_lights:
            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]]]
        for point_light in point_lights:
            self.scene.add_point_light(point_light[0], point_light[1], shadow=shadow)

        self.update_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 = BaseViewer.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", type: str = "cube"):
        global render_pause
        modelname = {
            # 'functional': 'functional.glb',
            # 'contact': 'gripper.glb'
        }.get(type, 'base.glb')
        modelpath = Path("assets/objects/vis_box") / modelname

        builder = self.scene.create_actor_builder()
        builder.set_physx_body_type("static")
        builder.add_visual_from_file(filename=str(modelpath))
        builder.set_initial_pose(pose)
        builder.set_name(name)
        render_pause = True
        builder.build()
        render_pause = False

    def clear_scene(self):
        global render_pause
        render_pause = True
        self.scene.clear()
        render_pause = False
        self.update_render()

    def update_render(self):
        global render_pause
        if not render_pause and not self.viewer.closed:
            self.scene.update_render()
            self.viewer.render()

    def save_config(self):
        with open(self.config_path, "w", encoding="utf-8") as f:
            json.dump(self.actor.config, f, ensure_ascii=False, indent=4)
        logging.info(f"Config saved to {self.config_path}")

    def main(self, pose, modelname, modelid, inherit_config: dict = None):
        ...

    def load_actor(self, pose, inherit_config):
        ...

    def update_config(self):
        ...

    def visualize(self):
        ...


class ObjectViewer(BaseViewer):
    EMPTY_CONFIG = {
        "center": [],   # Center Point
        "extents": [],  # Bounding Box Extents
        "scale": [1.0, 1.0, 1.0],  # Scale
        "transform_matrix": np.eye(4).tolist(),  # Model to Axis Rotation Matrix, fixed as Identity Matrix

        # Target Point Matrix (multiple), special points that can be obtained during planning (e.g., cup handle)
        "target_pose": [],
        # Grasping Point Matrix (multiple), grasping points are the positions where the robotic arm grasps the object (e.g., cup mouth)
        "contact_points_pose": [],
        # Functional Point Matrix (multiple), functional points are the positions where the object interacts with other objects (e.g., hammer head)
        "functional_matrix": [],
        # Orientation Point Matrix (single), orientation points specify the orientation of the object (e.g., shoe head facing left)
        "orientation_point": [],
        # Grasping Point Groups (same group should have the same position, different directions)
        "contact_points_group": [],
        # The number should be the same as the number of groups, should be set to true
        "contact_points_mask": [],
        
        "target_point_description": [],      # Target Point Description
        "contact_points_description": [],    # Grasping Point Description
        "functional_point_description": [],  # Functional Point Description
        "orientation_point_description": [], # Orientation Point Description
    }
    POINTS = [
        ("target_pose", "target"),
        ("contact_points_pose", "contact"),
        ("functional_matrix", "functional"),
        ("orientation_point", "orientation"),
    ]

    def __init__(self):
        super().__init__()

    def main(self, pose, modelname, modelid, inherit_config: dict = None):
        global render_pause
        self.modelid = modelid
        self.modelname = modelname
        self.reset()
        self.load_actor(pose, inherit_config)
        self.visualize()

        self.active = True

        def render():
            while self.active:
                self.update_render()
            self.clear_scene()

        self.render = Thread(target=render)
        self.render.start()
        self.console()
        self.active = False
        self.render.join()

    def __del__(self):
        self.active = False
        if hasattr(self, 'render'):
            self.render.join()
        self.scene.clear()
        if self.viewer is not None and not self.viewer.closed:
            self.viewer.close()

    def load_actor(self, pose, inherit_config=None, inherit_type: Literal['force', 'advice'] = 'advice'):
        modeldir = Path("assets/objects") / self.modelname
        modelid = '' if self.modelid is None else self.modelid
        self.config_path = modeldir / f"model_data{modelid}.json"

        # try to load as glb
        collision = modeldir / "collision" / f"base{modelid}.glb"
        visual = modeldir / "visual" / f"base{modelid}.glb"
        if not collision.exists() or not visual.exists():
            # try to load as obj
            collision = modeldir / "collision" / f"textured{modelid}.obj"
            visual = modeldir / "visual" / f"textured{modelid}.obj"
            
            if not collision.exists() or not visual.exists():
                logging.error(
                    f"Model files not found in {modeldir}({modelid}): "
                    f"collision {collision.exists()}, visual {visual.exists()}"
                )
                return False

        if self.config_path.exists():
            try:
                actor_config = json.load(open(self.config_path, "r", encoding="utf-8"))
                if len(actor_config['orientation_point']) > 1:
                    actor_config['orientation_point'] = [actor_config['orientation_point']]
            except json.JSONDecodeError:
                logging.warning(f"Invalid JSON in {self.config_path}, using empty config.")
                actor_config = None
        else:
            actor_config = None

        if actor_config is None:
            if inherit_config is None:
                actor_config = deepcopy(self.EMPTY_CONFIG)
            else:
                actor_config = deepcopy(inherit_config)
        else:
            if inherit_config is not None and inherit_type == 'force':
                actor_config = deepcopy(inherit_config)
            else:
                actor_config = actor_config
        actor_config.update(self.get_shape_data(collision))

        scale = actor_config['scale']
        builder = self.scene.create_actor_builder()
        builder.set_physx_body_type("dynamic")
        builder.add_multiple_convex_collisions_from_file(filename=str(collision), scale=scale)
        builder.add_visual_from_file(filename=str(visual), scale=scale)
        builder.set_initial_pose(pose)
        mesh = builder.build(name=f'{self.modelname}({self.modelid})')
        self.actor = Actor(mesh, actor_config)
        return True

    def get_shape_data(self, modelpath: Path):
        with open(modelpath, "rb") as file_obj:
            mesh: trimesh.Geometry = trimesh.load(
                file_obj, file_type=modelpath.suffix.strip("."))

        box: trimesh.primitives.Box = mesh.bounding_box_oriented
        return {
            "center": box.centroid.tolist(),  # Center Point
            "extents": box.extents.tolist(),  # Bounding Box Extents
        }

    def visualize(self):
        for key, name in self.POINTS:
            for idx in range(len(self.actor.config.get(key, []))):
                self.add_visual_box(pose=self.actor.get_point(name, idx, 'pose'), name=f"{name}_{idx}", type=name)
        self.update_render()

    def update_config(self, save: bool = False):
        config = deepcopy(self.EMPTY_CONFIG)
        config.update(self.actor.config)
        for key, _ in self.POINTS:
            config[key] = []

        actor_mat = self.actor.get_pose().to_transformation_matrix()

        def get_mat(entity: sapien.Entity):
            nonlocal config, actor_mat
            mat = entity.get_pose().to_transformation_matrix()
            p = actor_mat[:3, :3].T @ (mat[:3, 3] - actor_mat[:3, 3])
            mat[:3, 3] = p / config["scale"]
            mat[:3, :3] = actor_mat[:3, :3].T @ mat[:3, :3]
            return np.around(mat, 5)

        for entity in self.scene.get_all_actors():
            for key, name in self.POINTS:
                if entity.get_name().startswith(name):
                    config[key].append(get_mat(entity).tolist())

        self.actor.config = config
        if save:
            self.save_config()

    def reset_scale(self, scale):
        if len(scale) != 3:
            scale = [scale[0], scale[0], scale[0]]
        self.actor.config["scale"] = scale
        self.update_config()
        logging.info("Reloading scene, please wait for about 10 seconds...")
        self.reset()
        self.load_actor(self.actor.get_pose(), self.actor.config)
        self.visualize()

    @staticmethod
    def parse_point(cmd: str, req_id: bool = True):
        parse_map = {
            'c': 'contact',
            't': 'target',
            'f': 'functional',
            'o': 'orientation',
            'contact': 'contact',
            'target': 'target',
            'functional': 'functional',
            'orientation': 'orientation'
        }
        if cmd.strip() == '':
            cmd = input("  >> (t)arget, (c)ontact, (f)unctional, (o)rientation:")
        cmd = cmd.strip().split(" ")
        if req_id:
            try:
                type, pid = parse_map[cmd[0]], int(cmd[1])
            except (IndexError, ValueError, KeyError):
                return None, None
            return type, pid
        else:
            if len(cmd) != 1:
                return None
            return parse_map.get(cmd[0], None)

    def get_points(self, type: str) -> list[sapien.Entity]:
        points = []
        for entity in self.scene.get_all_actors():
            if entity.get_name().startswith(type):
                points.append(entity)
        return points

    def get_next_id(self, type: str):
        points = self.get_points(type)
        max_id = -1
        for p in points:
            max_id = max(max_id, int(p.get_name().split("_")[-1]))
        return max_id + 1

    def console(self):
        global render_pause
        modified = 0
        try:
            while not self.viewer.closed:
                cmd = input("Input command: ")
                if self.viewer.closed:
                    logging.warning("Viewer has been closed manually.")
                    cmd = input("Please choose to reopen, exit with save or exit without save: (r/s/e) ")
                    cmd = cmd.strip().lower()
                    if cmd in ['r', 'reopen']:
                        self.open_viewer()
                    if cmd in ['s', 'save']:
                        self.update_config(True)
                        break
                    if cmd in ['e', 'exit']:
                        break

                modified += 1
                if cmd == "save":
                    self.update_config(True)
                    modified = 0
                elif cmd[:6] == "resize":
                    """
                    Usage:
                        resize <x_size> <y_size> <z_size>: Set scaling along x, y, z axes
                        resize <size>: Uniformly scale all three axes
                    Example:
                        resize 0.1
                    """
                    args = cmd[7:].strip().split(" ")
                    if len(args) == 1:
                        size = float(args[0])
                        self.reset_scale((size, size, size))
                    elif len(args) == 3:
                        x_size = float(args[0])
                        y_size = float(args[1])
                        z_size = float(args[2])
                        self.reset_scale((x_size, y_size, z_size))
                    modified = 0
                elif cmd[:6] == "create":
                    """
                    Usage:
                        create <type>: Create (t)arget, (c)ontact, (f)unctional, (o)rientation point
                        create: Waits for input of point name
                    Example:
                        create t
                        create f
                    """
                    type = self.parse_point(cmd[6:], req_id=False)
                    if type is None:
                        logging.warning("Invalid type.")
                        continue
                    pid = self.get_next_id(type)
                    if type == "orientation" and pid > 0:
                        logging.warning("The orientation point is unique, please modify the existing one.")
                    else:
                        self.add_visual_box(self.actor.get_pose(), name=f"{type}_{pid}", type=type)
                        logging.info(f"Successfully created {type}_{pid}")
                elif cmd[:5] == "clone":
                    """
                    Usage:
                        clone <type> <id>: Clone a specified type and ID point in place
                        clone: Waits for input of point type and ID
                    Example:
                        clone t 1: Clones target_1 to create a new target point (e.g., target_2)
                    """
                    type, idx = self.parse_point(cmd[5:], req_id=True)
                    if type is None or idx is None:
                        logging.warning("Invalid type or id.")
                        continue
                    for entity in self.scene.get_all_actors():
                        if entity.get_name() == cmd and "_" in cmd:
                            type = cmd.split("_")[0]
                            if type == "orientation":
                                logging.warning("Orientation point is unique, cloning not supported!")
                            else:
                                pid = self.get_next_id(type)
                                self.add_visual_box(entity.get_pose(), name=f"{type}_{pid}", type=type)
                                logging.info(f"Successfully cloned {type}_{idx} to {type}_{pid}")
                elif cmd[:6] == "rotate":
                    """
                    Usage:
                        rotate <id> <axis> <interval>: Rotate a specified contact point around its own axis by a given interval, generating points belonging to the same group
                    Example:
                        rotate 1 x 90: Rotates contact_1 around its x-axis every 90 degrees, creating three additional contact points, and writes the group into concat_points_group
                    """
                    cmd = cmd[6:].strip().split(" ")
                    if len(cmd) != 3:
                        continue

                    name, axis, interval = cmd
                    if not name.isdigit():
                        logging.warning("Invalid id, must be a number.")
                        continue
                    name = int(name)

                    axis = {
                        "x": [1, 0, 0],
                        "y": [0, 1, 0],
                        "z": [0, 0, 1],
                        "r": [1, 0, 0],
                        "g": [0, 1, 0],
                        "b": [0, 0, 1],
                    }.get(axis, None)
                    if axis is None:
                        continue
                    axis = np.array(axis)

                    interval = int(interval)
                    for entity in self.scene.get_all_actors():
                        if entity.get_name() == f"contact_{name}":
                            pos_mat = entity.get_pose().to_transformation_matrix()
                            group_list = [name]
                            for i in range(interval, 360, interval):
                                pid = self.get_next_id("contact")
                                new_mat = np.eye(4)
                                new_mat[:3, 3] = pos_mat[:3, 3]

                                new_rot = axis * np.radians(i)
                                rot_mat = t3d.euler.euler2mat(new_rot[0], new_rot[1], new_rot[2])
                                new_mat[:3, :3] = pos_mat[:3, :3] @ rot_mat

                                pose = sapien.Pose(
                                    new_mat[:3, 3],
                                    t3d.quaternions.mat2quat(new_mat[:3, :3]),
                                )
                                self.add_visual_box(pose, name=f"contact_{pid}", type="contact")
                                group_list.append(pid)
                            self.actor.config["contact_points_group"].append(group_list)
                            self.actor.config["contact_points_mask"].append(True)
                            logging.info(f"Successfully rotated contact_{name} group, created {len(group_list)} points")
                            break
                elif cmd[:5] == "align":
                    """
                    Usage:
                        align: Aligns all group points' positions to the first point in the group
                    """
                    concat_points = {int(i.get_name().split("_")[-1]): i for i in self.get_points("contact")}
                    for group in self.actor.config["contact_points_group"]:
                        base = concat_points[group[0]]
                        base_p = base.get_pose().p
                        for pt in group[1:]:
                            concat_points[pt].set_pose(sapien.Pose(base_p, concat_points[pt].get_pose().q))
                        logging.info(f"Aligning group with {', '.join([str(i) for i in group[1:]])} to {group[0]}")
                elif cmd[:6] == "remove":
                    """
                    Usage:
                        remove <type> <id>: Removes a point with the specified name
                        remove: Waits for input of point name
                    Example:
                        remove t 0
                    """
                    type, idx = self.parse_point(cmd[6:], req_id=True)
                    if type is None or idx is None:
                        logging.warning("Invalid type or id.")
                        continue

                    for entity in self.scene.get_all_actors():
                        if entity.get_name() == f'{type}_{idx}':
                            render_pause = True
                            self.scene.remove_actor(entity)
                            render_pause = False
                            logging.info(f"Successfully removed {type}_{idx}")
                            break
                elif cmd == "exit":
                    if modified > 1:
                        cmd = input(
                            f'You have made {modified-1} changes without save, do you want to save them? (y/n/others to abort)'
                        )
                        if cmd.strip().lower() == 'y':
                            self.update_config(True)
                            break
                        elif cmd.strip().lower() == 'n':
                            break
                        else:
                            logging.info("Operation has been aborted.")
                    else:
                        break
                else:
                    modified -= 1
        except KeyboardInterrupt:
            pass


class URDFViewer(BaseViewer):
    EMPTY_CONFIG = {
        "scale": 1.0,  # Scale
        "transform_matrix": np.eye(4).tolist(),
        # Expected loading position to model actual pose transformation matrix, for example:
        # transform_matrix @ cube's bottom center point pose = cube pose
        "init_qpos": [],  # Initial joint state
        # Marker point matrix (multiple), marker points are special points that can be accessed during planning (e.g., cup handle)
        "target_points": [],
        # Grasping point matrix (multiple), grasping points are the positions where the robotic arm grasps the object (e.g., cup mouth)
        "contact_points": [],
        # Functional point matrix (multiple), functional points are the positions where the object interacts with other objects (e.g., hammer head)
        "functional_points": [],
        # Orientation point matrix (single), orientation points specify the orientation of the object (e.g., shoe head facing left)
        "orientation_point": [],
        # Grasping point groups (same group should have the same position, different directions)
        "contact_points_group": [],
        # The number should be the same as the number of groups, should be set to true
        "contact_points_mask": [],

        "target_points_description": [],      # Marker point description
        "contact_points_description": [],     # Grasping point description
        "functional_points_description": [],  # Functional point description
        "orientation_point_description": [],  # Orientation point description
    }
    """
        For each point:
        {
            "matrix": np.eye(4)    # 4x4 transformation matrix
            "base"  : "base_name"  # Base link name
        }
    """
    POINTS = [
        ("target_points", "target"),
        ("contact_points", "contact"),
        ("functional_points", "functional"),
        ("orientation_point", "orientation"),
    ]

    def __init__(self):
        super().__init__()

    def main(self, pose, modelname, modelid, inherit_config: dict = None):
        self.modelid = modelid
        self.modelname = modelname
        self.reset()
        self.load_actor(pose, inherit_config)
        self.visualize()

        self.active = True

        def render():
            while self.active and not self.viewer.closed:
                self.update_render()
            self.clear_scene()

        self.render = Thread(target=render)
        self.render.start()
        self.console()
        self.active = False
        self.render.join()

    def __del__(self):
        self.active = False
        if hasattr(self, 'render'):
            self.render.join()
        self.scene.clear()
        if self.viewer is not None and not self.viewer.closed:
            self.viewer.close()

    def load_actor(self, pose, inherit_config=None, inherit_type: Literal['force', 'advice'] = 'advice'):
        modeldir = Path("assets/objects") / self.modelname / str(self.modelid)
        self.config_path = modeldir / f"model_data.json"

        if self.config_path.exists():
            try:
                actor_config = json.load(open(self.config_path, "r", encoding="utf-8"))
            except json.JSONDecodeError:
                logging.warning(f"Invalid JSON in {self.config_path}, using empty config.")
                actor_config = None
        else:
            actor_config = None

        if actor_config is None:
            if inherit_config is None:
                actor_config = deepcopy(self.EMPTY_CONFIG)
            else:
                actor_config = deepcopy(inherit_config)
        else:
            if inherit_config is not None and inherit_type == 'force':
                actor_config = deepcopy(inherit_config)
            else:
                actor_config = actor_config

        loader: sapien.URDFLoader = self.scene.create_urdf_loader()
        loader.scale = actor_config["scale"]
        loader.fix_root_link = False
        loader.load_multiple_collisions_from_file = True
        actor: sapien.physx.PhysxArticulation = loader.load_multiple(str(modeldir / "mobility.urdf"))[0][0]
        actor.set_name(f"{self.modelname}({self.modelid})")
        actor.set_pose(self.get_real_pose(pose, np.array(actor_config.get("transform_matrix", np.eye(4)))))

        self.actor = ArticulationActor(actor, actor_config)
        for joint in self.actor.actor.get_joints():
            joint.set_drive_properties(
                damping=1000,
                stiffness=0,
            )
        if (self.actor.config.get("init_qpos") is not None and len(self.actor.config["init_qpos"]) > 0):
            self.actor.set_qpos(np.array(self.actor.config["init_qpos"]))
        return True

    def get_real_pose(self, pose: sapien.Pose, trans_matrix):
        pose_matrix = pose.to_transformation_matrix()
        return sapien.Pose(
            p=pose_matrix[:3, 3] + trans_matrix[:3, 3],
            q=t3d.quaternions.mat2quat(trans_matrix[:3, :3] @ pose_matrix[:3, :3]),
        )

    def visualize(self):
        for key, name in self.POINTS:
            for idx in range(len(self.actor.config.get(key, []))):
                self.add_visual_box(pose=self.actor.get_point(name, idx, 'pose'),
                                    name=f"{name}_{idx}<{self.actor.config[key][idx]['base']}>",
                                    type=name)
        self.update_render()

    def get_link(self, link_name: str):
        for link in self.actor.actor.get_links():
            if link.get_name() == link_name:
                return link
        return self.actor.actor

    def get_link_dict(self):
        link_dict = {}
        for link in self.actor.actor.get_links():
            link_dict[link.get_name()] = link
        return link_dict

    def get_base_name(self, point_name: str):
        res = re.search(r'(.*?)<(.*?)>', point_name)
        return res.group(2) if res else None

    def get_id(self, point_name: str):
        res = re.search(r'_(\d+)', point_name)
        return int(res.group(1)) if res else None

    def update_config(self, save: bool = False):
        config = deepcopy(self.EMPTY_CONFIG)
        config.update(self.actor.config)
        for key, _ in self.POINTS:
            config[key] = []

        link_dict = self.get_link_dict()

        def get_mat(entity: sapien.Entity, base: str):
            nonlocal config, link_dict
            mat = entity.get_pose().to_transformation_matrix()
            base_link = link_dict.get(base, self.actor)
            base_mat = base_link.get_pose().to_transformation_matrix()

            p = base_mat[:3, :3].T @ (mat[:3, 3] - base_mat[:3, 3])
            mat[:3, 3] = p / config["scale"]
            mat[:3, :3] = base_mat[:3, :3].T @ mat[:3, :3]
            return np.around(mat, 5)

        for entity in self.scene.get_all_actors():
            e_name = entity.get_name()
            for key, name in self.POINTS:
                if e_name.startswith(name):
                    base_name = self.get_base_name(e_name)
                    config[key].append({
                        "matrix": get_mat(entity, base_name).tolist(),
                        "base": base_name,
                    })

        self.actor.config = config
        if save:
            self.save_config()

    def reset_scale(self, scale):
        if not isinstance(scale, float) \
            and not isinstance(scale, int):
            scale = float(scale[0])
        self.actor.config["scale"] = scale
        self.update_config()
        logging.info("Reloading scene, please wait for about 10 seconds...")
        self.reset()
        self.load_actor(self.actor.get_pose(), self.actor.config)
        self.visualize()

    @staticmethod
    def parse_point(cmd: str, req_id: bool = True):
        parse_map = {
            'c': 'contact',
            't': 'target',
            'f': 'functional',
            'o': 'orientation',
            'contact': 'contact',
            'target': 'target',
            'functional': 'functional',
            'orientation': 'orientation'
        }
        if cmd.strip() == '':
            cmd = input("  >> (t)arget, (c)ontact, (f)unctional, (o)rientation:")
        cmd = cmd.strip().split(" ")

        if req_id:
            try:
                if len(cmd) == 2:
                    type, pid, base = parse_map[cmd[0]], int(cmd[1]), None
                else:
                    type, pid, base = parse_map[cmd[0]], int(cmd[1]), cmd[2]
            except (IndexError, ValueError, KeyError):
                return None, None, None
            return type, pid, base
        else:
            if len(cmd) != 2:
                return None, None
            return parse_map.get(cmd[0], None), cmd[1]

    def get_points(self, type: str) -> list[sapien.Entity]:
        points = []
        for entity in self.scene.get_all_actors():
            if entity.get_name().startswith(type):
                points.append(entity)
        return points

    def get_next_id(self, type: str):
        points = self.get_points(type)
        max_id = -1
        for p in points:
            res = re.search('_(\d+)<(.*?)>', p.get_name())
            max_id = max(max_id, int(res.group(1)))
        return max_id + 1

    def console(self):
        global render_pause
        modified = 0
        try:
            while not self.viewer.closed:
                cmd = input("Input command: ")
                if self.viewer.closed:
                    logging.warning("Viewer has been closed manually.")
                    cmd = input("Please choose to reopen, exit with save or exit without save: (r/s/e) ")
                    cmd = cmd.strip().lower()
                    if cmd in ['r', 'reopen']:
                        self.open_viewer()
                    if cmd in ['s', 'save']:
                        self.update_config(True)
                        break
                    if cmd in ['e', 'exit']:
                        break
                
                modified += 1
                if cmd == "save":
                    self.update_config(True)
                    modified = 0
                elif cmd[:6] == "resize":
                    """
                    Usage:
                        resize <size>: Synchronize the scaling of all three axes of the object
                    Example:
                        resize 0.1
                    """
                    args = cmd[7:].strip().split(" ")
                    size = float(args[0])
                    self.reset_scale(size)
                    modified = 0
                elif cmd == "run":
                    """
                    Get stable points through steps, press Ctrl+C to stop
                    """
                    try:
                        self.run = True
                        while True:
                            pass
                    except KeyboardInterrupt:
                        self.run = False
                        self.actor.config["transform_matrix"] = (
                            self.actor.get_pose().to_transformation_matrix().tolist())
                        self.actor.config["transform_matrix"][0][3] = 0
                        self.actor.config["transform_matrix"][1][3] = 0
                        self.update_config(True)
                    except Exception as e:
                        logging.warning(f"Error: {e}")
                elif cmd == "qpos":
                    """
                    Get current joint state
                    """
                    qpos = self.actor.get_qpos()
                    self.actor.config["init_qpos"] = qpos.tolist()
                    self.update_config(True)
                elif cmd[:4] == "mass":
                    '''
                    Set joint mass
                    '''
                    mass = cmd[5:].split(' ')
                    links = [(link.get_name(), link) for link in self.actor.actor.get_links()]
                    if len(mass) != len(links):
                        logging.warning(f"Mass list length({len(mass)}) does not match link count({len(links)}).")
                        continue
                    links.sort(key=lambda x: x[0])
                    self.actor.config['mass'] = {}
                    idx = 0
                    for name, link in links:
                        if name == 'base': continue
                        self.actor.config['mass'][name] = float(mass[idx])
                        idx += 1
                    self.save_config()
                elif cmd[:6] == "create":
                    """
                    Usage:
                        create <type> <base_link>: Create (t)arget, (c)ontact, (f)unctional, (o)rientation points
                        create: Wait for point name input
                    Example:
                        create t link_1
                    """
                    type, base = self.parse_point(cmd[6:], req_id=False)
                    if type is None or base is None:
                        logging.warning("Invalid type or base.")
                        continue

                    pid = self.get_next_id(type)
                    if type == "orientation" and pid > 0:
                        logging.warning("Orientation point is unique, please modify the existing one.")
                        continue

                    base_link = self.get_link(base)
                    if base_link is None:
                        logging.warning(f"Base link '{base}' not found.")
                    else:
                        self.add_visual_box(self.actor.get_pose(), name=f"{type}_{pid}<{base}>", type=type)
                        logging.info(f"Successfully created {type}_{pid}")
                elif cmd[:6] == "rebase":
                    """
                    Usage:
                        rebase <type> <id> <base_link>: Modify the base link of the specified point
                    Example:
                        rebase c 0 link1
                    """
                    type, pid, base = self.parse_point(cmd[6:], req_id=True)
                    if type is None or pid is None or base is None:
                        logging.warning("Invalid type, id or base link.")
                        continue

                    for entity in self.get_points(type):
                        name = entity.get_name()
                        if name.startswith(f"{type}_{pid}"):
                            new_name = f"{type}_{pid}<{base}>"
                            entity.set_name(new_name)
                            logging.info(f"Successfully rebased {name} to {new_name}")
                elif cmd[:5] == "clone":
                    """
                    Usage:
                        clone <type> <id>: Clone a specified type and ID point in place
                        clone: Wait for input of point type and ID
                    Example:
                        clone t 1: Clone target_1 to create a new target point (e.g., target_2)
                    """
                    type, pid, base = self.parse_point(cmd[5:], req_id=True)
                    if type is None or pid is None:
                        logging.warning("Invalid type or id.")
                        continue
                    if type == "orientation":
                        logging.warning("Orientation point is unique, cloning not supported!")
                        continue

                    name = f"{type}_{pid}"
                    for entity in self.scene.get_all_actors():
                        if entity.get_name().startswith(name):
                            pid = self.get_next_id(type)
                            base = self.get_base_name(entity.get_name())
                            self.add_visual_box(entity.get_pose(), name=f"{type}_{pid}<{base}>", type=type)
                            logging.info(f"Successfully cloned {name}<{base}> to {type}_{pid}<{base}>")
                elif cmd[:6] == "rotate":
                    """
                    Usage:
                        rotate <id> <axis> <interval>: Rotate a specified contact point around its own axis by a given interval, generating points belonging to the same group
                    Example:
                        rotate 1 x 90: Rotate contact_1 around its x-axis every 90 degrees, creating three additional contact points, and writes the group into concat_points_group
                    """
                    cmd = cmd[7:].strip().split(" ")
                    if len(cmd) != 3:
                        continue

                    name, axis, interval = cmd
                    if not name.isdigit():
                        logging.warning("Invalid id, must be a number.")
                        continue
                    name = int(name)

                    axis = {
                        "x": [1, 0, 0],
                        "y": [0, 1, 0],
                        "z": [0, 0, 1],
                        "r": [1, 0, 0],
                        "g": [0, 1, 0],
                        "b": [0, 0, 1],
                    }.get(axis, None)
                    if axis is None:
                        continue
                    axis = np.array(axis)

                    interval = int(interval)
                    for entity in self.scene.get_all_actors():
                        e_name = entity.get_name()
                        if e_name.startswith(f"contact_{name}"):
                            pos_mat = entity.get_pose().to_transformation_matrix()
                            base_name = self.get_base_name(e_name)
                            group_list = [name]
                            for i in range(interval, 360, interval):
                                pid = self.get_next_id("contact")
                                new_mat = np.eye(4)
                                new_mat[:3, 3] = pos_mat[:3, 3]

                                new_rot = axis * np.radians(i)
                                rot_mat = t3d.euler.euler2mat(new_rot[0], new_rot[1], new_rot[2])
                                new_mat[:3, :3] = pos_mat[:3, :3] @ rot_mat

                                pose = sapien.Pose(
                                    new_mat[:3, 3],
                                    t3d.quaternions.mat2quat(new_mat[:3, :3]),
                                )
                                self.add_visual_box(pose, name=f"contact_{pid}<{base_name}>", type="contact")
                                group_list.append(pid)
                            self.actor.config["contact_points_group"].append(group_list)
                            self.actor.config["contact_points_mask"].append(True)
                            logging.info(f"Successfully rotated {e_name} group, created {len(group_list)} points")
                elif cmd[:5] == "align":
                    """
                    Usage:
                        align: Align all group points' positions to the first point in the group
                    """
                    concat_points = {self.get_id(i.get_name()): i for i in self.get_points("contact")}
                    for group in self.actor.config["contact_points_group"]:
                        base = concat_points[group[0]]
                        base_p = base.get_pose().p
                        for pt in group[1:]:
                            concat_points[pt].set_pose(sapien.Pose(base_p, concat_points[pt].get_pose().q))
                        logging.info(f"Aligning group with {', '.join([str(i) for i in group[1:]])} to {group[0]}")
                elif cmd[:6] == "remove":
                    """
                    Usage:
                        remove <type> <id>: Remove a point with the specified name
                        remove: Wait for input of point name
                    Example:
                        remove t 0
                    """
                    type, idx, _ = self.parse_point(cmd[6:], req_id=True)
                    if type is None or idx is None:
                        logging.warning("Invalid type or id.")
                        continue

                    for entity in self.get_points(type):
                        name = entity.get_name()
                        if name.startswith(f"{type}_{idx}"):
                            render_pause = True
                            self.scene.remove_actor(entity)
                            render_pause = False
                            logging.info(f"Successfully removed {name}")
                            break
                elif cmd == "exit":
                    if modified > 1:
                        cmd = input(
                            f'You have made {modified-1} changes without save, do you want to save them? (y/n/others to abort)'
                        )
                        if cmd.strip().lower() == 'y':
                            self.update_config(True)
                            break
                        elif cmd.strip().lower() == 'n':
                            break
                        else:
                            logging.info("Operation has been aborted.")
                    else:
                        break
                else:
                    modified -= 1
                    if cmd != 'help':
                        logging.info(f"Unknown command: {cmd}")
                    help_info = ""
        except KeyboardInterrupt:
            pass

def auto_loader(model_name: str):
    model_dir = Path("./assets/objects/") / model_name
    collision = model_dir / "collision"
    visual    = model_dir / "visual"

    if not collision.exists():
        # URDF
        id_list = [
            int(i.name) for i in list(model_dir.iterdir()) \
                if i.is_dir() and i.name != 'visual'
        ]
        logging.info(f"<URDF> Found {len(id_list)} valid models: {id_list}")
        return URDFViewer(), id_list, sapien.Pose([0, 0, 0], [1, 0, 0, 0])
    else:
        collision_list = [
            int(re.search(r'\d+', i.name).group()) \
                for i in list(collision.iterdir()) \
                    if i.suffix in ['.obj', '.glb'] 
        ]
        visual_list = [
            int(re.search(r'\d+', i.name).group()) \
                for i in list(visual.iterdir()) \
                    if i.suffix in ['.obj', '.glb'] 
        ]
        id_list = list(
            set(visual_list) & set(collision_list))
        logging.info(f"<OBJECT> Found {len(id_list)} valid models: {id_list}")
        return ObjectViewer(), id_list, sapien.Pose([0, 0, 0], [0.707, 0.707, 0, 0])


def main(model_name: str, start: int = 0):
    try:
        viewer, id_list, init_pose = auto_loader(model_name)
    except Exception as e:
        logging.error(f"Failed to load model {model_name}: {e}")
        return

    for oid in id_list:
        if oid < start:
            continue
        os.environ["MODEL_NAME"] = f"{model_name}/{oid}"
        os.environ["MODEL_ID"] = "None"

        inherit_config = None
        try:
            logging.info(f'Annotating {model_name}({oid})')
            viewer.main(
                pose=init_pose,
                modelname=f'{model_name}',
                modelid=oid,
                inherit_config=inherit_config)
            inherit_config = viewer.actor.config
        except KeyboardInterrupt:
            break


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO, format='[{levelname:^8}] {message}', style="{")
    parser = argparse.ArgumentParser(description="Annotation Tool")
    parser.add_argument("model_name", type=str, help="Model Name")
    parser.add_argument("-s", "--start", type=int, default=0, help="Start ID")
    args = parser.parse_args()
    main(args.model_name, args.start)