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


class move_can_pot(Base_Task):

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

    def load_actors(self):
        self.pot_id = np.random.randint(0, 7)
        self.pot = rand_create_sapien_urdf_obj(
            scene=self,
            modelname="060_kitchenpot",
            modelid=self.pot_id,
            xlim=[0.0, 0.0],
            ylim=[0.0, 0.0],
            rotate_rand=True,
            rotate_lim=[0, 0, np.pi / 8],
            qpos=[0, 0, 0, 1],
        )
        pot_pose = self.pot.get_pose()
        rand_pos = rand_pose(
            xlim=[-0.3, 0.3],
            ylim=[0.05, 0.15],
            qpos=[0.5, 0.5, 0.5, 0.5],
            rotate_rand=True,
            rotate_lim=[0, np.pi / 4, 0],
        )
        while abs(rand_pos.p[0]) < 0.2 or (((pot_pose.p[0] - rand_pos.p[0])**2 +
                                            (pot_pose.p[1] - rand_pos.p[1])**2) < 0.09):
            rand_pos = rand_pose(
                xlim=[-0.3, 0.3],
                ylim=[0.05, 0.15],
                qpos=[0.5, 0.5, 0.5, 0.5],
                rotate_rand=True,
                rotate_lim=[0, np.pi / 4, 0],
            )
        id_list = [0, 2, 4, 5, 6]
        self.can_id = np.random.choice(id_list)
        self.can = create_actor(
            scene=self,
            pose=rand_pos,
            modelname="105_sauce-can",
            convex=True,
            model_id=self.can_id,
        )
        self.arm_tag = ArmTag("right" if self.can.get_pose().p[0] > 0 else "left")
        self.add_prohibit_area(self.pot, padding=0.03)
        self.add_prohibit_area(self.can, padding=0.1)
        pot_x, pot_y = self.pot.get_pose().p[0], self.pot.get_pose().p[1]
        if self.arm_tag == "left":
            self.prohibited_area.append([pot_x - 0.15, pot_y - 0.1, pot_x, pot_y + 0.1])
        else:
            self.prohibited_area.append([pot_x, pot_y - 0.1, pot_x + 0.15, pot_y + 0.1])
        self.orig_z = self.pot.get_pose().p[2]

        # Get pot's current pose and calculate target pose for placing the can
        pot_pose = self.pot.get_pose()
        self.target_pose = sapien.Pose(
            [
                pot_pose.p[0] - 0.18 if self.arm_tag == "left" else pot_pose.p[0] + 0.18,
                pot_pose.p[1],
                0.741 + self.table_z_bias,
            ],
            pot_pose.q,
        )

    def play_once(self):
        arm_tag = self.arm_tag
        # Grasp the can with specified pre-grasp distance
        self.move(self.grasp_actor(self.can, arm_tag=arm_tag, pre_grasp_dis=0.05))
        # Move the can backward and upward
        self.move(self.move_by_displacement(arm_tag, y=-0.1, z=0.1))

        # Place the can near the pot at calculated target pose
        self.move(self.place_actor(
            self.can,
            target_pose=self.target_pose,
            arm_tag=arm_tag,
            pre_dis=0.05,
            dis=0.0,
        ))

        self.info["info"] = {
            "{A}": f"060_kitchenpot/base{self.pot_id}",
            "{B}": f"105_sauce-can/base{self.can_id}",
            "{a}": str(arm_tag),
        }
        return self.info

    def check_success(self):
        pot_pose = self.pot.get_pose().p
        can_pose = self.can.get_pose().p
        can_pose_rpy = t3d.euler.quat2euler(self.can.get_pose().q)
        x_rotate = can_pose_rpy[0] * 180 / np.pi
        y_rotate = can_pose_rpy[1] * 180 / np.pi
        eps = np.array([0.2, 0.035, 15, 15])
        dis = (pot_pose[0] - can_pose[0] if self.arm_tag == "left" else can_pose[0] - pot_pose[0])
        check = True if dis > 0 else False
        return (np.all(np.array([
            abs(dis),
            np.abs(pot_pose[1] - can_pose[1]),
            abs(x_rotate - 90),
            abs(y_rotate),
        ]) < eps) and check and can_pose[2] <= self.orig_z + 0.001 and self.robot.is_left_gripper_open()
                and self.robot.is_right_gripper_open())
