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


class shake_bottle_horizontally(Base_Task):

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

    def load_actors(self):
        self.id_list = [i for i in range(20)]
        rand_pos = rand_pose(
            xlim=[-0.15, 0.15],
            ylim=[-0.15, -0.05],
            zlim=[0.785],
            qpos=[0, 0, 1, 0],
            rotate_rand=True,
            rotate_lim=[0, 0, np.pi / 4],
        )
        while abs(rand_pos.p[0]) < 0.1:
            rand_pos = rand_pose(
                xlim=[-0.15, 0.15],
                ylim=[-0.15, -0.05],
                zlim=[0.785],
                qpos=[0, 0, 1, 0],
                rotate_rand=True,
                rotate_lim=[0, 0, np.pi / 4],
            )
        self.bottle_id = np.random.choice(self.id_list)
        self.bottle = create_actor(
            scene=self,
            pose=rand_pos,
            modelname="001_bottle",
            convex=True,
            model_id=self.bottle_id,
        )
        self.bottle.set_mass(0.01)
        self.add_prohibit_area(self.bottle, padding=0.05)

    def play_once(self):
        # Determine which arm to use based on bottle position
        arm_tag = ArmTag("right" if self.bottle.get_pose().p[0] > 0 else "left")

        # Grasp the bottle with specified pre-grasp distance
        self.move(self.grasp_actor(self.bottle, arm_tag=arm_tag, pre_grasp_dis=0.1))

        # Lift the bottle up while setting initial orientation (quaternion)
        target_quat = [0.707, 0, 0, 0.707]
        self.move(self.move_by_displacement(arm_tag=arm_tag, z=0.1, quat=target_quat))

        # Rotate the bottle 90 degrees around y-axis (π/2)
        y_rotation = t3d.euler.euler2quat(0, (np.pi / 2), 0)
        rotated_q = t3d.quaternions.qmult(y_rotation, target_quat)
        target_quat = [-rotated_q[1], rotated_q[0], rotated_q[3], -rotated_q[2]]
        self.move(self.move_by_displacement(arm_tag=arm_tag, quat=target_quat))

        # Prepare two quaternions for shaking motion
        quat1 = deepcopy(target_quat)
        quat2 = deepcopy(target_quat)
        # First shake rotation: 157.5 degrees (7π/8) clockwise around y-axis
        y_rotation = t3d.euler.euler2quat(0, (np.pi / 8) * 7, 0)
        rotated_q = t3d.quaternions.qmult(y_rotation, quat1)
        quat1 = [-rotated_q[1], rotated_q[0], rotated_q[3], -rotated_q[2]]

        # Second shake rotation: 157.5 degrees (7π/8) counter-clockwise around y-axis
        y_rotation = t3d.euler.euler2quat(0, -7 * (np.pi / 8), 0)
        rotated_q = t3d.quaternions.qmult(y_rotation, quat2)
        quat2 = [-rotated_q[1], rotated_q[0], rotated_q[3], -rotated_q[2]]
        # Move up with first rotation
        self.move(self.move_by_displacement(arm_tag=arm_tag, z=0., quat=quat1))
        # Move down with opposite rotation
        self.move(self.move_by_displacement(arm_tag=arm_tag, z=-0.03, quat=quat2))

        # Perform shaking motion (up and down while rotating) three times
        for _ in range(2):
            # Move up with first rotation
            self.move(self.move_by_displacement(arm_tag=arm_tag, z=0.05, quat=quat1))
            # Move down with opposite rotation
            self.move(self.move_by_displacement(arm_tag=arm_tag, z=-0.05, quat=quat2))

        # Return to the orientation after initial tilt (before shaking)
        self.move(self.move_by_displacement(arm_tag=arm_tag, quat=target_quat))

        self.info["info"] = {
            "{A}": f"001_bottle/base{self.bottle_id}",
            "{a}": str(arm_tag),
        }
        return self.info

    def check_success(self):
        bottle_pose = self.bottle.get_pose().p
        return bottle_pose[2] > 0.8 + self.table_z_bias
