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


class pick_dual_bottles(Base_Task):

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

    def load_actors(self):
        self.bottle1 = rand_create_actor(
            self,
            xlim=[-0.25, -0.05],
            ylim=[0.03, 0.23],
            modelname="001_bottle",
            rotate_rand=True,
            rotate_lim=[0, 1, 0],
            qpos=[0.66, 0.66, -0.25, -0.25],
            convex=True,
            model_id=13,
        )

        self.bottle2 = rand_create_actor(
            self,
            xlim=[0.05, 0.25],
            ylim=[0.03, 0.23],
            modelname="001_bottle",
            rotate_rand=True,
            rotate_lim=[0, 1, 0],
            qpos=[0.65, 0.65, 0.27, 0.27],
            convex=True,
            model_id=16,
        )

        render_freq = self.render_freq
        self.render_freq = 0
        for _ in range(4):
            self.together_open_gripper(save_freq=None)
        self.render_freq = render_freq

        self.add_prohibit_area(self.bottle1, padding=0.1)
        self.add_prohibit_area(self.bottle2, padding=0.1)
        target_posi = [-0.2, -0.2, 0.2, -0.02]
        self.prohibited_area.append(target_posi)
        self.left_target_pose = [-0.06, -0.105, 1, 0, 1, 0, 0]
        self.right_target_pose = [0.06, -0.105, 1, 0, 1, 0, 0]

    def play_once(self):
        # Determine which arm to use for each bottle based on their x-coordinate position
        bottle1_arm_tag = ArmTag("left")
        bottle2_arm_tag = ArmTag("right")

        # Simultaneously grasp both bottles with their respective arms
        self.move(
            self.grasp_actor(self.bottle1, arm_tag=bottle1_arm_tag, pre_grasp_dis=0.08),
            self.grasp_actor(self.bottle2, arm_tag=bottle2_arm_tag, pre_grasp_dis=0.08),
        )

        # Simultaneously lift both bottles up by 0.1 meters
        self.move(
            self.move_by_displacement(arm_tag=bottle1_arm_tag, z=0.1),
            self.move_by_displacement(arm_tag=bottle2_arm_tag, z=0.1),
        )

        # Simultaneously place both bottles at their target positions
        self.move(
            self.place_actor(
                self.bottle1,
                target_pose=self.left_target_pose,
                arm_tag=bottle1_arm_tag,
                functional_point_id=0,
                pre_dis=0.0,
                dis=0.0,
                is_open=False,
            ),
            self.place_actor(
                self.bottle2,
                target_pose=self.right_target_pose,
                arm_tag=bottle2_arm_tag,
                functional_point_id=0,
                pre_dis=0.0,
                dis=0.0,
                is_open=False,
            ),
        )

        self.info["info"] = {"{A}": f"001_bottle/base13", "{B}": f"001_bottle/base16"}
        return self.info

    def check_success(self):
        bottle1_target = self.left_target_pose[:2]
        bottle2_target = self.right_target_pose[:2]
        eps = 0.1
        bottle1_pose = self.bottle1.get_functional_point(0)
        bottle2_pose = self.bottle2.get_functional_point(0)
        if bottle1_pose[2] < 0.78 or bottle2_pose[2] < 0.78:
            self.actor_pose = False
        return (abs(bottle1_pose[0] - bottle1_target[0]) < eps and abs(bottle1_pose[1] - bottle1_target[1]) < eps
                and bottle1_pose[2] > 0.89 and abs(bottle2_pose[0] - bottle2_target[0]) < eps
                and abs(bottle2_pose[1] - bottle2_target[1]) < eps and bottle2_pose[2] > 0.89)
