"""Pull a cabinet drawer open (RoboTwin 036_cabinet, three prismatic drawers).

A slider, not a hinge: the drawer front travels in a straight line, so the pull must stay on that
line. The same "follow the joint's own path" idea as microwave_door, with the arc replaced by a
line -- which is exactly why the two share one solver.
"""
from ..registry import register_task
from ..envs import YamTaskEnv
from ..solvers import articulate
from .. import conditions as C


@register_task("drawer_open")
class DrawerOpenTask(YamTaskEnv):
    title = "PULL THE DRAWER OPEN"
    tags = ["articulated", "slider", "constrained-motion", "robotwin-asset"]
    instruction = {"default": "Open the drawer."}

    rw_articulations = {"rw_cabinet": "036_cabinet/46653_clean.usd:0.20"}
    body_xy = (-0.02, -0.24)
    # Derived from the asset, not guessed: after the Y-up roll the drawer front sits at
    # sim y = -0.093 and stands 1.6 cm proud of the cabinet body, and joint 0 drives the
    # BOTTOM drawer, whose top edge is 10.7 cm above the table.
    front_dy = -0.085
    grip_z = 0.100
    press_down = 0.0
    # lean the wrist most of the way to horizontal so the fingers straddle the drawer's lip
    grip_tilt = 62.0
    pull_distance = 0.13
    gripper_effort, gripper_damping = 80, 90

    def _load_scene(self):
        self._placed = []
        self.scene.place_articulation({"name": "rw_cabinet", "bbox": [[-0.4283, -0.8033, -0.4334], [0.426, 0.8112, 0.4655]], "scale": 0.2, "xy": self.body_xy, "half": 0.14,
                                       "xy_jitter": (0.015, 0.012)})
        art = self.scene.articulations.get("rw_cabinet")
        if art is not None:
            print(f"[task] cabinet 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_cabinet', i).round(3)}",
                      flush=True)

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

    def solve(self):
        bx, by = self.scene.regions["rw_cabinet"]["xy"]
        # Take the drawer front from the MEASURED body, not a guessed offset: every link of a
        # closed cabinet reports the body origin, so there is nothing to read off the links, and
        # a hardcoded -9 cm was simply not where this cabinet's face is.
        print(f"[task] drawer lip at dy={self.front_dy:+.3f}, grip_z={self.grip_z:.3f}, "
              f"press {self.press_down*1000:.0f} mm", flush=True)
        return articulate.solve(self, fixture="rw_cabinet", kind="slider",
                                grip_xy=(bx, by+self.front_dy), grip_z=self.grip_z,
                                press=self.press_down, grip_tilt=self.grip_tilt,
                                open_jaw=False, jaw="x",
                                direction=(0.0, -1.0), distance=self.pull_distance,
                                approach_from=(0.0, -1.0))

    def evaluate(self):
        return self.check(
            C.labelled("drawer drawn out", C.joint_driven("rw_cabinet", 0, by=0.05)))