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


class place_object_scale(Base_Task):

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

    def load_actors(self):
        rand_pos = rand_pose(
            xlim=[-0.25, 0.25],
            ylim=[-0.2, 0.05],
            qpos=[0.5, 0.5, 0.5, 0.5],
            rotate_rand=True,
            rotate_lim=[0, 3.14, 0],
        )
        while abs(rand_pos.p[0]) < 0.02:
            rand_pos = rand_pose(
                xlim=[-0.25, 0.25],
                ylim=[-0.2, 0.05],
                qpos=[0.5, 0.5, 0.5, 0.5],
                rotate_rand=True,
                rotate_lim=[0, 3.14, 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"]

        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)

        if rand_pos.p[0] > 0:
            xlim = [0.02, 0.25]
        else:
            xlim = [-0.25, -0.02]
        target_rand_pose = rand_pose(
            xlim=xlim,
            ylim=[-0.2, 0.05],
            qpos=[0.5, 0.5, 0.5, 0.5],
            rotate_rand=True,
            rotate_lim=[0, 3.14, 0],
        )
        while (np.sqrt((target_rand_pose.p[0] - rand_pos.p[0])**2 + (target_rand_pose.p[1] - rand_pos.p[1])**2) < 0.15):
            target_rand_pose = rand_pose(
                xlim=xlim,
                ylim=[-0.2, 0.05],
                qpos=[0.5, 0.5, 0.5, 0.5],
                rotate_rand=True,
                rotate_lim=[0, 3.14, 0],
            )

        self.scale_id = np.random.choice([0, 1, 5, 6], 1)[0]

        self.scale = create_actor(
            scene=self,
            pose=target_rand_pose,
            modelname="072_electronicscale",
            model_id=self.scale_id,
            convex=True,
        )
        self.scale.set_mass(0.05)

        self.add_prohibit_area(self.object, padding=0.05)
        self.add_prohibit_area(self.scale, padding=0.05)

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

        # Grasp the object with the selected arm
        self.move(self.grasp_actor(self.object, arm_tag=self.arm_tag))

        # Lift the object up by 0.15 meters in z-axis
        self.move(self.move_by_displacement(arm_tag=self.arm_tag, z=0.15))

        # Place the object on the scale's functional point with free constraint,
        # using pre-placement distance of 0.05m and final placement distance of 0.005m
        self.move(
            self.place_actor(
                self.object,
                arm_tag=self.arm_tag,
                target_pose=self.scale.get_functional_point(0),
                constrain="free",
                pre_dis=0.05,
                dis=0.005,
            ))

        # Record information about the objects and arm used for the task
        self.info["info"] = {
            "{A}": f"072_electronicscale/base{self.scale_id}",
            "{B}": f"{self.selected_modelname}/base{self.selected_model_id}",
            "{a}": str(self.arm_tag),
        }
        return self.info

    def check_success(self):
        object_pose = self.object.get_pose().p
        scale_pose = self.scale.get_functional_point(0)
        distance_threshold = 0.035
        distance = np.linalg.norm(np.array(scale_pose[:2]) - np.array(object_pose[:2]))
        check_arm = (self.is_left_gripper_open if self.arm_tag == "left" else self.is_right_gripper_open)
        return (distance < distance_threshold and object_pose[2] > (scale_pose[2] - 0.01) and check_arm())
