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


class grab_roller(Base_Task):

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

    def load_actors(self):
        ori_qpos = [[0.5, 0.5, 0.5, 0.5], [0.5, 0.5, 0.5, 0.5], [0, 0, 0.707, 0.707]]
        self.model_id = np.random.choice([0, 2], 1)[0]
        rand_pos = rand_pose(
            xlim=[-0.15, 0.15],
            ylim=[-0.25, -0.05],
            qpos=ori_qpos[self.model_id],
            rotate_rand=True,
            rotate_lim=[0, 0.8, 0],
        )
        self.roller = create_actor(
            scene=self,
            pose=rand_pos,
            modelname="102_roller",
            convex=True,
            model_id=self.model_id,
        )

        self.add_prohibit_area(self.roller, padding=0.1)

    def play_once(self):
        # Initialize arm tags for left and right arms
        left_arm_tag = ArmTag("left")
        right_arm_tag = ArmTag("right")

        # Grasp the roller with both arms simultaneously at different contact points
        self.move(
            self.grasp_actor(self.roller, left_arm_tag, pre_grasp_dis=0.08, contact_point_id=0),
            self.grasp_actor(self.roller, right_arm_tag, pre_grasp_dis=0.08, contact_point_id=1),
        )

        # Lift the roller to height 0.85 by moving both arms upward simultaneously
        self.move(
            self.move_by_displacement(left_arm_tag, z=0.85 - self.roller.get_pose().p[2]),
            self.move_by_displacement(right_arm_tag, z=0.85 - self.roller.get_pose().p[2]),
        )

        # Record information about the roller in the info dictionary
        self.info["info"] = {"{A}": f"102_roller/base{self.model_id}"}
        return self.info

    def check_success(self):
        roller_pose = self.roller.get_pose().p
        return (self.is_left_gripper_close() and self.is_right_gripper_close() and roller_pose[2] > 0.8)
