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


class place_cans_plasticbox(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.5, 0.5, 0.5, 0.5],
            rotate_rand=True,
            rotate_lim=[0, 0, 0],
        )

        self.plasticbox_id = np.random.choice([3, 5], 1)[0]

        self.plasticbox = create_actor(
            scene=self,
            pose=rand_pos_1,
            modelname="062_plasticbox",
            convex=True,
            model_id=self.plasticbox_id,
        )
        self.plasticbox.set_mass(0.05)

        rand_pos_2 = rand_pose(
            xlim=[-0.25, -0.15],
            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, 5, 6], 1)[0]

        self.object1 = create_actor(
            scene=self,
            pose=rand_pos_2,
            modelname="071_can",
            convex=True,
            model_id=self.object1_id,
        )
        self.object1.set_mass(0.05)

        rand_pos_3 = rand_pose(
            xlim=[0.15, 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.object2_id = np.random.choice([0, 1, 2, 3, 5, 6], 1)[0]

        self.object2 = create_actor(
            scene=self,
            pose=rand_pos_3,
            modelname="071_can",
            convex=True,
            model_id=self.object2_id,
        )
        self.object2.set_mass(0.05)

        self.add_prohibit_area(self.plasticbox, padding=0.1)
        self.add_prohibit_area(self.object1, padding=0.05)
        self.add_prohibit_area(self.object2, padding=0.05)

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

        # Grasp both objects with dual arms
        self.move(
            self.grasp_actor(self.object1, arm_tag=arm_tag_left, pre_grasp_dis=0.1),
            self.grasp_actor(self.object2, arm_tag=arm_tag_right, pre_grasp_dis=0.1),
        )

        # Lift up both arms after grasping
        self.move(
            self.move_by_displacement(arm_tag=arm_tag_left, z=0.2),
            self.move_by_displacement(arm_tag=arm_tag_right, z=0.2),
        )

        # Place left object into plastic box at target point 1
        self.move(
            self.place_actor(
                self.object1,
                arm_tag=arm_tag_left,
                target_pose=self.plasticbox.get_functional_point(1),
                constrain="free",
                pre_dis=0.1,
            ))

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

        # Left arm moves back to origin while right arm places object into plastic box at target point 0
        self.move(
            self.back_to_origin(arm_tag=arm_tag_left),
            self.place_actor(
                self.object2,
                arm_tag=arm_tag_right,
                target_pose=self.plasticbox.get_functional_point(0),
                constrain="free",
                pre_dis=0.1,
            ),
        )

        self.move(self.move_by_displacement(arm_tag=arm_tag_right, z=0.08))
        # Right arm moves back to original position
        self.move(self.back_to_origin(arm_tag=arm_tag_right))

        self.info["info"] = {
            "{A}": f"071_can/base{self.object1_id}",
            "{B}": f"062_plasticbox/base{self.plasticbox_id}",
            "{C}": f"071_can/base{self.object2_id}",
        }
        return self.info

    def check_success(self):
        plasticbox_functional_points_0 = self.plasticbox.get_functional_point(0)[0:2]
        plasticbox_functional_points_1 = self.plasticbox.get_functional_point(1)[0:2]
        dis1 = min(np.linalg.norm(self.object1.get_pose().p[0:2] - plasticbox_functional_points_0),
                   np.linalg.norm(self.object1.get_pose().p[0:2] - plasticbox_functional_points_1))
        dis2 = min(np.linalg.norm(self.object2.get_pose().p[0:2] - plasticbox_functional_points_0),
                   np.linalg.norm(self.object2.get_pose().p[0:2] - plasticbox_functional_points_1))
        threshold = 0.04
        return dis1 < threshold and dis2 < threshold and self.is_left_gripper_open() and self.is_right_gripper_open()