"""Pour: carry a cup of beads over a bowl and rotate the wrist to tip them out."""
from ..registry import register_task
from ..envs import YamTaskEnv
from ..solvers import pour
from .. import conditions as C


@register_task("pour_cup")
class PourCupTask(YamTaskEnv):
    title = "POUR THE CUP INTO THE BOWL"
    tags = ["pouring", "trajectory", "wrist-rotation", "robotwin-asset"]
    instruction = {"default": "Pour the cup out into the bowl."}

    # 0.05 gave a 4.7 cm cup: too narrow for even two 1.8 cm beads, never mind four
    rw_objects = {"rw_cup": "021_cup/base0.usd:0.062:0.10"}
    cup_xy = (-0.03, 0.10)
    bowl = {"usd": "002_bowl/base1_mesh.usd", "scale": 0.062, "rpy": (90, 0, 0), "xy": (0.02, -0.24)}
    n_beads = 2
    bead_drop_z = 0.030          # start the beads INSIDE the cup, not around its foot
    pour_deg = 105.0
    gripper_effort, gripper_damping = 48, 66

    def _load_scene(self):
        self.scene.spawn_container({"name": "bowl", **self.bowl, "xy_jitter": (0.015, 0.015)})
        self._placed = [self.scene.place_object(
            {"name": "rw_cup", "xy": self.cup_xy, "xy_jitter": (0.015, 0.010)})]
        for i in range(self.n_beads):
            # z_offset lifts each bead into the cup's mouth. Without it reseat_objects seats them
            # at table height -- i.e. on the table AROUND the cup -- and the episode mimes a pour
            # with an empty cup (it reported 0/4 poured while looking correct).
            self._placed.append(self.scene.place_object(
                {"name": f"bead{i}", "xy": (self.cup_xy[0]+0.011*(i-1), self.cup_xy[1]),
                 "z_offset": self.bead_drop_z, "xy_jitter": (0.002, 0.002)}))

    def solve(self):
        return pour.solve(self, obj="rw_cup", target="bowl", pour_deg=self.pour_deg,
                          beads=[f"bead{i}" for i in range(self.n_beads)], max_gap=0.035)

    def evaluate(self):
        return self.check(C.labelled(
            "beads poured into the bowl",
            C.count_in_region([f"bead{i}" for i in range(self.n_beads)], "bowl",
                              at_least=max(1, self.n_beads//2), pad=0.05)))