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


class place_empty_cup(Base_Task):

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

    def load_actors(self):
        tag = np.random.randint(0, 2)
        cup_xlim = [[0.15, 0.3], [-0.3, -0.15]]
        coaster_lim = [[-0.05, 0.1], [-0.1, 0.05]]
        self.cup = rand_create_actor(
            self,
            xlim=cup_xlim[tag],
            ylim=[-0.2, 0.05],
            modelname="021_cup",
            rotate_rand=False,
            qpos=[0.5, 0.5, 0.5, 0.5],
            convex=True,
            model_id=0,
        )
        cup_pose = self.cup.get_pose().p

        coaster_pose = rand_pose(
            xlim=coaster_lim[tag],
            ylim=[-0.2, 0.05],
            rotate_rand=False,
            qpos=[0.5, 0.5, 0.5, 0.5],
        )

        while np.sum(pow(cup_pose[:2] - coaster_pose.p[:2], 2)) < 0.01:
            coaster_pose = rand_pose(
                xlim=coaster_lim[tag],
                ylim=[-0.2, 0.05],
                rotate_rand=False,
                qpos=[0.5, 0.5, 0.5, 0.5],
            )
        self.coaster = create_actor(
            self,
            pose=coaster_pose,
            modelname="019_coaster",
            convex=True,
            model_id=0,
            is_static=True
        )

        self.add_prohibit_area(self.cup, padding=0.05)
        self.add_prohibit_area(self.coaster, padding=0.05)
        self.delay(2)
        cup_pose = self.cup.get_pose().p

    def play_once(self):
        # Get the current pose of the cup
        cup_pose = self.cup.get_pose().p
        # Determine which arm to use based on cup's x position (right if positive, left if negative)
        arm_tag = ArmTag("right" if cup_pose[0] > 0 else "left")

        # Close the gripper to prepare for grasping
        self.move(self.close_gripper(arm_tag, pos=0.6))
        # Grasp the cup using the selected arm
        self.move(
            self.grasp_actor(
                self.cup,
                arm_tag,
                pre_grasp_dis=0.1,
                contact_point_id=[0, 2][int(arm_tag == "left")],
            ))
        # Lift the cup up by 0.08 meters along z-axis
        self.move(self.move_by_displacement(arm_tag, z=0.08, move_axis="arm"))

        # Get coaster's functional point as target pose
        target_pose = self.coaster.get_functional_point(0)
        # Place the cup onto the coaster
        self.move(self.place_actor(
            self.cup,
            arm_tag,
            target_pose=target_pose,
            functional_point_id=0,
            pre_dis=0.05,
        ))
        # Lift the arm slightly (0.05m) after placing to avoid collision
        self.move(self.move_by_displacement(arm_tag, z=0.05, move_axis="arm"))

        self.info["info"] = {"{A}": "021_cup/base0", "{B}": "019_coaster/base0"}
        return self.info

    def check_success(self):
        # eps = [0.03, 0.03, 0.015]
        eps = 0.035
        cup_pose = self.cup.get_functional_point(0, "pose").p
        coaster_pose = self.coaster.get_functional_point(0, "pose").p
        return (
            # np.all(np.abs(cup_pose - coaster_pose) < eps)
            np.sum(pow(cup_pose[:2] - coaster_pose[:2], 2)) < eps**2 and abs(cup_pose[2] - coaster_pose[2]) < 0.015
            and self.is_left_gripper_open() and self.is_right_gripper_open())
