"""Two-stage articulated task: pull the drawer open, then put a milk carton inside it.

Neither half is interesting alone -- the point is the ORDER. The drawer has to be open before
anything can go in, so the second skill's workspace is created by the first one, and a policy
that gets the order wrong drives the carton into a closed drawer front.
"""
from ..registry import register_task
from ..envs import YamTaskEnv
from ..solvers import articulate, shelf
from .. import conditions as C


@register_task("drawer_store")
class DrawerStoreTask(YamTaskEnv):
    title = "OPEN THE DRAWER, PUT THE CARTON IN"
    tags = ["articulated", "slider", "sequential", "pick-place", "robotwin-asset"]
    instruction = {"default": "Open the drawer and put the milk carton inside."}

    rw_articulations = {"rw_cabinet": "036_cabinet/46653_clean.usd:0.20"}
    rw_objects = {"rw_carton": "038_milk-box/base0.usd:0.045:0.08"}
    body_xy = (-0.02, -0.26)
    carton_xy = (0.06, 0.06)
    spawn_jitter = (0.020, 0.015)
    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.14
    gripper_effort, gripper_damping = 80, 90

    def _load_scene(self):
        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.13,
                                       "xy_jitter": (0.015, 0.012)})
        self._placed = [self.scene.place_object(
            {"name": "rw_carton", "xy": self.carton_xy, "xy_jitter": self.spawn_jitter})]
        art = self.scene.articulations.get("rw_cabinet")
        if art is not None:
            print(f"[task] cabinet joints={list(art.data.joint_names)}", 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 solve(self):
        bx, by = self.scene.regions["rw_cabinet"]["xy"]
        opened = 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))
        if opened["moved"] < 0.04:
            self.recorder.phase = "ABORT: the drawer never opened, so nothing can go in"
            return {"opened": False, **opened}
        # the drawer's mouth is now `moved` metres out in front of the cabinet body
        self.scene.regions["drawer_mouth"] = {
            "xy": (bx, by+self.front_dy-opened["moved"]*0.45),
            "half": 0.05, "shelf_z": self.scene.regions["rw_cabinet"]["top_z"]+self.grip_z-0.02}
        put = shelf.solve(self, obj="rw_carton", target="drawer_mouth", standoff=0.14)
        return {"opened": True, **opened, **put}

    def evaluate(self):
        return self.check(
            C.labelled("drawer opened", C.joint_driven("rw_cabinet", 0, by=0.05)),
            C.labelled("carton went into the drawer",
                       C.object_moved_to("rw_carton", "drawer_mouth", tol=0.09)))