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


class place_container_plate(Base_Task):

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

    def load_actors(self):
        container_pose = rand_pose(
            xlim=[-0.28, 0.28],
            ylim=[-0.1, 0.05],
            rotate_rand=False,
            qpos=[0.5, 0.5, 0.5, 0.5],
        )
        while abs(container_pose.p[0]) < 0.2:
            container_pose = rand_pose(
                xlim=[-0.28, 0.28],
                ylim=[-0.1, 0.05],
                rotate_rand=False,
                qpos=[0.5, 0.5, 0.5, 0.5],
            )
        id_list = {"002_bowl": [1, 2, 3, 5], "021_cup": [1, 2, 3, 4, 5, 6, 7]}
        self.actor_name = np.random.choice(["002_bowl", "021_cup"])
        self.container_id = np.random.choice(id_list[self.actor_name])
        self.container = create_actor(
            self,
            pose=container_pose,
            modelname=self.actor_name,
            model_id=self.container_id,
            convex=True,
        )

        x = 0.05 if self.container.get_pose().p[0] > 0 else -0.05
        self.plate_id = 0
        pose = rand_pose(
            xlim=[x - 0.03, x + 0.03],
            ylim=[-0.15, -0.1],
            rotate_rand=False,
            qpos=[0.5, 0.5, 0.5, 0.5],
        )
        self.plate = create_actor(
            self,
            pose=pose,
            modelname="003_plate",
            scale=[0.025, 0.025, 0.025],
            is_static=True,
            convex=True,
        )
        self.add_prohibit_area(self.container, padding=0.1)
        self.add_prohibit_area(self.plate, padding=0.1)

    def play_once(self):
        # Get container's position to determine which arm to use
        container_pose = self.container.get_pose().p
        # Select arm based on container's x position (right if positive, left if negative)
        arm_tag = ArmTag("right" if container_pose[0] > 0 else "left")

        # Grasp the container using selected arm with specific contact point
        self.move(
            self.grasp_actor(
                self.container,
                arm_tag=arm_tag,
                contact_point_id=[0, 2][int(arm_tag == "left")],
                pre_grasp_dis=0.1,
            ))
        # Lift the container up by 0.1m along z-axis
        self.move(self.move_by_displacement(arm_tag, z=0.1, move_axis="arm"))

        # Place the container onto the plate's functional point
        self.move(
            self.place_actor(
                self.container,
                target_pose=self.plate.get_functional_point(0),
                arm_tag=arm_tag,
                functional_point_id=0,
                pre_dis=0.12,
                dis=0.03,
            ))
        # Move the arm up by 0.1m after placing
        self.move(self.move_by_displacement(arm_tag, z=0.08, move_axis="arm"))

        # Record information about the objects and arm used
        self.info["info"] = {
            "{A}": f"003_plate/base{self.plate_id}",
            "{B}": f"{self.actor_name}/base{self.container_id}",
            "{a}": str(arm_tag),
        }
        return self.info

    def check_success(self):
        container_pose = self.container.get_pose().p
        target_pose = self.plate.get_pose().p
        eps = np.array([0.05, 0.05, 0.03])
        return (np.all(abs(container_pose[:3] - target_pose) < eps) and self.is_left_gripper_open()
                and self.is_right_gripper_open())
