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


class place_burger_fries(Base_Task):

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

    def load_actors(self):
        rand_pos_1 = rand_pose(
            xlim=[-0.0, 0.0],
            ylim=[-0.15, -0.1],
            qpos=[0.706527, 0.706483, -0.0291356, -0.0291767],
            rotate_rand=True,
            rotate_lim=[0, 0, 0],
        )
        self.tray_id = np.random.choice([0, 1, 2, 3], 1)[0]
        self.tray = create_actor(
            scene=self,
            pose=rand_pos_1,
            modelname="008_tray",
            convex=True,
            model_id=self.tray_id,
            scale=(2.0, 2.0, 2.0),
            is_static=True,
        )
        self.tray.set_mass(0.05)

        rand_pos_2 = rand_pose(
            xlim=[-0.3, -0.25],
            ylim=[-0.15, -0.07],
            qpos=[0.5, 0.5, 0.5, 0.5],
            rotate_rand=True,
            rotate_lim=[0, 0, 0],
        )
        self.object1_id = np.random.choice([0, 1, 2, 3, 4, 5], 1)[0]
        self.hamburg = create_actor(
            scene=self,
            pose=rand_pos_2,
            modelname="006_hamburg",
            convex=True,
            model_id=self.object1_id,
        )
        self.hamburg.set_mass(0.05)

        rand_pos_3 = rand_pose(
            xlim=[0.2, 0.3],
            ylim=[-0.15, -0.07],
            qpos=[1.0, 0.0, 0.0, 0.0],
            rotate_rand=True,
            rotate_lim=[0, 0, 0],
        )
        self.object2_id = np.random.choice([0, 1], 1)[0]
        self.frenchfries = create_actor(
            scene=self,
            pose=rand_pos_3,
            modelname="005_french-fries",
            convex=True,
            model_id=self.object2_id,
        )
        self.frenchfries.set_mass(0.05)

        self.add_prohibit_area(self.tray, padding=0.1)
        self.add_prohibit_area(self.hamburg, padding=0.05)
        self.add_prohibit_area(self.frenchfries, padding=0.05)

    def play_once(self):
        arm_tag_left = ArmTag("left")
        arm_tag_right = ArmTag("right")

        # Dual grasp of hamburg and french fries
        self.move(
            self.grasp_actor(self.hamburg, arm_tag=arm_tag_left, pre_grasp_dis=0.1),
            self.grasp_actor(self.frenchfries, arm_tag=arm_tag_right, pre_grasp_dis=0.1),
        )

        # Move up before placing
        self.move(
            self.move_by_displacement(arm_tag=arm_tag_left, z=0.1),
            self.move_by_displacement(arm_tag=arm_tag_right, z=0.1),
        )

        # Get target poses from tray for placing
        tray_place_pose_left = self.tray.get_functional_point(0)
        tray_place_pose_right = self.tray.get_functional_point(1)

        # Place hamburg on tray
        self.move(
            self.place_actor(self.hamburg,
                             arm_tag=arm_tag_left,
                             target_pose=tray_place_pose_left,
                             functional_point_id=0,
                             constrain="free",
                             pre_dis=0.1,
                             pre_dis_axis='fp'), )

        # Move up after placing
        self.move(self.move_by_displacement(arm_tag=arm_tag_left, z=0.08), )

        self.move(
            self.place_actor(self.frenchfries,
                             arm_tag=arm_tag_right,
                             target_pose=tray_place_pose_right,
                             functional_point_id=0,
                             constrain="free",
                             pre_dis=0.1,
                             pre_dis_axis='fp'),
            self.back_to_origin(arm_tag=arm_tag_left),
        )

        self.move(self.move_by_displacement(arm_tag=arm_tag_right, z=0.08))

        self.info['info'] = {
            "{A}": f"006_hamburg/base{self.object1_id}",
            "{B}": f"008_tray/base{self.tray_id}",
            "{C}": f"005_french-fries/base{self.object2_id}",
        }
        return self.info

    def check_success(self):
        dis1 = np.linalg.norm(
            self.tray.get_functional_point(0, "pose").p[0:2] - self.hamburg.get_functional_point(0, "pose").p[0:2])
        dis2 = np.linalg.norm(
            self.tray.get_functional_point(1, "pose").p[0:2] - self.frenchfries.get_functional_point(0, "pose").p[0:2])
        threshold = 0.08
        return dis1 < threshold and dis2 < threshold and self.is_left_gripper_open() and self.is_right_gripper_open()
