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


class press_stapler(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.1, 0.05],
            qpos=[0.5, 0.5, 0.5, 0.5],
            rotate_rand=True,
            rotate_lim=[0, np.pi, 0],
        )

        self.stapler_id = np.random.choice([0, 1, 2, 3, 4, 5, 6], 1)[0]
        self.stapler = create_actor(self,
                                    pose=rand_pos,
                                    modelname="048_stapler",
                                    convex=True,
                                    model_id=self.stapler_id,
                                    is_static=True)

        self.add_prohibit_area(self.stapler, padding=0.05)

    def play_once(self):
        # Determine which arm to use based on stapler's position (left if negative x, right otherwise)
        arm_tag = ArmTag("left" if self.stapler.get_pose().p[0] < 0 else "right")

        # Move arm to the overhead position of the stapler and close the gripper
        self.move(self.grasp_actor(self.stapler, arm_tag=arm_tag, pre_grasp_dis=0.1, grasp_dis=0.1, contact_point_id=2))
        self.move(self.close_gripper(arm_tag=arm_tag))

        # Move the stapler down slightly to press it
        self.move(
            self.grasp_actor(self.stapler, arm_tag=arm_tag, pre_grasp_dis=0.02, grasp_dis=0.02, contact_point_id=2))

        self.info["info"] = {"{A}": f"048_stapler/base{self.stapler_id}", "{a}": str(arm_tag)}
        return self.info

    def check_success(self):
        if self.stage_success_tag:
            return True
        stapler_pose = self.stapler.get_contact_point(2)[:3]
        positions = self.get_gripper_actor_contact_position("048_stapler")
        eps = [0.03, 0.03]
        for position in positions:
            if (np.all(np.abs(position[:2] - stapler_pose[:2]) < eps) and abs(position[2] - stapler_pose[2]) < 0.03):
                self.stage_success_tag = True
                return True
        return False
