from ._base_task import Base_Task
from .utils import *
import sapien
import math
from copy import deepcopy
import numpy as np


class place_bread_basket(Base_Task):

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

    def load_actors(self):
        rand_pos = rand_pose(
            xlim=[0.0, 0.0],
            ylim=[-0.2, -0.2],
            qpos=[0.5, 0.5, 0.5, 0.5],
            rotate_rand=True,
            rotate_lim=[0, 3.14, 0],
        )
        id_list = [0, 1, 2, 3, 4]
        self.basket_id = np.random.choice(id_list)
        self.breadbasket = create_actor(
            scene=self,
            pose=rand_pos,
            modelname="076_breadbasket",
            convex=True,
            model_id=self.basket_id,
        )

        breadbasket_pose = self.breadbasket.get_pose()
        self.bread: list[Actor] = []
        self.bread_id = []

        for i in range(2):
            rand_pos = rand_pose(
                xlim=[-0.27, 0.27],
                ylim=[-0.2, 0.05],
                qpos=[0.707, 0.707, 0.0, 0.0],
                rotate_rand=True,
                rotate_lim=[0, np.pi / 4, 0],
            )
            try_num = 0
            while True:
                pd = True
                try_num += 1
                if try_num > 50:
                    try_num = -1
                    break
                try_num0 = 0
                while (abs(rand_pos.p[0]) < 0.15 or ((rand_pos.p[0] - breadbasket_pose.p[0])**2 +
                                                     (rand_pos.p[1] - breadbasket_pose.p[1])**2) < 0.01):
                    try_num0 += 1
                    rand_pos = rand_pose(
                        xlim=[-0.27, 0.27],
                        ylim=[-0.2, 0.05],
                        qpos=[0.707, 0.707, 0.0, 0.0],
                        rotate_rand=True,
                        rotate_lim=[0, np.pi / 4, 0],
                    )
                    if try_num0 > 50:
                        try_num = -1
                        break
                if try_num == -1:
                    break
                for j in range(len(self.bread)):
                    peer_pose = self.bread[j].get_pose()
                    if ((peer_pose.p[0] - rand_pos.p[0])**2 + (peer_pose.p[1] - rand_pos.p[1])**2) < 0.01:
                        pd = False
                        break
                if pd:
                    break
            if try_num == -1:
                break
            id_list = [0, 1, 3, 5, 6]
            self.bread_id.append(np.random.choice(id_list))
            bread_actor = create_actor(
                scene=self,
                pose=rand_pos,
                modelname="075_bread",
                convex=True,
                model_id=self.bread_id[i],
            )
            self.bread.append(bread_actor)

        for i in range(len(self.bread)):
            self.add_prohibit_area(self.bread[i], padding=0.03)

        self.add_prohibit_area(self.breadbasket, padding=0.05)

    def play_once(self):

        def remove_bread(id, num):
            arm_tag = ArmTag("right" if self.bread[id].get_pose().p[0] > 0 else "left")

            # Grasp the bread
            self.move(self.grasp_actor(self.bread[id], arm_tag=arm_tag, pre_grasp_dis=0.07))
            # Move up a little
            self.move(self.move_by_displacement(arm_tag=arm_tag, z=0.1, move_axis="arm"))

            # Get bread basket's functional point as target pose
            breadbasket_pose = self.breadbasket.get_functional_point(0)
            # Place the bread into the bread basket
            self.move(
                self.place_actor(
                    self.bread[id],
                    arm_tag=arm_tag,
                    target_pose=breadbasket_pose,
                    constrain="free",
                    pre_dis=0.12,
                ))
            if num == 0:
                # Move up further after placing first bread
                self.move(self.move_by_displacement(arm_tag=arm_tag, z=0.15, move_axis="arm"))
            else:
                # Open gripper to place the second bread
                self.move(self.open_gripper(arm_tag=arm_tag))

        def remove():
            # Determine which bread is on the left
            id = 0 if self.bread[0].get_pose().p[0] < 0 else 1

            # Simultaneously grasp both breads with dual arms
            self.move(
                self.grasp_actor(self.bread[id], arm_tag="left", pre_grasp_dis=0.05),
                self.grasp_actor(self.bread[id ^ 1], arm_tag="right", pre_grasp_dis=0.07),
            )

            # Lift both arms slightly after grasping
            self.move(
                self.move_by_displacement(arm_tag="left", z=0.05, move_axis="arm"),
                self.move_by_displacement(arm_tag="right", z=0.05, move_axis="arm"),
            )

            breadbasket_pose = self.breadbasket.get_functional_point(0)
            # Place first bread into the basket using left arm
            self.move(
                self.place_actor(
                    self.bread[id],
                    arm_tag="left",
                    target_pose=breadbasket_pose,
                    constrain="free",
                    pre_dis=0.13,
                ))
            # Move left arm up a little
            self.move(self.move_by_displacement(arm_tag="left", z=0.1, move_axis="arm"))

            # Move left arm away while placing second bread with right arm, avoiding collision
            self.move(
                self.back_to_origin(arm_tag="left"),
                self.place_actor(
                    self.bread[id ^ 1],
                    arm_tag="right",
                    target_pose=breadbasket_pose,
                    constrain="free",
                    pre_dis=0.13,
                    dis=0.05,  # Move right arm slightly away to avoid collision
                ),
            )

        arm_info = None
        # Check if there's only one bread or both are on the same side
        if (len(self.bread) <= 1 or (self.bread[0].get_pose().p[0] * self.bread[1].get_pose().p[0]) > 0):
            if len(self.bread) == 1:
                # Handle single bread case
                remove_bread(0, 0)
                arm_info = "left" if self.bread[0].get_pose().p[0] < 0 else "right"
            else:
                # When two breads are present but on the same side, pick the front one first
                id = (0 if self.bread[0].get_pose().p[1] < self.bread[1].get_pose().p[1] else 1)
                arm_info = "left" if self.bread[0].get_pose().p[0] < 0 else "right"
                remove_bread(id, 0)
                remove_bread(id ^ 1, 1)
        else:
            # Dual-arm removal when breads are on opposite sides
            remove()
            arm_info = "dual"

        self.info["info"] = {
            "{A}": f"076_breadbasket/base{self.basket_id}",
            "{B}": f"075_bread/base{self.bread_id[0]}",
            "{a}": arm_info,
        }
        if len(self.bread) == 2:
            self.info["info"]["{C}"] = f"075_bread/base{self.bread_id[1]}"

        return self.info

    def check_success(self):
        breadbasket_pose = self.breadbasket.get_pose().p
        eps1 = 0.05
        check = True
        for i in range(len(self.bread)):
            pose = self.bread[i].get_pose().p
            if np.all(abs(pose[:2] - breadbasket_pose[:2]) < np.array([eps1, eps1])) and (pose[2]
                                                                                          > 0.73 + self.table_z_bias):
                continue
            else:
                check = False

        return (check and self.robot.is_left_gripper_open() and self.robot.is_right_gripper_open())
