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


class place_object_basket(Base_Task):

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

    def load_actors(self):
        self.arm_tag = ArmTag({0: "left", 1: "right"}[np.random.randint(0, 2)])
        self.basket_name = "110_basket"
        self.basket_id = np.random.randint(0, 2)
        toycar_dict = {
            "081_playingcards": [0, 1, 2],
            "057_toycar": [0, 1, 2, 3, 4, 5],
        }
        self.object_name = ["081_playingcards", "057_toycar"][np.random.randint(0, 2)]
        self.object_id = toycar_dict[self.object_name][np.random.randint(0, len(toycar_dict[self.object_name]))]
        if self.arm_tag == "left":  # toycar on left
            self.basket = rand_create_actor(
                scene=self,
                modelname=self.basket_name,
                model_id=self.basket_id,
                xlim=[0.02, 0.02],
                ylim=[-0.08, -0.05],
                qpos=[0.5, 0.5, 0.5, 0.5],
                convex=True,
            )
            self.object = rand_create_actor(
                scene=self,
                modelname=self.object_name,
                model_id=self.object_id,
                xlim=[-0.25, -0.2],
                ylim=[-0.1, 0.1],
                rotate_rand=True,
                rotate_lim=[0, np.pi / 6, 0],
                qpos=[0.707225, 0.706849, -0.0100455, -0.00982061],
                convex=True,
            )
        else:  # toycar on right
            self.basket = rand_create_actor(
                scene=self,
                modelname=self.basket_name,
                model_id=self.basket_id,
                xlim=[-0.02, -0.02],
                ylim=[-0.08, -0.05],
                qpos=[0.5, 0.5, 0.5, 0.5],
                convex=True,
            )
            self.object = rand_create_actor(
                scene=self,
                modelname=self.object_name,
                model_id=self.object_id,
                xlim=[0.2, 0.25],
                ylim=[-0.1, 0.1],
                rotate_rand=True,
                rotate_lim=[0, np.pi / 6, 0],
                qpos=[0.707225, 0.706849, -0.0100455, -0.00982061],
                convex=True,
            )
        self.basket.set_mass(0.5)
        self.object.set_mass(0.01)
        self.object_start_height = self.object.get_pose().p[2]
        self.start_height = self.basket.get_pose().p[2]
        self.add_prohibit_area(self.object, padding=0.1)
        self.add_prohibit_area(self.basket, padding=0.05)

    def play_once(self):
        # Grasp the toy car
        self.move(self.grasp_actor(self.object, arm_tag=self.arm_tag))

        # Lift the toy car up
        self.move(self.move_by_displacement(arm_tag=self.arm_tag, z=0.15))

        # Get functional points of basket for placing
        f0 = np.array(self.basket.get_functional_point(0))
        f1 = np.array(self.basket.get_functional_point(1))
        place_pose = (f0 if np.linalg.norm(f0[:2] - self.object.get_pose().p[:2])
                      < np.linalg.norm(f1[:2] - self.object.get_pose().p[:2]) else f1)
        place_pose[:2] = f0[:2] if place_pose is f0 else f1[:2]
        place_pose[3:] = (-1, 0, 0, 0) if self.arm_tag == "left" else (0.05, 0, 0, 0.99)

        # Place the toy car in the basket
        self.move(self.place_actor(
            self.object,
            arm_tag=self.arm_tag,
            target_pose=place_pose,
            dis=0.02,
            is_open=False,
        ))

        if not self.plan_success:
            self.plan_success = True  # Try new way
            # Move up and away (recovery motion when plan fails)
            place_pose[0] += -0.15 if self.arm_tag == "left" else 0.15
            place_pose[2] += 0.15
            self.move(self.move_to_pose(arm_tag=self.arm_tag, target_pose=place_pose))

            # Lower down (recovery motion when plan fails)
            place_pose[2] -= 0.05
            self.move(self.move_to_pose(arm_tag=self.arm_tag, target_pose=place_pose))

            # Open gripper to release object
            self.move(self.open_gripper(arm_tag=self.arm_tag))

            # Move arm away and grasp the basket with opposite arm (recovery strategy)
            self.move(
                self.back_to_origin(arm_tag=self.arm_tag),
                self.grasp_actor(self.basket, arm_tag=self.arm_tag.opposite, pre_grasp_dis=0.02),
            )
        else:
            # Open gripper to release object
            self.move(self.open_gripper(arm_tag=self.arm_tag))
            # lift arm up, to avoid collision with the basket
            self.move(self.move_by_displacement(arm_tag=self.arm_tag, z=0.08))
            # Move arm away and grasp the basket with opposite arm
            self.move(
                self.back_to_origin(arm_tag=self.arm_tag),
                self.grasp_actor(self.basket, arm_tag=self.arm_tag.opposite, pre_grasp_dis=0.08),
            )

        # Lift basket a bit after grasping
        self.move(
            self.move_by_displacement(
                arm_tag=self.arm_tag.opposite,
                x=0.05 if self.arm_tag.opposite == "right" else -0.05,
                z=0.05,
            ))

        self.info["info"] = {
            "{A}": f"{self.object_name}/base{self.object_id}",
            "{B}": f"{self.basket_name}/base{self.basket_id}",
            "{a}": str(self.arm_tag),
            "{b}": str(self.arm_tag.opposite),
        }
        return self.info

    def check_success(self):
        toy_p = self.object.get_pose().p
        basket_p = self.basket.get_pose().p
        basket_axis = (self.basket.get_pose().to_transformation_matrix()[:3, :3] @ np.array([[0, 1, 0]]).T)
        obj_contact_table = not self.check_actors_contact(self.object_name, "table")
        obj_contact_basket = self.check_actors_contact(self.object_name, self.basket_name)
        return (basket_p[2] - self.start_height > 0.02 and \
                toy_p[2] - self.object_start_height > 0.02 and \
                np.dot(basket_axis.reshape(3), [0, 0, 1]) > 0.5 and \
                np.sum(np.sqrt((toy_p - basket_p)**2)) < 0.15 and \
                obj_contact_table and obj_contact_basket)
