"""Pour from a kettle into a cup: carry it over and rotate the wrist to tip the contents out.

The pour goes into a cup, which is small enough that the wrist rotation has to be aimed rather
than merely performed.
"""
from ..registry import register_task
from ..envs import YamTaskEnv
from ..solvers import pour
from .. import conditions as C


@register_task("pour_kettle")
class PourKettleTask(YamTaskEnv):
    title = "POUR THE KETTLE INTO THE CUP"
    tags = ["pouring", "trajectory", "wrist-rotation", "robotwin-asset"]
    instruction = {"default": "Pour water from the kettle into the cup."}

    rw_objects = {"rw_kettle": "091_kettle/base0.usd:0.052:0.12"}
    # the +4.5 cm grip offset pushes the actual grip point away from the shoulder, so the
    # kettle has to start closer than it looks: at (-0.03,0.11) the grip sat 0.39 m out
    kettle_xy = (0.0, 0.03)
    # pour INTO A CUP. The _mesh variant is required: the convex-decomposition build fills the
    # cavity solid, so the beads would land on top of the cup instead of dropping into it.
    cup = {"usd": "021_cup/base0_mesh.usd", "scale": 0.085, "rpy": (90, 0, 0), "xy": (0.02, -0.22)}
    # No simulated contents, deliberately. A vessel that gets LIFTED has to be a convex dynamic
    # body -- PhysX will not carry a hollow triangle mesh -- so the kettle's interior is solid and
    # a bead placed "inside" is pushed out onto the table before the episode even starts (they
    # were measured sitting at table height every run). Success is therefore the POUR itself:
    # kettle over the cup, tipped past the angle its contents would leave it at.
    pour_deg = 108.0
    min_pour_tilt = 75.0
    # Measured off the mesh, not guessed. Two separate problems, both invisible in the numbers:
    #   * the mesh origin is at its EDGE again -- the body's centroid is 4.5 cm off in y, so the
    #     jaws were closing beside the kettle ("did not rise")
    #   * the body is a smooth 7 cm shell that slips; the top 25% is a 1.5 x 3.5 cm neck, which
    #     is a real pinch
    grip_offset = (0.0, 0.045)
    # the 1.5 cm neck stalls the jaws but slips out under load; grip the 7.2 cm body
    # instead, which the corrected y-offset now actually lands on
    grip_top = None
    gripper_effort, gripper_damping = 78, 86

    def _load_scene(self):
        self.scene.spawn_container({"name": "cup", **self.cup, "xy_jitter": (0.015, 0.015)})
        self._placed = [self.scene.place_object(
            {"name": "rw_kettle", "xy": self.kettle_xy, "xy_jitter": (0.015, 0.010)})]

    def solve(self):
        pos = self.scene.object_pos("rw_kettle")
        print(f"[task] kettle {self.scene.object_size('rw_kettle').round(3)} at {pos.round(3)}",
              flush=True)
        return pour.solve(self, obj="rw_kettle", target="cup", pour_deg=self.pour_deg,
                          max_gap=0.048, jaw="x", grasp_offset=self.grip_offset,
                          grasp_top=self.grip_top)

    def evaluate(self):
        return self.check(
            C.labelled("kettle lifted", C.object_lifted("rw_kettle", 0.05)),
            C.labelled("kettle tipped out over the cup",
                       C.object_tipped_over("rw_kettle", "cup", min_tilt_deg=self.min_pour_tilt)))