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


class rotate_qrcode(Base_Task):

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

    def load_actors(self):
        qrcode_pose = rand_pose(
            xlim=[-0.25, 0.25],
            ylim=[-0.2, 0.0],
            qpos=[0, 0, 0.707, 0.707],
            rotate_rand=True,
            rotate_lim=[0, 0.7, 0],
        )
        while abs(qrcode_pose.p[0]) < 0.05:
            qrcode_pose = rand_pose(
                xlim=[-0.25, 0.25],
                ylim=[-0.2, 0.0],
                qpos=[0, 0, 0.707, 0.707],
                rotate_rand=True,
                rotate_lim=[0, 0.7, 0],
            )

        self.model_id = np.random.choice([0, 1, 2, 3], 1)[0]
        self.qrcode = create_actor(
            self,
            pose=qrcode_pose,
            modelname="070_paymentsign",
            convex=True,
            model_id=self.model_id,
        )

        self.add_prohibit_area(self.qrcode, padding=0.12)
        # Define target placement position based on arm tag (left or right side of table)
        target_x = -0.2 if self.qrcode.get_pose().p[0] < 0 else 0.2
        self.target_pose = [target_x, -0.15, 0.74 + self.table_z_bias, 1, 0, 0, 0]

    def play_once(self):
        # Determine which arm to use based on QR code position (left if on left side, right otherwise)
        arm_tag = ArmTag("left" if self.qrcode.get_pose().p[0] < 0 else "right")

        # Grasp the QR code with specified pre-grasp distance
        self.move(self.grasp_actor(self.qrcode, arm_tag=arm_tag, pre_grasp_dis=0.05))

        # Lift the QR code vertically by 0.07 meters
        self.move(self.move_by_displacement(arm_tag=arm_tag, z=0.07))

        # Place the QR code at the target position with specified placement parameters
        self.move(
            self.place_actor(
                self.qrcode,
                arm_tag=arm_tag,
                target_pose=self.target_pose,
                pre_dis=0.07,
                dis=0.01,
                constrain="align",
            ))

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

    def check_success(self):
        qrcode_quat = self.qrcode.get_pose().q
        qrcode_pos = self.qrcode.get_pose().p
        target_quat = [0.707, 0.707, 0, 0]
        if qrcode_quat[0] < 0:
            qrcode_quat = qrcode_quat * -1
        eps = 0.05
        return (np.all(np.abs(qrcode_quat - target_quat) < eps) and qrcode_pos[2] < 0.75 + self.table_z_bias
                and self.is_left_gripper_open() and self.is_right_gripper_open())
