from pathlib import Path
import pybullet as p
import pybullet_data
import rospkg
import pickle
import os

from active_grasp.bbox import AABBox
from robot_helpers.bullet import *
from robot_helpers.io import load_yaml
from robot_helpers.model import KDLModel
from robot_helpers.spatial import Rotation
from vgn.perception import UniformTSDFVolume
from vgn.utils import find_urdfs, view_on_sphere
from vgn.detection import VGN, select_local_maxima

# import vgn.visualizer as vis

rospack = rospkg.RosPack()
pkg_root = Path(rospack.get_path("active_grasp"))
urdfs_dir = pkg_root / "assets"


class Simulation:
    """Robot is placed s.t. world and base frames are the same"""

    def __init__(self, gui, scene_id, vgn_path):
        self.configure_physics_engine(gui, 60, 4)
        self.configure_visualizer()
        self.seed()
        self.load_robot()
        self.load_vgn(Path(vgn_path))
        self.scene = get_scene(scene_id)
        self.target_uid = None

    def configure_physics_engine(self, gui, rate, sub_step_count):
        self.rate = rate
        self.dt = 1.0 / self.rate
        p.connect(p.GUI if gui else p.DIRECT)
        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        p.setPhysicsEngineParameter(fixedTimeStep=self.dt, numSubSteps=sub_step_count)
        p.setGravity(0.0, 0.0, -9.81)

    def configure_visualizer(self):
        # p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0)
        p.resetDebugVisualizerCamera(1.2, 30, -30, [0.4, 0.0, 0.2])

    def seed(self, seed=None):
        self.rng = np.random.default_rng(seed) if seed else np.random

    def load_robot(self):
        panda_urdf_path = urdfs_dir / "franka/panda_arm_hand.urdf"
        self.arm = BtPandaArm(panda_urdf_path)
        print(f"Robot base uid: {self.arm.uid}")
        
        # 打印机械臂的所有部件信息
        num_joints = p.getNumJoints(self.arm.uid)
        print(f"Number of joints: {num_joints}")
        for i in range(num_joints):
            joint_info = p.getJointInfo(self.arm.uid, i)
            link_name = joint_info[12].decode('utf-8')
            print(f"Joint {i}: {link_name}")
        
        self.gripper = BtPandaGripper(self.arm)
        self.model = KDLModel.from_urdf_file(
            panda_urdf_path, self.arm.base_frame, self.arm.ee_frame
        )
        self.camera = BtCamera(320, 240, 0.96, 0.01, 1.0, self.arm.uid, 11)

    def load_vgn(self, model_path):
        self.vgn = VGN(model_path)

    def reset(self):
        # Wait until the simulation is stable
        secs = 3.0
        sleep_ticks = self.rate * secs
        for i in range(int(sleep_ticks)):
            p.stepSimulation()
        self.target_uid = None
        # Reset the scene
        valid = False
        while not valid:
            self.set_arm_configuration([0.0, -1.39, 0.0, -2.36, 0.0, 1.57, 0.79])
            self.scene.clear()
            q = self.scene.generate(self.rng)
            self.set_arm_configuration(q)
            uid = self.select_target()
            if uid == -1:
                continue
            support_id = self.select_support()  
            bbox = self.get_target_bbox(uid)
            #import ipdb; ipdb.set_trace()
            valid = True
            # valid = self.check_for_grasps(bbox)
        return bbox

    def set_arm_configuration(self, q):
        #import ipdb; ipdb.set_trace()
        for i, q_i in enumerate(q):
            p.resetJointState(self.arm.uid, i, q_i, 0)
        p.resetJointState(self.arm.uid, 9, 0.04, 0)
        p.resetJointState(self.arm.uid, 10, 0.04, 0)
        self.gripper.set_desired_width(0.4)

    def select_target(self):
        if self.target_uid is not None:
            p.changeVisualShape(self.target_uid, -1, rgbaColor=[1, 0, 0, 1])
            return self.target_uid
        _, _, mask = self.camera.get_image()
        uids, counts = np.unique(mask, return_counts=True)
        
        mask = np.isin(uids, self.scene.object_uids)  # remove ids of the floor, etc
        uids, counts = uids[mask], counts[mask]
        # If no objects are detected, try again
        container_idx = np.where(np.isin(uids, self.scene.container_uid))[0]
        object_uids = np.delete(uids, container_idx)
        object_counts = np.delete(counts, container_idx)
        target_uid = object_uids[np.argmin(object_counts)]
        if(len(object_uids) == 0 and len(object_counts) == 0):
            return -1
        
        print("select target uid: "+str(target_uid))
        p.changeVisualShape(target_uid, -1, rgbaColor=[1, 0, 0, 1])
        self.target_uid = target_uid
        return target_uid
    
    def select_support(self):
        support_id = self.scene.support_uid
        return support_id

    def get_target_bbox(self, uid):
        aabb_min, aabb_max = p.getAABB(uid)
        # enlarge the bounding box
        #import ipdb; ipdb.set_trace()
        aabb_min = np.asarray(aabb_min)
        aabb_max = np.asarray(aabb_max)
        aabb_min -= 0.0
        aabb_max += 0.0

        return AABBox(aabb_min, aabb_max)

    def check_for_grasps(self, bbox):
        origin = Transform.from_translation(self.scene.origin)
        origin.translation[2] -= 0.05
        center = Transform.from_translation(self.scene.center)

        # First, reconstruct the scene from many views
        tsdf = UniformTSDFVolume(self.scene.length, 40)
        r = 2.0 * self.scene.length
        theta = np.pi / 4.0
        phis = np.linspace(0.0, 2.0 * np.pi, 5)
        for view in [view_on_sphere(center, r, theta, phi) for phi in phis]:
            depth_img = self.camera.get_image(view)[1]
            tsdf.integrate(depth_img, self.camera.intrinsic, view.inv() * origin)
        voxel_size, tsdf_grid = tsdf.voxel_size, tsdf.get_grid()

        # Then check whether VGN can find any grasps on the target
        out = self.vgn.predict(tsdf_grid)
        grasps, qualities = select_local_maxima(voxel_size, out, threshold=0.9)

        # vis.scene_cloud(voxel_size, tsdf.get_scene_cloud())
        # vis.grasps(grasps, qualities, 0.05)
        # vis.show()

        for grasp in grasps:
            pose = origin * grasp.pose
            tip = pose.rotation.apply([0, 0, 0.05]) + pose.translation
            if bbox.is_inside(tip):
                return True
        return False

    def step(self):
        p.stepSimulation()


class Scene:
    def __init__(self):
        self.support_urdf = urdfs_dir / "plane/model.urdf"
        self.support_uid = -1
        self.object_uids = []
        self.container_uid = -1
        self.container_list = [
                                "chair",
                                "table",
                                "sofa",
                                "picnic_basket",
                                "cabinet",
                                "flower_pot",
                                "blender_cabinet"
                            ]

    def clear(self):
        self.remove_support()
        self.remove_all_objects()

    def generate(self, rng):
        raise NotImplementedError

    def add_support(self, pos):
        self.support_uid = p.loadURDF(str(self.support_urdf), pos, globalScaling=0.3)
        print('support id: '+str(self.support_uid))

    def remove_support(self):
        p.removeBody(self.support_uid)

    def add_object(self, urdf, ori, pos, scale=1.0):
        uid = p.loadURDF(str(urdf), pos, ori.as_quat(), globalScaling=scale)
        print('add object uid: '+str(uid))
        import ipdb; ipdb.set_trace()
        
        self.object_uids.append(uid)
        return uid

    def remove_object(self, uid):
        p.removeBody(uid)
        self.object_uids.remove(uid)

    def remove_all_objects(self):
        print(f"Objects to remove: {self.object_uids}")
        for uid in list(self.object_uids):
            # 获取物体信息
            body_info = p.getBodyInfo(uid)
            print(f'Removing object uid: {uid}, Name: {body_info[1].decode("utf-8")}')
            self.remove_object(uid)


class YamlScene(Scene):
    def __init__(self, config_name):
        super().__init__()
        #import ipdb; ipdb.set_trace()
        self.config_path = pkg_root / "cfg/sim" / config_name

    def load_config(self):
        self.scene = load_yaml(self.config_path)
        self.center = np.asarray(self.scene["center"])
        self.length = 0.3
        self.origin = self.center - np.r_[0.5 * self.length, 0.5 * self.length, 0.0]

    def generate(self, rng):
        self.load_config()
        self.add_support(self.center)
        for object in self.scene["objects"]:
            urdf = urdfs_dir / object["object_id"] / "model.urdf"
            ori = Rotation.from_euler("xyz", object["rpy"], degrees=True)
            pos = self.center + np.asarray(object["xyz"])
            scale = object.get("scale", 1)
            if randomize := object.get("randomize", False):
                angle = rng.uniform(-randomize["rot"], randomize["rot"])
                ori = Rotation.from_euler("z", angle, degrees=True) * ori
                b = np.asarray(randomize["pos"])
                pos += rng.uniform(-b, b)
            self.add_object(urdf, ori, pos, scale)
        for _ in range(60):
            p.stepSimulation()
        return self.scene["q"]

class PickleScene(Scene):
    def __init__(self, config_name):
        super().__init__()
        self.config_path = pkg_root / "cfg/sim" / config_name
        self.center = np.r_[0.5, 0.0, 0.0]
        self.length = 0.3
        self.origin = self.center - np.r_[0.5 * self.length, 0.5 * self.length, 0.0]
        

    def load_scene(self, scene_dir, model_dir, rot_z=0):
        scene_path = os.path.join(scene_dir, "scene.pickle")
        self.scene = pickle.load(open(scene_path, "rb"))
        self.fall_objects_dict = self.get_fall_objects_dict(scene_dir)
        self._initialize_pybullet_scene(model_dir, rot_z)

    def _initialize_pybullet_scene(self, model_dir, rot_z=0):
        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        # 修改平面加载,添加碰撞属性
        planeId = p.loadURDF("plane100.urdf", useMaximalCoordinates=True, flags=p.URDF_USE_IMPLICIT_CYLINDER)
        p.changeDynamics(planeId, -1, 
            lateralFriction=1.0,
            restitution=0.5,
            contactStiffness=1000000,
            contactDamping=100
        )
        
        # 创建旋转矩阵
        rot_matrix = Rotation.from_euler("z", rot_z, degrees=True)
        
        for obj_name in self.scene.keys():
            if obj_name in self.fall_objects_dict:
                continue
            # 处理旋转
            origin_orientation = self.scene[obj_name]["rotation"] # quaternion
            origin_orientation = Rotation.from_quat(origin_orientation)
            orientation = rot_matrix * origin_orientation
            orientation = orientation.as_quat()
            
            # 处理位置
            original_position = self.scene[obj_name]["position"]
            # 将位置向量绕z轴旋转相同的角度
            rotated_position = rot_matrix.apply(original_position)
            position = rotated_position + self.center
            
            class_name = obj_name[:-4]
            #obj_path = os.path.join(model_dir, class_name, obj_name, "model.obj")
            
            urdf_path = os.path.join(model_dir, class_name, obj_name, "model.urdf")
            #import ipdb; ipdb.set_trace()
            uid = p.loadURDF(urdf_path, position, orientation, globalScaling=1.0)
            # uid = self._load_obj_to_pybullet(
            #     obj_file_path=obj_path,
            #     position=position,
            #     orientation=orientation,
            #     scale=1.0,
            # )
            
            self.object_uids.append(uid)
            for container in self.container_list:
                if container in obj_name:
                    self.container_uid = uid
                    print('container id: '+str(self.container_uid))

    def get_fall_objects_dict(self, scene_path):
        fall_objects_path = os.path.join(scene_path, "fall_objects.pickle")
        fall_objects_dict = pickle.load(open(fall_objects_path, "rb"))
        return fall_objects_dict
    
    def _load_obj_to_pybullet(self, obj_file_path, position, orientation, scale):
        visual_ind = p.createVisualShape(
            shapeType=p.GEOM_MESH,
            fileName=obj_file_path,
            rgbaColor=[1, 1, 1, 1],
            specularColor=[0.4, 0.4, 0],
            visualFramePosition=[0, 0, 0],
            meshScale=scale,
        )
        collision_ind = p.createCollisionShape(
            shapeType=p.GEOM_MESH,
            fileName=obj_file_path,
            meshScale=scale,
        )
        uid = p.createMultiBody(
            baseMass=0,
            baseVisualShapeIndex=visual_ind,
            baseCollisionShapeIndex=collision_ind,
            basePosition=position,
            baseOrientation=orientation,
            useMaximalCoordinates=True,
        )
        return uid

    def generate(self, rng):
        model_dir = "/media/hofee/data/data/nbv/nbv_objects_urdf"
        self.load_scene(self.config_path, model_dir, rot_z=90)
        self.add_support(self.center - np.array([0.0, 0.0, 0.05]))
        # for _ in range(60):
        #     p.stepSimulation()
            
        q = np.asarray([0.0, -1.39, 0.0, -2.36, 0.0, 1.57, 0.79])
        q += np.asarray(rng.uniform(-0.08, 0.08, 7))
        return q.tolist()

def get_scene(scene_id):
    if scene_id.endswith(".yaml"):
        return YamlScene(scene_id)
    else:
        return PickleScene(scene_id)
