from ._base_task import Base_Task
from .utils import *
import sapien
import math
from ._GLOBAL_CONFIGS import *
from copy import deepcopy


class move_playingcard_away(Base_Task):

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

    def load_actors(self):
        rand_pos = rand_pose(
            xlim=[-0.1, 0.1],
            ylim=[-0.2, 0.05],
            qpos=[0.5, 0.5, 0.5, 0.5],
            rotate_rand=True,
            rotate_lim=[0, 3.14, 0],
        )
        while abs(rand_pos.p[0]) < 0.05:
            rand_pos = rand_pose(
                xlim=[-0.1, 0.1],
                ylim=[-0.2, 0.05],
                qpos=[0.5, 0.5, 0.5, 0.5],
                rotate_rand=True,
                rotate_lim=[0, 3.14, 0],
            )

        self.playingcards_id = np.random.choice([0, 1, 2], 1)[0]
        self.playingcards = create_actor(
            scene=self,
            pose=rand_pos,
            modelname="081_playingcards",
            convex=True,
            model_id=self.playingcards_id,
        )

        self.prohibited_area.append([-100, -0.3, 100, 0.1])
        self.add_prohibit_area(self.playingcards, padding=0.1)

        self.target_pose = self.playingcards.get_pose() # TODO

    def play_once(self):
        # Determine which arm to use based on playing cards position
        arm_tag = ArmTag("right" if self.playingcards.get_pose().p[0] > 0 else "left")

        # Grasp the playing cards with specified arm
        self.move(self.grasp_actor(self.playingcards, arm_tag=arm_tag, pre_grasp_dis=0.1, grasp_dis=0.01))
        # Move the playing cards horizontally (right if right arm, left if left arm)
        self.move(self.move_by_displacement(arm_tag, x=0.3 if arm_tag == "right" else -0.3))
        # Open gripper to release the playing cards
        self.move(self.open_gripper(arm_tag))

        self.info["info"] = {
            "{A}": f"081_playingcards/base{self.playingcards_id}",
            "{a}": str(arm_tag),
        }
        return self.info

    def check_success(self):
        playingcards_pose = self.playingcards.get_pose().p
        edge_x = 0.23

        return (np.all(abs(playingcards_pose[0]) > abs(edge_x)) and self.robot.is_left_gripper_open()
                and self.robot.is_right_gripper_open())
