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


class move_pillbottle_pad(Base_Task):

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

    def load_actors(self):
        rand_pos = rand_pose(
            xlim=[-0.25, 0.25],
            ylim=[-0.1, 0.1],
            qpos=[0.5, 0.5, 0.5, 0.5],
            rotate_rand=False,
        )
        while abs(rand_pos.p[0]) < 0.05:
            rand_pos = rand_pose(
                xlim=[-0.25, 0.25],
                ylim=[-0.1, 0.1],
                qpos=[0.5, 0.5, 0.5, 0.5],
                rotate_rand=False,
            )

        self.pillbottle_id = np.random.choice([1, 2, 3, 4, 5], 1)[0]
        self.pillbottle = create_actor(
            scene=self,
            pose=rand_pos,
            modelname="080_pillbottle",
            convex=True,
            model_id=self.pillbottle_id,
        )
        self.pillbottle.set_mass(0.05)

        if rand_pos.p[0] > 0:
            xlim = [0.05, 0.25]
        else:
            xlim = [-0.25, -0.05]
        target_rand_pose = rand_pose(
            xlim=xlim,
            ylim=[-0.2, 0.1],
            qpos=[1, 0, 0, 0],
            rotate_rand=False,
        )
        while (np.sqrt((target_rand_pose.p[0] - rand_pos.p[0])**2 + (target_rand_pose.p[1] - rand_pos.p[1])**2) < 0.1):
            target_rand_pose = rand_pose(
                xlim=xlim,
                ylim=[-0.2, 0.1],
                qpos=[1, 0, 0, 0],
                rotate_rand=False,
            )
        half_size = [0.04, 0.04, 0.0005]
        self.pad = create_box(
            scene=self,
            pose=target_rand_pose,
            half_size=half_size,
            color=(0, 0, 1),
            name="box",
            is_static=True,
        )
        self.add_prohibit_area(self.pillbottle, padding=0.05)
        self.add_prohibit_area(self.pad, padding=0.1)

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

        # Grasp the pillbottle
        self.move(self.grasp_actor(self.pillbottle, arm_tag=arm_tag, pre_grasp_dis=0.06, gripper_pos=0))

        # Lift up the pillbottle by 0.1 meters in z-axis
        self.move(self.move_by_displacement(arm_tag=arm_tag, z=0.05))

        # Get the target pose for placing the pillbottle
        target_pose = self.pad.get_functional_point(1)
        # Place the pillbottle at the target pose
        self.move(
            self.place_actor(self.pillbottle,
                             arm_tag=arm_tag,
                             target_pose=target_pose,
                             pre_dis=0.05,
                             dis=0,
                             functional_point_id=0,
                             pre_dis_axis='fp'))

        self.info["info"] = {
            "{A}": f"080_pillbottle/base{self.pillbottle_id}",
            "{a}": str(arm_tag),
        }

        return self.info

    def check_success(self):
        pillbottle_pos = self.pillbottle.get_pose().p
        target_pos = self.pad.get_pose().p
        eps1 = 0.03
        return (np.all(abs(pillbottle_pos[:2] - target_pos[:2]) < np.array([eps1, eps1]))
                and np.abs(self.pillbottle.get_pose().p[2] - (0.741 + self.table_z_bias)) < 0.005
                and self.robot.is_left_gripper_open() and self.robot.is_right_gripper_open())
