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


class place_object_stand(Base_Task):

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

    def load_actors(self):
        rand_pos = rand_pose(
            xlim=[-0.28, 0.28],
            ylim=[-0.05, 0.05],
            qpos=[0.707, 0.707, 0.0, 0.0],
            rotate_rand=True,
            rotate_lim=[0, np.pi / 3, 0],
        )
        while abs(rand_pos.p[0]) < 0.2:
            rand_pos = rand_pose(
                xlim=[-0.28, 0.28],
                ylim=[-0.05, 0.05],
                qpos=[0.707, 0.707, 0.0, 0.0],
                rotate_rand=True,
                rotate_lim=[0, np.pi / 3, 0],
            )

        def get_available_model_ids(modelname):
            asset_path = os.path.join("assets/objects", modelname)
            json_files = glob.glob(os.path.join(asset_path, "model_data*.json"))

            available_ids = []
            for file in json_files:
                base = os.path.basename(file)
                try:
                    idx = int(base.replace("model_data", "").replace(".json", ""))
                    available_ids.append(idx)
                except ValueError:
                    continue

            return available_ids

        object_list = [
            "047_mouse",
            "048_stapler",
            "050_bell",
            "073_rubikscube",
            "057_toycar",
            "079_remotecontrol",
        ]
        self.selected_modelname = np.random.choice(object_list)
        available_model_ids = get_available_model_ids(self.selected_modelname)
        if not available_model_ids:
            raise ValueError(f"No available model_data.json files found for {self.selected_modelname}")
        self.selected_model_id = np.random.choice(available_model_ids)
        self.object = create_actor(
            scene=self,
            pose=rand_pos,
            modelname=self.selected_modelname,
            convex=True,
            model_id=self.selected_model_id,
        )
        self.object.set_mass(0.05)

        object_pos = self.object.get_pose()
        if object_pos.p[0] > 0:
            xlim = [0.0, 0.05]
        else:
            xlim = [-0.05, 0.0]
        target_rand_pos = rand_pose(
            xlim=xlim,
            ylim=[-0.15, -0.1],
            qpos=[0.707, 0.707, 0.0, 0.0],
            rotate_rand=True,
            rotate_lim=[0, np.pi / 6, 0],
        )
        while ((object_pos.p[0] - target_rand_pos.p[0])**2 + (object_pos.p[1] - target_rand_pos.p[1])**2) < 0.01:
            target_rand_pos = rand_pose(
                xlim=xlim,
                ylim=[-0.15, -0.1],
                qpos=[0.707, 0.707, 0.0, 0.0],
                rotate_rand=True,
                rotate_lim=[0, np.pi / 6, 0],
            )
        id_list = [0, 1, 2, 3, 4]
        self.displaystand_id = np.random.choice(id_list)
        self.displaystand = create_actor(
            scene=self,
            pose=target_rand_pos,
            modelname="074_displaystand",
            convex=True,
            model_id=self.displaystand_id,
        )

        self.object.set_mass(0.01)
        self.displaystand.set_mass(0.01)

        self.add_prohibit_area(self.displaystand, padding=0.05)
        self.add_prohibit_area(self.object, padding=0.1)

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

        # Grasp the object with specified arm
        self.move(self.grasp_actor(self.object, arm_tag=arm_tag, pre_grasp_dis=0.1))
        # Lift the object up by 0.06 meters in z-direction
        self.move(self.move_by_displacement(arm_tag=arm_tag, z=0.06))

        # Get the target pose from display stand's functional point
        displaystand_pose = self.displaystand.get_functional_point(0)

        # Place the object onto the display stand with free constraint
        self.move(
            self.place_actor(
                self.object,
                arm_tag=arm_tag,
                target_pose=displaystand_pose,
                constrain="free",
                pre_dis=0.07,
            ))

        # Store information about the objects and arm used in the info dictionary
        self.info["info"] = {
            "{A}": f"{self.selected_modelname}/base{self.selected_model_id}",
            "{B}": f"074_displaystand/base{self.displaystand_id}",
            "{a}": str(arm_tag),
        }
        return self.info

    def check_success(self):
        object_pose = self.object.get_pose().p
        displaystand_pose = self.displaystand.get_pose().p
        eps1 = 0.03
        return (np.all(abs(object_pose[:2] - displaystand_pose[:2]) < np.array([eps1, eps1]))
                and self.robot.is_left_gripper_open() and self.robot.is_right_gripper_open())
