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


class put_bottles_dustbin(Base_Task):

    def setup_demo(self, **kwags):
        super()._init_task_env_(table_xy_bias=[0.3, 0], **kwags)

    def load_actors(self):
        pose_lst = []

        def create_bottle(model_id):
            bottle_pose = rand_pose(
                xlim=[-0.25, 0.3],
                ylim=[0.03, 0.23],
                rotate_rand=False,
                rotate_lim=[0, 1, 0],
                qpos=[0.707, 0.707, 0, 0],
            )
            tag = True
            gen_lim = 100
            i = 1
            while tag and i < gen_lim:
                tag = False
                if np.abs(bottle_pose.p[0]) < 0.05:
                    tag = True
                for pose in pose_lst:
                    if (np.sum(np.power(np.array(pose[:2]) - np.array(bottle_pose.p[:2]), 2)) < 0.0169):
                        tag = True
                        break
                if tag:
                    i += 1
                    bottle_pose = rand_pose(
                        xlim=[-0.25, 0.3],
                        ylim=[0.03, 0.23],
                        rotate_rand=False,
                        rotate_lim=[0, 1, 0],
                        qpos=[0.707, 0.707, 0, 0],
                    )
            pose_lst.append(bottle_pose.p[:2])
            bottle = create_actor(
                self,
                bottle_pose,
                modelname="114_bottle",
                convex=True,
                model_id=model_id,
            )

            return bottle

        self.bottles = []
        self.bottles_data = []
        self.bottle_id = [1, 2, 3]
        self.bottle_num = 3
        for i in range(self.bottle_num):
            bottle = create_bottle(self.bottle_id[i])
            self.bottles.append(bottle)
            self.add_prohibit_area(bottle, padding=0.1)

        self.dustbin = create_actor(
            self.scene,
            pose=sapien.Pose([-0.45, 0, 0], [0.5, 0.5, 0.5, 0.5]),
            modelname="011_dustbin",
            convex=True,
            is_static=True,
        )
        self.delay(2)
        self.right_middle_pose = [0, 0.0, 0.88, 0, 1, 0, 0]

    def play_once(self):
        # Sort bottles based on their x and y coordinates
        bottle_lst = sorted(self.bottles, key=lambda x: [x.get_pose().p[0] > 0, x.get_pose().p[1]])

        for i in range(self.bottle_num):
            bottle = bottle_lst[i]
            # Determine which arm to use based on bottle's x position
            arm_tag = ArmTag("left" if bottle.get_pose().p[0] < 0 else "right")

            delta_dis = 0.06

            # Define end position for left arm
            left_end_action = Action("left", "move", [-0.35, -0.1, 0.93, 0.65, -0.25, 0.25, 0.65])

            if arm_tag == "left":
                # Grasp the bottle with left arm
                self.move(self.grasp_actor(bottle, arm_tag=arm_tag, pre_grasp_dis=0.1))
                # Move left arm up
                self.move(self.move_by_displacement(arm_tag, z=0.1))
                # Move left arm to end position
                self.move((ArmTag("left"), [left_end_action]))
            else:
                # Grasp the bottle with right arm while moving left arm to origin
                right_action = self.grasp_actor(bottle, arm_tag=arm_tag, pre_grasp_dis=0.1)
                right_action[1][0].target_pose[2] += delta_dis
                right_action[1][1].target_pose[2] += delta_dis
                self.move(right_action, self.back_to_origin("left"))
                # Move right arm up
                self.move(self.move_by_displacement(arm_tag, z=0.1))
                # Place the bottle at middle position with right arm
                self.move(
                    self.place_actor(
                        bottle,
                        target_pose=self.right_middle_pose,
                        arm_tag=arm_tag,
                        functional_point_id=0,
                        pre_dis=0.0,
                        dis=0.0,
                        is_open=False,
                        constrain="align",
                    ))
                # Grasp the bottle with left arm (adjusted height)
                left_action = self.grasp_actor(bottle, arm_tag="left", pre_grasp_dis=0.1)
                left_action[1][0].target_pose[2] -= delta_dis
                left_action[1][1].target_pose[2] -= delta_dis
                self.move(left_action)
                # Open right gripper
                self.move(self.open_gripper(ArmTag("right")))
                # Move left arm to end position while moving right arm to origin
                self.move((ArmTag("left"), [left_end_action]), self.back_to_origin("right"))
            # Open left gripper
            self.move(self.open_gripper("left"))

        self.info["info"] = {
            "{A}": f"114_bottle/base{self.bottle_id[0]}",
            "{B}": f"114_bottle/base{self.bottle_id[1]}",
            "{C}": f"114_bottle/base{self.bottle_id[2]}",
            "{D}": f"011_dustbin/base0",
        }
        return self.info

    def stage_reward(self):
        taget_pose = [-0.45, 0]
        eps = np.array([0.221, 0.325])
        reward = 0
        reward_step = 1 / 3
        for i in range(self.bottle_num):
            bottle_pose = self.bottles[i].get_pose().p
            if (np.all(np.abs(bottle_pose[:2] - taget_pose) < eps) and bottle_pose[2] > 0.2 and bottle_pose[2] < 0.7):
                reward += reward_step
        return reward

    def check_success(self):
        taget_pose = [-0.45, 0]
        eps = np.array([0.221, 0.325])
        for i in range(self.bottle_num):
            bottle_pose = self.bottles[i].get_pose().p
            if (np.all(np.abs(bottle_pose[:2] - taget_pose) < eps) and bottle_pose[2] > 0.2 and bottle_pose[2] < 0.7):
                continue
            return False
        return True
