"""Wipe a stain off the table with an eraser held in the gripper (RLBench wipe_desk).

The reward here is the *path*, not a final object pose: the eraser has to stay pressed on the
table and cover the stain in lawn-mower lanes, which is a long continuous trajectory rather than
a pick-and-place.
"""
from ..registry import register_task
from ..envs import YamTaskEnv
from ..solvers import tool_use
from .. import conditions as C


@register_task("eraser_wipe")
class EraserWipeTask(YamTaskEnv):
    title = "WIPE THE STAIN WITH THE ERASER"
    tags = ["tool-use", "trajectory", "contact-rich", "robotwin-asset"]
    instruction = {"default": "Wipe the stain off the table."}

    # 0.09 gave a 17 cm long block -- too long to sweep inside a 13 cm stain, and the jaws had to
    # clamp its 6.4 cm waist while 5 cm of overhang levered it out
    rw_objects = {"rw_eraser": "117_whiteboard-eraser/base0.usd:0.055:0.05"}
    eraser_xy = (-0.05, 0.13)
    # roll 90 lays the eraser PAD-DOWN. The mesh's thin axis is Y and its flat pad is the low-Y
    # face, so left unrotated the eraser stands on its long edge and wipes with its side.
    eraser_roll = 90
    board_xy = (0.03, -0.13)
    board_w, board_d = 0.30, 0.22
    # the smudge to be wiped, drawn ON the board
    stain_size = 0.15
    stain_color = (0.46, 0.46, 0.48)
    # a soft block: clamping hard squashes it out of the jaws
    gripper_effort, gripper_damping = 52, 70

    def _load_scene(self):
        b = self.scene.build_whiteboard({"name": "board", "xy": self.board_xy,
                                         "width": self.board_w, "depth": self.board_d,
                                         "xy_jitter": (0.012, 0.012)})
        # the grey smudge sits ON the board's surface, so the wipe happens on the board
        self.scene.build_marker({"name": "stain", "xy": b["xy"], "size": self.stain_size,
                                 "aspect": 0.72, "color": self.stain_color,
                                 "z": b["top_z"]+0.001})
        self._placed = [self.scene.place_object(
            {"name": "rw_eraser", "xy": self.eraser_xy, "roll": self.eraser_roll,
             "xy_jitter": (0.015, 0.012)})]

    def solve(self):
        return tool_use.solve(self, tool="rw_eraser", target="stain", kind="wipe")

    def evaluate(self):
        # the eraser must have been carried over the stain and set down there, not dropped en route
        return self.check(
            C.labelled("eraser ended over the stain", C.object_moved_to("rw_eraser", "stain", tol=0.11)),
            C.labelled("eraser stayed flat on the table", C.object_level("rw_eraser", 30.0)))