from ._base_task import Base_Task
from .utils import *
from ._GLOBAL_CONFIGS import *


class scan_object(Base_Task):

    def setup_demo(self, **kwags):
        super()._init_task_env_(**kwags)

    def load_actors(self):
        tag = np.random.randint(2)
        if tag == 0:
            scanner_x_lim = [-0.25, -0.05]
            object_x_lim = [0.05, 0.25]
        else:
            scanner_x_lim = [0.05, 0.25]
            object_x_lim = [-0.25, -0.05]
        scanner_pose = rand_pose(
            xlim=scanner_x_lim,
            ylim=[-0.15, -0.05],
            qpos=[0, 0, 0.707, 0.707],
            rotate_rand=True,
            rotate_lim=[0, 1.2, 0],
        )
        self.scanner_id = np.random.choice([0, 1, 2, 3, 4], 1)[0]
        self.scanner = create_actor(
            scene=self.scene,
            pose=scanner_pose,
            modelname="024_scanner",
            convex=True,
            model_id=self.scanner_id,
        )

        object_pose = rand_pose(
            xlim=object_x_lim,
            ylim=[-0.2, 0.0],
            qpos=[0.5, 0.5, 0.5, 0.5],
            rotate_rand=True,
            rotate_lim=[0, 1.2, 0],
        )
        self.object_id = np.random.choice([0, 1, 2, 3, 4, 5], 1)[0]
        self.object = create_actor(
            scene=self.scene,
            pose=object_pose,
            modelname="112_tea-box",
            convex=True,
            model_id=self.object_id,
        )
        self.add_prohibit_area(self.scanner, padding=0.1)
        self.add_prohibit_area(self.object, padding=0.1)
        target_posi = [-0.2, -0.03, 0.2, -0.01]
        self.prohibited_area.append(target_posi)
        self.left_object_target_pose = [-0.03, -0.02, 0.95, 0.707, 0, -0.707, 0]
        self.right_object_target_pose = [0.03, -0.02, 0.95, 0.707, 0, 0.707, 0]

    def play_once(self):
        scanner_arm_tag = ArmTag("left" if self.scanner.get_pose().p[0] < 0 else "right")
        object_arm_tag = scanner_arm_tag.opposite

        # Move the scanner and object to the gripper
        self.move(
            self.grasp_actor(self.scanner, arm_tag=scanner_arm_tag, pre_grasp_dis=0.08),
            self.grasp_actor(self.object, arm_tag=object_arm_tag, pre_grasp_dis=0.08),
        )
        self.move(
            self.move_by_displacement(arm_tag=scanner_arm_tag, x=0.05 if scanner_arm_tag == "right" else -0.05, z=0.13),
            self.move_by_displacement(arm_tag=object_arm_tag, x=0.05 if object_arm_tag == "right" else -0.05, z=0.13),
        )
        # Get object target pose and place the object
        object_target_pose = (self.right_object_target_pose
                              if object_arm_tag == "right" else self.left_object_target_pose)
        self.move(
            self.place_actor(
                self.object,
                arm_tag=object_arm_tag,
                target_pose=object_target_pose,
                pre_dis=0.0,
                dis=0.0,
                is_open=False,
            ))

        # Move the scanner to align with the object
        self.move(
            self.place_actor(
                self.scanner,
                arm_tag=scanner_arm_tag,
                target_pose=self.object.get_functional_point(1),
                functional_point_id=0,
                pre_dis=0.05,
                dis=0.05,
                is_open=False,
            ))

        self.info["info"] = {
            "{A}": f"112_tea-box/base{self.object_id}",
            "{B}": f"024_scanner/base{self.scanner_id}",
            "{a}": str(object_arm_tag),
            "{b}": str(scanner_arm_tag),
        }
        return self.info

    def check_success(self):
        object_pose = self.object.get_pose().p
        scanner_func_pose = self.scanner.get_functional_point(0)
        target_vec = t3d.quaternions.quat2mat(scanner_func_pose[-4:]) @ np.array([0, 0, -1])
        obj2scanner_vec = scanner_func_pose[:3] - object_pose
        dis = np.sum(target_vec * obj2scanner_vec)
        object_pose1 = object_pose + dis * target_vec
        eps = 0.025
        return (np.all(np.abs(object_pose1 - scanner_func_pose[:3]) < eps) and dis > 0 and dis < 0.07
                and self.is_left_gripper_close() and self.is_right_gripper_close())
