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


class handover_mic(Base_Task):

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

    def load_actors(self):
        rand_pos = rand_pose(
            xlim=[-0.2, 0.2],
            ylim=[-0.05, 0.0],
            qpos=[0.707, 0.707, 0, 0],
            rotate_rand=False,
        )
        while abs(rand_pos.p[0]) < 0.15:
            rand_pos = rand_pose(
                xlim=[-0.2, 0.2],
                ylim=[-0.05, 0.0],
                qpos=[0.707, 0.707, 0, 0],
                rotate_rand=False,
            )
        self.microphone_id = np.random.choice([0, 4, 5], 1)[0]

        self.microphone = create_actor(
            scene=self,
            pose=rand_pos,
            modelname="018_microphone",
            convex=True,
            model_id=self.microphone_id,
        )

        self.add_prohibit_area(self.microphone, padding=0.07)
        self.handover_middle_pose = [0, -0.05, 0.98, 0, 1, 0, 0]
        self.grasp_arm_tag = ArmTag("right" if self.microphone.get_pose().p[0] > 0 else "left")
        self.handover_arm_tag = self.grasp_arm_tag.opposite

    def play_once(self):
        # Determine the arm to grasp the microphone based on its position
        grasp_arm_tag = ArmTag("right" if self.microphone.get_pose().p[0] > 0 else "left")
        # The opposite arm will be used for the handover
        handover_arm_tag = grasp_arm_tag.opposite

        # Move the grasping arm to the microphone's position and grasp it
        self.move(
            self.grasp_actor(
                self.microphone,
                arm_tag=grasp_arm_tag,
                contact_point_id=[1, 9, 10, 11, 12, 13, 14, 15],
                pre_grasp_dis=0.1,
            ))
        # Move the handover arm to a position suitable for handing over the microphone
        self.move(
            self.move_by_displacement(
                grasp_arm_tag,
                z=0.12,
                quat=(GRASP_DIRECTION_DIC["front_right"]
                      if grasp_arm_tag == "left" else GRASP_DIRECTION_DIC["front_left"]),
                move_axis="arm",
            ))
        
        # Move the handover arm to the middle position for handover
        self.move(
            self.place_actor(
                self.microphone,
                arm_tag=grasp_arm_tag,
                target_pose=self.handover_middle_pose,
                functional_point_id=0,
                pre_dis=0.0,
                dis=0.0,
                is_open=False,
                constrain="free",
            ))
        # Move the handover arm to grasp the microphone from the grasping arm
        self.move(
            self.grasp_actor(
                self.microphone,
                arm_tag=handover_arm_tag,
                contact_point_id=[0, 2, 3, 4, 5, 6, 7, 8],
                pre_grasp_dis=0.1,
            ))
        # Move the grasping arm to open the gripper and lift the microphone
        self.move(self.open_gripper(grasp_arm_tag))
        # Move the handover arm to lift the microphone to a height of 0.98
        self.move(
            self.move_by_displacement(grasp_arm_tag, z=0.07, move_axis="arm"),
            self.move_by_displacement(handover_arm_tag, x=0.05 if handover_arm_tag == "right" else -0.05),
        )

        self.info["info"] = {
            "{A}": f"018_microphone/base{self.microphone_id}",
            "{a}": str(grasp_arm_tag),
            "{b}": str(handover_arm_tag),
        }
        return self.info

    def check_success(self):
        microphone_pose = self.microphone.get_functional_point(0)
        contact = self.get_gripper_actor_contact_position("018_microphone")
        if len(contact) == 0:
            return False
        close_gripper_func = self.is_left_gripper_close if self.handover_arm_tag == "left" else self.is_right_gripper_close
        open_gripper_func = self.is_left_gripper_open if self.grasp_arm_tag == "left" else self.is_right_gripper_open
        tag = microphone_pose[0] < 0 if self.handover_arm_tag == "left" else microphone_pose[0] > 0
        return (close_gripper_func() and open_gripper_func() and microphone_pose[2] > 0.92 and tag)
