dovla_cil/tasks/predicates.py

from __future__ import annotations

import math
from typing import Any

from dovla_cil.tasks.schema import RelationSpec, TaskSpec

SUPPORTED_PREDICATES = {
    "inside",
    "near",
    "next_to",
    "left_of",
    "right_of",
    "behind",
    "in_front_of",
    "lifted",
    "opened",
    "closed",
    "grasped",
}


def evaluate_predicate(predicate: RelationSpec, symbolic_state: dict[str, Any]) -> bool:
    """Evaluate one symbolic relation over a lightweight object state.

    Supported state shapes are intentionally permissive for simulator adapters. Object state can
    live either at `state["objects"][object_id]` or directly at `state[object_id]`. Explicit
    relations can also be provided as `state["relations"][relation_name] = [[arg0, arg1], ...]`.
    """

    relation = predicate.name
    args = predicate.args
    if relation not in SUPPORTED_PREDICATES:
        raise KeyError(f"Unsupported predicate: {relation}")
    if _explicit_relation(symbolic_state, relation, args):
        return True

    if relation == "inside":
        _expect_arity(predicate, 2)
        obj_state = _object_state(symbolic_state, args[0])
        return obj_state.get("inside") == args[1] or obj_state.get("container") == args[1]

    if relation in {"near", "next_to"}:
        _expect_arity(predicate, 2)
        threshold = float(symbolic_state.get("near_threshold", 0.5))
        return _distance(symbolic_state, args[0], args[1]) <= threshold

    if relation in {"left_of", "right_of", "behind", "in_front_of"}:
        _expect_arity(predicate, 2)
        obj_pos = _position(symbolic_state, args[0])
        ref_pos = _position(symbolic_state, args[1])
        if relation == "left_of":
            return obj_pos[0] < ref_pos[0]
        if relation == "right_of":
            return obj_pos[0] > ref_pos[0]
        if relation == "behind":
            return obj_pos[1] > ref_pos[1]
        return obj_pos[1] < ref_pos[1]

    if relation == "lifted":
        _expect_arity(predicate, 1)
        obj_state = _object_state(symbolic_state, args[0])
        if "lifted" in obj_state:
            return bool(obj_state["lifted"])
        return _position(symbolic_state, args[0])[2] > float(symbolic_state.get("lifted_z", 0.1))

    if relation == "opened":
        _expect_arity(predicate, 1)
        return _is_open(_object_state(symbolic_state, args[0]))

    if relation == "closed":
        _expect_arity(predicate, 1)
        state = _object_state(symbolic_state, args[0])
        if "closed" in state:
            return bool(state["closed"])
        return not _is_open(state)

    if relation == "grasped":
        _expect_arity(predicate, 1)
        return bool(_object_state(symbolic_state, args[0]).get("grasped", False))

    raise AssertionError(f"Unhandled predicate: {relation}")


def evaluate_task_success(task: TaskSpec, symbolic_state: dict[str, Any]) -> bool:
    return all(
        evaluate_predicate(predicate, symbolic_state) for predicate in task.success_predicates
    )


def near_target(position: float, target_position: float, tolerance: float) -> bool:
    return abs(float(position) - float(target_position)) <= float(tolerance)


def moved_toward_target(
    start_position: float, end_position: float, target_position: float, *, min_delta: float = 1e-6
) -> bool:
    start_distance = abs(float(target_position) - float(start_position))
    end_distance = abs(float(target_position) - float(end_position))
    return end_distance < start_distance - min_delta


def _expect_arity(predicate: RelationSpec, arity: int) -> None:
    if len(predicate.args) != arity:
        raise ValueError(
            f"Predicate {predicate.name!r} expects {arity} args, got {len(predicate.args)}"
        )


def _object_state(symbolic_state: dict[str, Any], object_id: str) -> dict[str, Any]:
    objects = symbolic_state.get("objects", {})
    if isinstance(objects, dict) and object_id in objects:
        value = objects[object_id]
    elif object_id in symbolic_state:
        value = symbolic_state[object_id]
    else:
        raise KeyError(f"Object {object_id!r} is missing from symbolic state")
    if not isinstance(value, dict):
        raise KeyError(f"Object {object_id!r} is missing from symbolic state")
    return value


def _position(symbolic_state: dict[str, Any], object_id: str) -> tuple[float, float, float]:
    state = _object_state(symbolic_state, object_id)
    raw = state.get("position", state.get("pose", state.get("xyz")))
    if raw is None:
        raw = [state.get("x", 0.0), state.get("y", 0.0), state.get("z", 0.0)]
    if isinstance(raw, dict):
        return (
            float(raw.get("x", 0.0)),
            float(raw.get("y", 0.0)),
            float(raw.get("z", 0.0)),
        )
    if isinstance(raw, list | tuple):
        values = [float(value) for value in raw]
        while len(values) < 3:
            values.append(0.0)
        return values[0], values[1], values[2]
    raise ValueError(f"Object {object_id!r} position must be a list, tuple, or dict")


def _distance(symbolic_state: dict[str, Any], left: str, right: str) -> float:
    left_pos = _position(symbolic_state, left)
    right_pos = _position(symbolic_state, right)
    return math.sqrt(sum((left_pos[index] - right_pos[index]) ** 2 for index in range(3)))


def _is_open(state: dict[str, Any]) -> bool:
    if "opened" in state:
        return bool(state["opened"])
    if "open" in state:
        return bool(state["open"])
    if state.get("joint_state") == "open":
        return True
    if "openness" in state:
        return float(state["openness"]) > 0.5
    return False


def _explicit_relation(symbolic_state: dict[str, Any], relation: str, args: list[str]) -> bool:
    relations = symbolic_state.get("relations", {})
    if not isinstance(relations, dict):
        return False
    candidates = relations.get(relation, [])
    if isinstance(candidates, dict):
        candidates = candidates.values()
    normalized_args = tuple(args)
    for candidate in candidates:
        if tuple(candidate) == normalized_args:
            return True
    return False