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


class place_bread_skillet(Base_Task):

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

    def load_actors(self):
        id_list = [0, 1, 3, 5, 6]
        self.bread_id = np.random.choice(id_list)
        rand_pos = rand_pose(
            xlim=[-0.28, 0.28],
            ylim=[-0.2, 0.05],
            qpos=[0.707, 0.707, 0.0, 0.0],
            rotate_rand=True,
            rotate_lim=[0, np.pi / 4, 0],
        )
        while abs(rand_pos.p[0]) < 0.2:
            rand_pos = rand_pose(
                xlim=[-0.28, 0.28],
                ylim=[-0.2, 0.05],
                qpos=[0.707, 0.707, 0.0, 0.0],
                rotate_rand=True,
                rotate_lim=[0, np.pi / 4, 0],
            )
        self.bread = create_actor(
            self,
            pose=rand_pos,
            modelname="075_bread",
            model_id=self.bread_id,
            convex=True,
        )

        xlim = [0.15, 0.25] if rand_pos.p[0] < 0 else [-0.25, -0.15]
        self.model_id_list = [0, 1, 2, 3]
        self.skillet_id = np.random.choice(self.model_id_list)
        rand_pos = rand_pose(
            xlim=xlim,
            ylim=[-0.2, 0.05],
            qpos=[0, 0, 0.707, 0.707],
            rotate_rand=True,
            rotate_lim=[0, np.pi / 6, 0],
        )
        self.skillet = create_actor(
            self,
            pose=rand_pos,
            modelname="106_skillet",
            model_id=self.skillet_id,
            convex=True,
        )

        self.bread.set_mass(0.001)
        self.skillet.set_mass(0.01)
        self.add_prohibit_area(self.bread, padding=0.03)
        self.add_prohibit_area(self.skillet, padding=0.05)

    def play_once(self):
        # Determine which arm to use based on skillet position (right if on positive x, left otherwise)
        arm_tag = ArmTag("right" if self.skillet.get_pose().p[0] > 0 else "left")

        # Grasp the skillet and bread simultaneously with dual arms
        self.move(
            self.grasp_actor(self.skillet, arm_tag=arm_tag, pre_grasp_dis=0.07, gripper_pos=0),
            self.grasp_actor(self.bread, arm_tag=arm_tag.opposite, pre_grasp_dis=0.07, gripper_pos=0),
        )

        # Lift both arms
        self.move(
            self.move_by_displacement(arm_tag=arm_tag, z=0.1, move_axis="arm"),
            self.move_by_displacement(arm_tag=arm_tag.opposite, z=0.1),
        )

        # Define a custom target pose for the skillet based on its side (left or right)
        target_pose = self.get_arm_pose(arm_tag=arm_tag)
        if arm_tag == "left":
            # Set specific position and orientation for left arm
            target_pose[:2] = [-0.1, -0.05]
            target_pose[2] -= 0.05
            target_pose[3:] = [-0.707, 0, -0.707, 0]
        else:
            # Set specific position and orientation for right arm
            target_pose[:2] = [0.1, -0.05]
            target_pose[2] -= 0.05
            target_pose[3:] = [0, 0.707, 0, -0.707]

        # Place the skillet to the defined target pose
        self.move(self.move_to_pose(arm_tag=arm_tag, target_pose=target_pose))

        # Get the functional point of the skillet as placement target for the bread
        target_pose = self.skillet.get_functional_point(0)

        # Place the bread onto the skillet
        self.move(
            self.place_actor(
                self.bread,
                target_pose=target_pose,
                arm_tag=arm_tag.opposite,
                constrain="free",
                pre_dis=0.05,
                dis=0.05,
            ))

        self.info["info"] = {
            "{A}": f"106_skillet/base{self.skillet_id}",
            "{B}": f"075_bread/base{self.bread_id}",
            "{a}": str(arm_tag),
        }
        return self.info

    def check_success(self):
        target_pose = self.skillet.get_functional_point(0)
        bread_pose = self.bread.get_pose().p
        return (np.all(abs(target_pose[:2] - bread_pose[:2]) < [0.035, 0.035])
                and target_pose[2] > 0.76 + self.table_z_bias and bread_pose[2] > 0.76 + self.table_z_bias)
