"""Stack two green blocks on a marked target region (RLBench stack_blocks)."""
from ..registry import register_task
from ..envs import YamTaskEnv
from ..solvers import stack
from .. import conditions as C


@register_task("stack_blocks")
class StackBlocksTask(YamTaskEnv):
    title = "STACK THE BLOCKS ON THE TARGET"
    tags = ["stacking", "precision", "primitive"]
    instruction = {"default": "Stack the two green blocks on the target square."}

    block_spawns = {"cube_g": (-0.07, 0.10), "cube_b": (0.07, 0.08)}
    spawn_jitter = (0.020, 0.015)
    target_xy = (0.00, -0.16)
    target_size = 0.12
    # solid 5 cm cubes: a firm clamp, they do not deform
    gripper_effort, gripper_damping = 78, 88

    def _load_scene(self):
        self.scene.build_marker({"name": "target", "xy": self.target_xy, "size": self.target_size,
                                 "xy_jitter": (0.015, 0.015)})
        self._placed = [self.scene.place_object({"name": n, "xy": xy,
                                                 "xy_jitter": self.spawn_jitter})
                        for n, xy in self.block_spawns.items()]

    def solve(self):
        return stack.solve(self, objects=list(self.block_spawns), target="target")

    def evaluate(self):
        a, b = list(self.block_spawns)
        return self.check(
            C.labelled("bottom block on target", C.object_in_region(a, "target", pad=0.04)),
            C.labelled("top block on target", C.object_in_region(b, "target", pad=0.04)),
            C.labelled("blocks stacked", C.objects_stacked(a, b, min_dz=0.03)))