"""Flip a wall switch (RoboTwin 056_switch, hinged toggle button).

The smallest articulated task in the suite and the one with the least margin: the toggle is a
couple of centimetres across, nothing is grasped, and the whole motion is a short press with the
closed gripper. Success is the joint moving, not anything moving in space.
"""
from ..registry import register_task
from ..envs import YamTaskEnv
from ..solvers import articulate
from .. import conditions as C


@register_task("switch_toggle")
class SwitchToggleTask(YamTaskEnv):
    title = "FLIP THE SWITCH"
    tags = ["articulated", "hinge", "precision", "no-grasp", "robotwin-asset"]
    instruction = {"default": "Flip the switch."}

    rw_articulations = {"rw_switch": "056_switch/100880_clean.usd:0.13"}
    body_xy = (0.00, -0.14)
    press_dy = 0.0
    press_depth = 0.045             # how far to drive the fingertip down onto the button
    # The button link's frame IS the hinge, so pressing exactly there applies zero moment and the
    # joint does not budge (measured: 0.000 rad). Press off to one side of the hinge line.
    press_lever = 0.016
    gripper_effort, gripper_damping = 80, 90

    def _load_scene(self):
        self._placed = []
        self.scene.place_articulation({"name": "rw_switch", "bbox": [[-0.5435, -0.8475, -0.9522], [0.5292, 0.7599, 0.7389]], "scale": 0.13, "xy": self.body_xy, "half": 0.06,
                                       "xy_jitter": (0.012, 0.010)})
        art = self.scene.articulations.get("rw_switch")
        if art is not None:
            print(f"[task] switch joints={list(art.data.joint_names)} "
                  f"limits={art.data.joint_pos_limits[0].cpu().numpy().round(3).tolist()}", flush=True)
            for i, nm in enumerate(art.data.body_names):
                print(f"[task]   link {i} {nm}: {self.scene.link_pos('rw_switch', i).round(3)}",
                      flush=True)

    def _initialize_episode(self):
        for _ in range(60):
            self.step()

    def solve(self):
        # Press the ACTUAL button, read off its link. The button sits 6.3 cm off the body centre
        # in x -- pressing the body centre pressed the switch plate and nothing moved. Its hinge
        # axis is horizontal, so the motion is a straight downward press, not an arc.
        b = self.scene.link_pos("rw_switch", 1)
        print(f"[task] toggle button at {b.round(3)} ({b[2]-0.45:.3f} m above the table)", flush=True)
        return articulate.solve(self, fixture="rw_switch", kind="press",
                                grip_xy=(float(b[0]), float(b[1])+self.press_lever),
                                grip_z=float(b[2]-0.45)+0.012,
                                distance=self.press_depth,
                                approach_from=(0.0, -1.0), open_jaw=False)

    def evaluate(self):
        return self.check(
            C.labelled("switch toggled", C.joint_driven("rw_switch", 0, by=0.25)))