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


class click_alarmclock(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.0],
            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.05:
            rand_pos = rand_pose(
                xlim=[-0.25, 0.25],
                ylim=[-0.2, 0.0],
                qpos=[0.5, 0.5, 0.5, 0.5],
                rotate_rand=True,
                rotate_lim=[0, 3.14, 0],
            )

        self.alarmclock_id = np.random.choice([1, 3], 1)[0]
        self.alarm = create_actor(
            scene=self,
            pose=rand_pos,
            modelname="046_alarm-clock",
            convex=True,
            model_id=self.alarmclock_id,
            is_static=True,
        )
        self.add_prohibit_area(self.alarm, padding=0.05)
        self.check_arm_function = self.is_left_gripper_close if self.alarm.get_pose().p[0] < 0 else self.is_right_gripper_close

    def play_once(self):
        # Determine which arm to use based on alarm clock's position (right if positive x, left otherwise)
        arm_tag = ArmTag("right" if self.alarm.get_pose().p[0] > 0 else "left")
    
        # Move the gripper above the top center of the alarm clock and close the gripper to simulate a click
        # Note: although the code structure resembles a grasp, it is used here to simulate a touch/click action
        # You can adjust API parameters to move above the top button and close the gripper (similar to grasp_actor)
        self.move((
            ArmTag(arm_tag),
            [
                Action(
                    arm_tag,
                    "move",
                    self.get_grasp_pose(self.alarm, pre_dis=0.1, contact_point_id=0, arm_tag=arm_tag)[:3] +
                    [0.5, -0.5, 0.5, 0.5],
                ),
                Action(arm_tag, "close", target_gripper_pos=0.0),
            ],
        ))
    
        # Move the gripper downward to press the top button of the alarm clock
        self.move(self.move_by_displacement(arm_tag, z=-0.065))
        # Check whether the simulated click action was successful
        self.check_success()
    
        # Move the gripper back to the original height (not lifting the alarm clock)
        self.move(self.move_by_displacement(arm_tag, z=0.065))
        # Optionally check success again
        self.check_success()
    
        # Record information about the alarm clock and the arm used
        self.info["info"] = {
            "{A}": f"046_alarm-clock/base{self.alarmclock_id}",
            "{a}": str(arm_tag),
        }
        return self.info


    def check_success(self):
        if self.stage_success_tag:
            return True
        if not self.check_arm_function():
            return False
        alarm_pose = self.alarm.get_contact_point(0)[:3]
        positions = self.get_gripper_actor_contact_position("046_alarm-clock")
        eps = [0.03, 0.03]
        for position in positions:
            if (np.all(np.abs(position[:2] - alarm_pose[:2]) < eps) and abs(position[2] - alarm_pose[2]) < 0.03):
                self.stage_success_tag = True
                return True
        return False
