"""ManiSkill-style peg insertion: put the peg into a square socket."""
from ..registry import register_task
from ..envs import YamTaskEnv
from ..solvers import insert
from .. import conditions as C


@register_task("peg_insert")
class PegInsertTask(YamTaskEnv):
    title = "PEG -> HOLE (insertion)"
    tags = ["insertion", "precision", "primitive"]
    instruction = {"default": "Put the peg in the hole.",
                   "specific": "Pick up the yellow peg and insert it into the square socket."}

    peg_spawn = (-0.03, 0.10)
    peg_spawn_jitter = (0.02, 0.015)
    socket_xy = (-0.02, -0.22)
    socket_jitter = (0.015, 0.015)
    hole_width = 0.042
    # grasp near the TOP so ~10 cm of shaft is left below the fingers to enter the hole
    grasp_top = 0.03
    gripper_effort, gripper_damping = 75, 85

    def _load_scene(self):
        self.scene.build_socket({"name": "socket", "xy": self.socket_xy,
                                 "xy_jitter": self.socket_jitter, "hole": self.hole_width})
        self._placed = [self.scene.place_object(
            {"name": "peg", "xy": self.peg_spawn, "xy_jitter": self.peg_spawn_jitter})]

    def solve(self):
        return insert.solve(self, obj="peg", target="socket", grasp_top=self.grasp_top)

    def evaluate(self):
        return self.check(C.labelled(
            "peg seated in hole", C.object_seated_in_hole("peg", "socket", xy_tol=0.02, half_h=0.065)))