from ._base_task import Base_Task
from .utils import *
import numpy as np
from ._GLOBAL_CONFIGS import *


class hanging_mug(Base_Task):

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

    def load_actors(self):
        self.mug_id = np.random.choice([i for i in range(10)])
        self.mug = rand_create_actor(
            self,
            xlim=[-0.25, -0.1],
            ylim=[-0.05, 0.05],
            ylim_prop=True,
            modelname="039_mug",
            rotate_rand=True,
            rotate_lim=[0, 1.57, 0],
            qpos=[0.707, 0.707, 0, 0],
            convex=True,
            model_id=self.mug_id,
        )

        rack_pose = rand_pose(
            xlim=[0.1, 0.3],
            ylim=[0.13, 0.17],
            rotate_rand=True,
            rotate_lim=[0, 0.2, 0],
            qpos=[-0.22, -0.22, 0.67, 0.67],
        )

        self.rack = create_actor(self, pose=rack_pose, modelname="040_rack", is_static=True, convex=True)

        self.add_prohibit_area(self.mug, padding=0.1)
        self.add_prohibit_area(self.rack, padding=0.1)
        self.middle_pos = [0.0, -0.15, 0.75, 1, 0, 0, 0]

    def play_once(self):
        # Initialize arm tags for grasping and hanging
        grasp_arm_tag = ArmTag("left")
        hang_arm_tag = ArmTag("right")

        # Move the grasping arm to the mug's position and grasp it
        self.move(self.grasp_actor(self.mug, arm_tag=grasp_arm_tag, pre_grasp_dis=0.05))
        self.move(self.move_by_displacement(arm_tag=grasp_arm_tag, z=0.08))

        # Move the grasping arm to a middle position before hanging
        self.move(
            self.place_actor(self.mug,
                             arm_tag=grasp_arm_tag,
                             target_pose=self.middle_pos,
                             pre_dis=0.05,
                             dis=0.0,
                             constrain="free"))
        self.move(self.move_by_displacement(arm_tag=grasp_arm_tag, z=0.1))

        # Grasp the mug with the hanging arm, and move the grasping arm back to its origin
        self.move(self.back_to_origin(grasp_arm_tag),
                  self.grasp_actor(self.mug, arm_tag=hang_arm_tag, pre_grasp_dis=0.05))
        self.move(self.move_by_displacement(arm_tag=hang_arm_tag, z=0.1, quat=GRASP_DIRECTION_DIC['front']))

        # Target pose for hanging the mug is the functional point of the rack
        target_pose = self.rack.get_functional_point(0)
        # Move the hanging arm to the target pose and hang the mug
        self.move(
            self.place_actor(self.mug,
                             arm_tag=hang_arm_tag,
                             target_pose=target_pose,
                             functional_point_id=0,
                             constrain="align",
                             pre_dis=0.05,
                             dis=-0.05,
                             pre_dis_axis='fp'))
        self.move(self.move_by_displacement(arm_tag=hang_arm_tag, z=0.1, move_axis='arm'))
        self.info["info"] = {"{A}": f"039_mug/base{self.mug_id}", "{B}": "040_rack/base0"}
        return self.info

    def check_success(self):
        mug_function_pose = self.mug.get_functional_point(0)[:3]
        rack_pose = self.rack.get_pose().p
        rack_function_pose = self.rack.get_functional_point(0)[:3]
        rack_middle_pose = (rack_pose + rack_function_pose) / 2
        eps = 0.02
        return (np.all(abs((mug_function_pose - rack_middle_pose)[:2]) < eps) and self.is_right_gripper_open()
                and mug_function_pose[2] > 0.86)
