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


class handover_block(Base_Task):

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

    def load_actors(self):
        rand_pos = rand_pose(
            xlim=[-0.25, -0.05],
            ylim=[0, 0.25],
            zlim=[0.842],
            qpos=[0.981, 0, 0, 0.195],
            rotate_rand=True,
            rotate_lim=[0, 0, 0.2],
        )
        self.box = create_box(
            scene=self,
            pose=rand_pos,
            half_size=(0.03, 0.03, 0.1),
            color=(1, 0, 0),
            name="box",
            boxtype="long",
        )

        rand_pos = rand_pose(
            xlim=[0.1, 0.25],
            ylim=[0.15, 0.2],
        )

        self.target_box = create_box(
            scene=self,
            pose=rand_pos,
            half_size=(0.05, 0.05, 0.005),
            color=(0, 0, 1),
            name="target_box",
            is_static=True,
        )

        self.add_prohibit_area(self.box, padding=0.1)
        self.add_prohibit_area(self.target_box, padding=0.1)
        self.block_middle_pose = [0, 0.0, 0.9, 0, 1, 0, 0]

    def play_once(self):
        # Determine which arm to use for grasping based on box position
        grasp_arm_tag = ArmTag("left" if self.box.get_pose().p[0] < 0 else "right")
        # The other arm will be used for placing
        place_arm_tag = grasp_arm_tag.opposite

        # Grasp the box with the selected arm
        self.move(
            self.grasp_actor(
                self.box,
                arm_tag=grasp_arm_tag,
                pre_grasp_dis=0.07,
                grasp_dis=0.0,
                contact_point_id=[0, 1, 2, 3],
            ))
        # Lift the box up
        self.move(self.move_by_displacement(grasp_arm_tag, z=0.1))
        # Place the box at initial position [0, 0., 0.9, 0, 1, 0, 0]
        self.move(
            self.place_actor(
                self.box,
                target_pose=self.block_middle_pose,
                arm_tag=grasp_arm_tag,
                functional_point_id=0,
                pre_dis=0,
                dis=0,
                is_open=False,
                constrain="free",
            ))

        # Grasp the box again with the other arm (for repositioning)
        self.move(
            self.grasp_actor(
                self.box,
                arm_tag=place_arm_tag,
                pre_grasp_dis=0.07,
                grasp_dis=0.0,
                contact_point_id=[4, 5, 6, 7],
            ))
        # Open the original grasping arm's gripper
        self.move(self.open_gripper(grasp_arm_tag))
        # Move the original arm up to release the box
        self.move(self.move_by_displacement(grasp_arm_tag, z=0.1, move_axis="arm"))
        # Perform two actions simultaneously:
        # 1. Return the original arm to its origin position
        # 2. Place the box at the target's functional point with precise alignment
        self.move(
            self.back_to_origin(grasp_arm_tag),
            self.place_actor(
                self.box,
                target_pose=self.target_box.get_functional_point(1, "pose"),
                arm_tag=place_arm_tag,
                functional_point_id=0,
                pre_dis=0.05,
                dis=0.,
                constrain="align",
                pre_dis_axis="fp",
            ),
        )

        return self.info

    def check_success(self):
        box_pos = self.box.get_functional_point(0, "pose").p
        target_pose = self.target_box.get_functional_point(1, "pose").p
        eps = [0.03, 0.03]
        return (np.all(np.abs(box_pos[:2] - target_pose[:2]) < eps) and abs(box_pos[2] - target_pose[2]) < 0.01
                and self.is_right_gripper_open())
