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


class place_fan(Base_Task):

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

    def load_actors(self):
        rand_pos = rand_pose(
            xlim=[-0.1, 0.1],
            ylim=[-0.15, -0.05],
            qpos=[0.0, 0.0, 0.707, 0.707],
            rotate_rand=True,
            rotate_lim=[0, 2 * np.pi, 0],
        )
        id_list = [4, 5]
        self.fan_id = np.random.choice(id_list)
        self.fan = create_actor(
            scene=self,
            pose=rand_pos,
            modelname="099_fan",
            convex=True,
            model_id=self.fan_id,
        )
        self.fan.set_mass(0.01)

        xlim = [0.15, 0.25] if self.fan.get_pose().p[0] > 0 else [-0.25, -0.15]
        rand_pos = rand_pose(
            xlim=xlim,
            ylim=[-0.15, -0.05],
        )

        colors = {
            "Red": (1, 0, 0),
            "Green": (0, 1, 0),
            "Blue": (0, 0, 1),
            "Yellow": (1, 1, 0),
            "Cyan": (0, 1, 1),
            "Magenta": (1, 0, 1),
            "Black": (0, 0, 0),
            "Gray": (0.5, 0.5, 0.5),
            "Orange": (1, 0.5, 0),
            "Purple": (0.5, 0, 0.5),
            "Brown": (0.65, 0.4, 0.16),
            "Pink": (1, 0.75, 0.8),
            "Lime": (0.5, 1, 0),
            "Olive": (0.5, 0.5, 0),
            "Teal": (0, 0.5, 0.5),
            "Maroon": (0.5, 0, 0),
            "Navy": (0, 0, 0.5),
            "Coral": (1, 0.5, 0.31),
            "Turquoise": (0.25, 0.88, 0.82),
            "Indigo": (0.29, 0, 0.51),
            "Beige": (0.96, 0.91, 0.81),
            "Tan": (0.82, 0.71, 0.55),
            "Silver": (0.75, 0.75, 0.75),
        }

        color_items = list(colors.items())
        idx = np.random.choice(len(color_items))
        self.color_name, self.color_value = color_items[idx]

        self.pad = create_box(
            scene=self.scene,
            pose=rand_pos,
            half_size=(0.05, 0.05, 0.001),
            color=self.color_value,
            name="box",
        )

        self.pad.set_mass(1)
        self.add_prohibit_area(self.fan, padding=0.07)
        self.prohibited_area.append([
            rand_pos.p[0] - 0.15,
            rand_pos.p[1] - 0.15,
            rand_pos.p[0] + 0.15,
            rand_pos.p[1] + 0.15,
        ])
        # Get the target pose for placing the fan from the pad's current pose
        target_pose = self.pad.get_pose().p
        self.target_pose = target_pose.tolist() + [1, 0, 0, 0]

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

        # Grasp the fan with the selected arm
        self.move(self.grasp_actor(self.fan, arm_tag=arm_tag, pre_grasp_dis=0.05))
        # Lift the fan slightly after grasping
        self.move(self.move_by_displacement(arm_tag=arm_tag, z=0.05))

        # Place the fan onto the pad with alignment constraint along specified axes
        self.move(
            self.place_actor(
                self.fan,
                arm_tag=arm_tag,
                target_pose=self.target_pose,
                constrain="align",
                pre_dis=0.04,
                dis=0.005,
            ))

        self.info["info"] = {
            "{A}": f"099_fan/base{self.fan_id}",
            "{B}": self.color_name,
            "{a}": str(arm_tag),
        }
        return self.info

    def check_success(self):
        fan_qpose = self.fan.get_pose().q
        fan_pose = self.fan.get_pose().p

        target_pose = self.target_pose[:3]
        target_qpose = np.array([0.707, 0.707, 0.0, 0.0])

        if fan_qpose[0] < 0:
            fan_qpose *= -1

        eps = np.array([0.05, 0.05, 0.05, 0.05])

        return (np.all(abs(fan_qpose - target_qpose) < eps[-4:]) and self.robot.is_left_gripper_open()
                and self.robot.is_right_gripper_open()) and (np.all(abs(fan_pose - target_pose) < np.array([0.04, 0.04, 0.04])))
