from ._base_task import Base_Task
from .utils import *
import sapien
import transforms3d as t3d
from ._GLOBAL_CONFIGS import *


class beat_block_hammer(Base_Task):

    def setup_demo(self, **kwags):
        self.grasp_pre_dis = kwags.get("grasp_pre_dis", 0.12)
        self.grasp_dis = kwags.get("grasp_dis", 0.01)
        self.grasp_close_pos = kwags.get("grasp_close_pos", 0.0)
        self.grasp_pose_offset = np.asarray(kwags.get("grasp_pose_offset", [0.0, 0.0, 0.0]), dtype=np.float64)
        self.grasp_approach_down_angle_deg = float(kwags.get("grasp_approach_down_angle_deg", 30.0))
        self.grasp_tcp_to_contact_dis = float(kwags.get("grasp_tcp_to_contact_dis", 0.12))
        self.post_grasp_contact_steps = kwags.get("post_grasp_contact_steps", 4)
        self.place_pre_dis = kwags.get("place_pre_dis", 0.06)
        self.place_raise_dis = kwags.get("place_raise_dis", 0.06)
        self.place_touch_z_offset = kwags.get("place_touch_z_offset", 0.004)
        self.place_xy_correction_threshold = kwags.get("place_xy_correction_threshold", 0.02)
        self.place_final_down_iters = kwags.get("place_final_down_iters", 5)
        self.force_arm_tag = kwags.get("force_arm_tag", None)
        self.force_block_arm_tag = kwags.get("force_block_arm_tag", self.force_arm_tag)
        self.physical_grasp = kwags.get("physical_grasp", False)
        self.hammer_mass = float(kwags.get("hammer_mass", 0.001))
        super()._init_task_env_(**kwags)

    def _step_with_robot_drives(self):
        self.robot._entity_qf(self.robot.left_entity)
        if self.robot.right_entity is not self.robot.left_entity:
            self.robot._entity_qf(self.robot.right_entity)
        self.scene.step()

    def _set_actor_visual_color(self, actor, color):
        for component in actor.actor.get_components():
            if not isinstance(component, sapien.render.RenderBodyComponent):
                continue
            for shape in component.render_shapes:
                material = shape.get_material()
                if hasattr(material, "set_base_color_texture"):
                    material.set_base_color_texture(None)
                material.set_base_color([*color[:3], 1])

    def _pose_for_actor_point_correction(self, arm_tag, actor_point, target_point):
        ee_pose = np.array(
            self.robot.get_left_ee_pose() if arm_tag == "left" else self.robot.get_right_ee_pose(),
            dtype=np.float64,
        )
        correction = np.asarray(target_point, dtype=np.float64) - np.asarray(actor_point, dtype=np.float64)
        target_pose = ee_pose.copy()
        target_pose[:3] += correction
        return target_pose

    def _apply_grasp_pose_offset(self, pre_grasp_pose, grasp_pose, arm_tag):
        if np.linalg.norm(self.grasp_pose_offset) == 0:
            return pre_grasp_pose, grasp_pose
        if pre_grasp_pose is None or grasp_pose is None:
            self.plan_success = False
            return pre_grasp_pose, grasp_pose
        grasp_pose_offset = self.grasp_pose_offset.copy()
        pre_grasp_pose = np.asarray(pre_grasp_pose, dtype=np.float64).copy()
        grasp_pose = np.asarray(grasp_pose, dtype=np.float64).copy()
        if pre_grasp_pose.ndim != 1 or grasp_pose.ndim != 1 or pre_grasp_pose.size < 3 or grasp_pose.size < 3:
            self.plan_success = False
            return None, None
        pre_grasp_pose[:3] += grasp_pose_offset
        grasp_pose[:3] += grasp_pose_offset
        return pre_grasp_pose.tolist(), grasp_pose.tolist()

    def _tilt_grasp_approach(self, pre_grasp_pose, grasp_pose):
        if pre_grasp_pose is None or grasp_pose is None:
            self.plan_success = False
            return pre_grasp_pose, grasp_pose
        pre_grasp_pose = np.asarray(pre_grasp_pose, dtype=np.float64).copy()
        grasp_pose = np.asarray(grasp_pose, dtype=np.float64).copy()
        approach_distance = np.linalg.norm(grasp_pose[:3] - pre_grasp_pose[:3])
        if approach_distance < 1e-6:
            return pre_grasp_pose.tolist(), grasp_pose.tolist()

        angle = np.deg2rad(self.grasp_approach_down_angle_deg)
        desired_dir = np.array([0.0, np.cos(angle), -np.sin(angle)], dtype=np.float64)
        desired_dir /= np.linalg.norm(desired_dir)

        current_dir = grasp_pose[:3] - pre_grasp_pose[:3]
        current_dir /= np.linalg.norm(current_dir)

        cross = np.cross(current_dir, desired_dir)
        dot = float(np.clip(np.dot(current_dir, desired_dir), -1.0, 1.0))
        if np.linalg.norm(cross) < 1e-8:
            if dot > 0:
                align_mat = np.eye(3)
            else:
                axis = np.cross(current_dir, [1.0, 0.0, 0.0])
                if np.linalg.norm(axis) < 1e-8:
                    axis = np.cross(current_dir, [0.0, 1.0, 0.0])
                axis /= np.linalg.norm(axis)
                align_mat = t3d.axangles.axangle2mat(axis, np.pi)
        else:
            cross_mat = np.array(
                [
                    [0.0, -cross[2], cross[1]],
                    [cross[2], 0.0, -cross[0]],
                    [-cross[1], cross[0], 0.0],
                ],
                dtype=np.float64,
            )
            align_mat = np.eye(3) + cross_mat + cross_mat @ cross_mat * ((1.0 - dot) / (np.linalg.norm(cross) ** 2))

        grasp_rot = t3d.quaternions.quat2mat(grasp_pose[-4:])
        tilted_rot = align_mat @ grasp_rot
        grasp_pose[-4:] = t3d.quaternions.mat2quat(tilted_rot)
        contact_point = grasp_pose[:3] + current_dir * (self.grasp_tcp_to_contact_dis + self.grasp_dis)
        pre_grasp_pose[-4:] = grasp_pose[-4:]
        grasp_pose[:3] = contact_point - desired_dir * (self.grasp_tcp_to_contact_dis + self.grasp_dis)
        pre_grasp_pose[:3] = contact_point - desired_dir * (self.grasp_tcp_to_contact_dis + self.grasp_pre_dis)
        return pre_grasp_pose.tolist(), grasp_pose.tolist()

    def _close_gripper_holding_arm(self, arm_tag, target_gripper_pos):
        if arm_tag == "left":
            gripper_result = self.set_gripper(left_pos=target_gripper_pos, set_tag="left")
        else:
            gripper_result = self.set_gripper(right_pos=target_gripper_pos, set_tag="right")
        control_seq = {
            "left_arm": None,
            "left_gripper": gripper_result if arm_tag == "left" else None,
            "right_arm": None,
            "right_gripper": gripper_result if arm_tag == "right" else None,
        }
        self.take_dense_action(control_seq)
        self.robot.hold_passive_gripper_mimic(arm_tag)
        return True

    def load_actors(self):
        table_center_x, table_center_y = self.table_xy_bias
        table_near_y = table_center_y - self.table_width / 2
        hammer_xy = [table_center_x, table_center_y]
        block_x_offset = [0.05, 0.25]
        block_y_from_near_edge = [0.15, 0.45]

        self.hammer = create_actor(
            scene=self,
            pose=sapien.Pose([hammer_xy[0], hammer_xy[1], 0.783], [0, 0, 0.995, 0.105]),
            modelname="020_hammer",
            convex=True,
            model_id=0,
        )
        self._set_actor_visual_color(self.hammer, (142 / 255, 144 / 255, 137 / 255))
        if self.force_block_arm_tag == "right":
            block_xlim = [hammer_xy[0] + block_x_offset[0], hammer_xy[0] + block_x_offset[1]]
        elif self.force_block_arm_tag == "left":
            block_xlim = [hammer_xy[0] - block_x_offset[1], hammer_xy[0] - block_x_offset[0]]
        else:
            block_xlim = [hammer_xy[0] + block_x_offset[0], hammer_xy[0] + block_x_offset[1]]
        block_ylim = [
            table_near_y + block_y_from_near_edge[0],
            table_near_y + block_y_from_near_edge[1],
        ]
        block_pose = rand_pose(
            xlim=block_xlim,
            ylim=block_ylim,
            zlim=[0.76],
            qpos=[1, 0, 0, 0],
            rotate_rand=True,
            rotate_lim=[0, 0, 0.5],
        )
        while (self.force_block_arm_tag is None and abs(block_pose.p[0]) < 0.05) or np.sum(pow(block_pose.p[:2], 2)) < 0.001:
            block_pose = rand_pose(
                xlim=block_xlim,
                ylim=block_ylim,
                zlim=[0.76],
                qpos=[1, 0, 0, 0],
                rotate_rand=True,
                rotate_lim=[0, 0, 0.5],
            )

        self.block = create_box(
            scene=self,
            pose=block_pose,
            half_size=(0.025, 0.025, 0.025),
            color=(63 / 255, 142 / 255, 67 / 255),
            name="box",
            is_static=True,
        )
        self.hammer.set_mass(self.hammer_mass)

        self.add_prohibit_area(self.hammer, padding=0.10)
        self.prohibited_area.append([
            block_pose.p[0] - 0.05,
            block_pose.p[1] - 0.05,
            block_pose.p[0] + 0.05,
            block_pose.p[1] + 0.05,
        ])

    def play_once(self):
        # Get the position of the block's functional point
        block_pose = self.block.get_functional_point(0, "pose").p
        # Determine which arm to use based on block position (left if block is on left side, else right)
        arm_tag = ArmTag(self.force_arm_tag if self.force_arm_tag is not None else ("left" if block_pose[0] < 0 else "right"))

        # Grasp the hammer with the selected arm
        if self.physical_grasp:
            pre_grasp_pose, grasp_pose = self.choose_grasp_pose(
                self.hammer,
                arm_tag=arm_tag,
                pre_dis=self.grasp_pre_dis,
                target_dis=self.grasp_dis,
            )
            pre_grasp_pose, grasp_pose = self._apply_grasp_pose_offset(pre_grasp_pose, grasp_pose, arm_tag)
            pre_grasp_pose, grasp_pose = self._tilt_grasp_approach(pre_grasp_pose, grasp_pose)
            if not self.plan_success:
                return self.info
            self.move((arm_tag, [Action(arm_tag, "move", target_pose=pre_grasp_pose)]))
            self.move(
                (
                    arm_tag,
                    [
                        Action(
                            arm_tag,
                            "move",
                            target_pose=grasp_pose,
                            constraint_pose=[1, 1, 1, 0, 0, 0],
                        ),
                    ],
                )
            )
            self._close_gripper_holding_arm(arm_tag, self.grasp_close_pos)
            for _ in range(self.post_grasp_contact_steps):
                self._step_with_robot_drives()
        else:
            self.move(
                self.grasp_actor(
                    self.hammer,
                    arm_tag=arm_tag,
                    pre_grasp_dis=self.grasp_pre_dis,
                    grasp_dis=self.grasp_dis,
                    gripper_pos=self.grasp_close_pos,
                ))
        # Move the hammer upwards
        if self.physical_grasp:
            self.move(self.move_by_displacement(arm_tag, z=0.07, move_axis="arm"))
        else:
            self.move(self.move_by_displacement(arm_tag, z=0.07, move_axis="arm"))

        # Place the hammer on the block's functional point (position 1)
        if self.physical_grasp:
            target_point = self.block.get_functional_point(1, "pose").p

            ee_pose = np.array(
                self.robot.get_left_ee_pose() if arm_tag == "left" else self.robot.get_right_ee_pose(),
                dtype=np.float64,
            )
            raise_pose = ee_pose.copy()
            raise_pose[2] += self.place_raise_dis
            if np.linalg.norm(raise_pose[:3] - ee_pose[:3]) > 1e-6:
                self.move((arm_tag, [Action(arm_tag, "move", target_pose=raise_pose)]))

            high_target_point = target_point.copy()
            high_target_point[2] += self.place_touch_z_offset + self.place_pre_dis
            hammer_point = self.hammer.get_functional_point(0, "pose").p
            high_place_pose = self._pose_for_actor_point_correction(
                arm_tag,
                hammer_point,
                high_target_point,
            )
            self.move((arm_tag, [Action(arm_tag, "move", target_pose=high_place_pose)]))

            for _ in range(self.place_final_down_iters):
                if self.check_success():
                    break
                ee_pose = np.array(
                    self.robot.get_left_ee_pose() if arm_tag == "left" else self.robot.get_right_ee_pose(),
                    dtype=np.float64,
                )
                hammer_point = self.hammer.get_functional_point(0, "pose").p
                xy_error = target_point[:2] - hammer_point[:2]
                if not np.all(np.abs(xy_error) <= self.place_xy_correction_threshold):
                    xy_place_pose = ee_pose.copy()
                    xy_place_pose[:2] += xy_error
                    self.move((arm_tag, [Action(arm_tag, "move", target_pose=xy_place_pose)]))
                    continue
                if self.check_actors_contact(self.hammer.get_name(), self.block.get_name()):
                    break
                z_error = target_point[2] + self.place_touch_z_offset - hammer_point[2]
                if abs(z_error) < 0.005:
                    break
                down_target_point = hammer_point.copy()
                down_target_point[2] = target_point[2] + self.place_touch_z_offset
                down_pose = self._pose_for_actor_point_correction(
                    arm_tag,
                    hammer_point,
                    down_target_point,
                )
                self.move((arm_tag, [Action(arm_tag, "move", target_pose=down_pose)]))
        else:
            self.move(
                self.place_actor(
                    self.hammer,
                    target_pose=self.block.get_functional_point(1, "pose"),
                    arm_tag=arm_tag,
                    functional_point_id=0,
                    pre_dis=self.place_pre_dis,
                    dis=0,
                    is_open=False,
                ))

        self.info["info"] = {"{A}": "020_hammer/base0", "{a}": str(arm_tag)}
        return self.info

    def _hammer_contacts_block(self):
        return self.check_actors_contact(self.hammer.get_name(), self.block.get_name())

    def check_strict_success(self):
        hammer_target_pose = self.hammer.get_functional_point(0, "pose").p
        block_pose = self.block.get_functional_point(1, "pose").p
        eps = np.array([0.02, 0.02])
        return np.all(abs(hammer_target_pose[:2] - block_pose[:2]) < eps) and self._hammer_contacts_block()

    def check_success(self):
        return self.check_strict_success()
