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


class open_laptop(Base_Task):

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

    def load_actors(self):
        self.model_name = "015_laptop"
        self.model_id = np.random.randint(0, 11)
        self.laptop: ArticulationActor = rand_create_sapien_urdf_obj(
            scene=self,
            modelname=self.model_name,
            modelid=self.model_id,
            xlim=[-0.05, 0.05],
            ylim=[-0.1, 0.05],
            rotate_rand=True,
            rotate_lim=[0, 0, np.pi / 3],
            qpos=[0.7, 0, 0, 0.7],
            fix_root_link=True,
        )
        limit = self.laptop.get_qlimits()[0]
        self.laptop.set_qpos([limit[0] + (limit[1] - limit[0]) * 0.2])
        self.laptop.set_mass(0.01)
        self.laptop.set_properties(1, 0)
        self.add_prohibit_area(self.laptop, padding=0.1)

    def play_once(self):
        face_prod = get_face_prod(self.laptop.get_pose().q, [1, 0, 0], [1, 0, 0])
        arm_tag = ArmTag("left" if face_prod > 0 else "right")
        self.arm_tag = arm_tag

        # Grasp the laptop
        self.move(self.grasp_actor(self.laptop, arm_tag=arm_tag, pre_grasp_dis=0.08, contact_point_id=0))

        for _ in range(15):
            # Get target rotation pose
            self.move(
                self.grasp_actor(
                    self.laptop,
                    arm_tag=arm_tag,
                    pre_grasp_dis=0.0,
                    grasp_dis=0.0,
                    contact_point_id=1,
                ))
            if not self.plan_success:
                break
            if self.check_success(target=0.5):
                break

        self.info["info"] = {
            "{A}": f"{self.model_name}/base{self.model_id}",
            "{a}": str(arm_tag),
        }
        return self.info

    def check_success(self, target=0.4):
        limit = self.laptop.get_qlimits()[0]
        qpos = self.laptop.get_qpos()
        rotate_pose = self.laptop.get_contact_point(1)
        tip_pose = (self.robot.get_left_tcp_pose() if self.arm_tag == "left" else self.robot.get_right_tcp_pose())
        dis = np.sqrt(np.sum((np.array(tip_pose[:3]) - np.array(rotate_pose[:3]))**2))
        return qpos[0] >= limit[0] + (limit[1] - limit[0]) * target and dis < 0.1
