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


class open_microwave(Base_Task):

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

    def load_actors(self):
        self.model_name = "044_microwave"
        self.model_id = np.random.randint(0, 2)
        self.microwave = rand_create_sapien_urdf_obj(
            scene=self,
            modelname=self.model_name,
            modelid=self.model_id,
            xlim=[-0.12, -0.02],
            ylim=[0.15, 0.2],
            zlim=[0.8, 0.8],
            qpos=[0.707, 0, 0, 0.707],
            fix_root_link=True,
        )
        self.microwave.set_mass(0.01)
        self.microwave.set_properties(0.0, 0.0)

        self.add_prohibit_area(self.microwave)
        self.prohibited_area.append([-0.25, -0.25, 0.25, 0.1])

    def play_once(self):
        arm_tag = ArmTag("left")

        # Grasp the microwave with pre-grasp displacement
        self.move(self.grasp_actor(self.microwave, arm_tag=arm_tag, pre_grasp_dis=0.08, contact_point_id=0))

        start_qpos = self.microwave.get_qpos()[0]
        for _ in range(50):
            # Rotate microwave
            self.move(
                self.grasp_actor(
                    self.microwave,
                    arm_tag=arm_tag,
                    pre_grasp_dis=0.0,
                    grasp_dis=0.0,
                    contact_point_id=4,
                ))

            new_qpos = self.microwave.get_qpos()[0]
            if new_qpos - start_qpos <= 0.001:
                break
            start_qpos = new_qpos
            if not self.plan_success:
                break
            if self.check_success(target=0.7):
                break

        if not self.check_success(target=0.7):
            self.plan_success = True  # Try new way
            # Open gripper
            self.move(self.open_gripper(arm_tag=arm_tag))
            self.move(self.move_by_displacement(arm_tag=arm_tag, y=-0.05, z=0.05))

            # Grasp at contact point 1
            self.move(self.grasp_actor(self.microwave, arm_tag=arm_tag, contact_point_id=1))

            # Grasp more tightly at contact point 1
            self.move(self.grasp_actor(
                self.microwave,
                arm_tag=arm_tag,
                pre_grasp_dis=0.02,
                contact_point_id=1,
            ))

            start_qpos = self.microwave.get_qpos()[0]
            for _ in range(30):
                # Rotate microwave using contact point 2
                self.move(
                    self.grasp_actor(
                        self.microwave,
                        arm_tag=arm_tag,
                        pre_grasp_dis=0.0,
                        grasp_dis=0.0,
                        contact_point_id=2,
                    ))

                new_qpos = self.microwave.get_qpos()[0]
                if new_qpos - start_qpos <= 0.001:
                    break
                start_qpos = new_qpos
                if not self.plan_success:
                    break
                if self.check_success(target=0.7):
                    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.6):
        limits = self.microwave.get_qlimits()
        qpos = self.microwave.get_qpos()
        return qpos[0] >= limits[0][1] * target
