"""
Language → pick/place resolution for the demo2 arm.

Maps a natural-language instruction to a concrete (pick_body, place_xyz) for the
goal-conditioned pick-place policy. The policy is already conditioned on pick &
place coordinates, so "language" only needs to resolve which object (by colour)
and where to put it (the pad, by default). Keyword-based; no model required.
"""
from __future__ import annotations

import mujoco

# colour word -> object body name in build_arm_scene
COLOR_BODY = {
    "green": "object", "red": "object_red", "blue": "object_blue",
}
_ALIASES = {"the one": "green", "it": "green", "object": "green", "ball": "green"}


def resolve(instruction, model, data, place_geom="place_pad"):
    """Return (pick_body, pick_xyz, place_xyz, color) for an instruction.

    Examples: "pick the red one and put it on the pad", "grab blue", "pick it up".
    Falls back to the green/default object when no colour is named.
    """
    t = (instruction or "").lower()
    color = next((c for c in COLOR_BODY if c in t), None)
    if color is None:
        for k, v in _ALIASES.items():
            if k in t:
                color = v
                break
    color = color or "green"
    pick_body = COLOR_BODY.get(color, "object")

    pid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, pick_body)
    if pid < 0:
        pick_body, color = "object", "green"
        pid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "object")
    pick_xyz = data.xpos[pid].copy()

    pp = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, place_geom)
    place_xyz = model.geom_pos[pp].copy()
    place_xyz[2] = 0.526
    return pick_body, pick_xyz, place_xyz, color