from ._base_task import Base_Task
from .utils import *


class turn_switch(Base_Task):

    def setup_demo(self, is_test=False, **kwargs):
        super()._init_task_env_(**kwargs)

    def load_actors(self):
        self.model_name = "056_switch"
        self.model_id = np.random.randint(0, 8)
        self.switch = rand_create_sapien_urdf_obj(
            scene=self,
            modelname=self.model_name,
            modelid=self.model_id,
            xlim=[-0.25, 0.25],
            ylim=[0.0, 0.1],
            zlim=[0.81, 0.84],
            rotate_rand=True,
            rotate_lim=[0, 0, np.pi / 4],
            qpos=[0.704141, 0, 0, 0.71006],
            fix_root_link=True,
        )
        self.prohibited_area.append([-0.4, -0.2, 0.4, 0.2])

    def play_once(self):
        switch_pose = self.switch.get_pose()
        face_dir = -switch_pose.to_transformation_matrix()[:3, 0]
        arm_tag = ArmTag("right" if face_dir[0] > 0 else "left")

        # close gripper
        self.move(self.close_gripper(arm_tag=arm_tag, pos=0))
        # move the gripper to turn off the switch
        self.move(self.grasp_actor(self.switch, arm_tag=arm_tag, pre_grasp_dis=0.04))

        self.info["info"] = {"{A}": f"056_switch/base{self.model_id}", "{a}": str(arm_tag)}
        return self.info

    def check_success(self):
        limit = self.switch.get_qlimits()[0]
        return self.switch.get_qpos()[0] >= limit[1] - 0.05
