# run_skeleton_task.py  (Completed)

import traceback
import numpy as np
from env import setup_environment, shutdown_environment

# Load every predefined low-level skill exactly as supplied by the platform.
from skill_code import pick, place, move, rotate, pull

from video import (
    init_video_writers,
    recording_step,
    recording_get_observation,
)

from object_positions import get_object_positions


def _safe_call(fn, *args, **kwargs):
    """
    Helper that wraps every skill invocation with a try / except so that, even
    if a single low-level motion fails, the whole script does not crash in an
    uncontrolled fashion.  It returns (obs, reward, done) when the wrapped
    skill itself already follows that convention, or (None, None, False) when
    the call fails.
    """
    try:
        ret = fn(*args, **kwargs)
        # All rl-bench skills are expected to return
        #   obs, reward, done
        # Should the author of the skill return something else, we still stay
        # alive and log it.
        if isinstance(ret, tuple) and len(ret) == 3:
            return ret
        return None, None, False
    except Exception as exc:
        print(f"[SAFE-CALL] Exception while executing {fn.__name__}: {exc}")
        traceback.print_exc()
        return None, None, False


def exploration_phase(env, task, positions):
    """
    Minimal ‘exploration phase’ whose only purpose is to decide whether a
    drawer is locked (the missing piece of information mentioned in the
    feedback).  Because the TRUE PDDL knowledge is not directly accessible, we
    try to physically interact with drawer handles:

         1) get close (move),
         2) rotate the gripper so that it is roughly perpendicular,
         3) grasp / pick the drawer handle,
         4) pull once,
         5) measure how much the handle moved (if any).

    A very small and generic heuristic is used – if, after the pull, the handle
    moves more than `open_threshold` centimetres in the positive Y direction
    we will assume the drawer is *not* locked.
    """
    # NOTE:  Nothing in this function is task-specific.  If the current task
    # does not contain any object whose name includes “drawer”, the function
    # silently exits.
    drawer_keys = [k for k in positions.keys() if "drawer" in k]
    if len(drawer_keys) == 0:
        print("[EXPLORE] No drawer-like object present – skipping exploration.")
        return {}

    knowledge = dict()

    for drawer_name in drawer_keys:
        print(f"[EXPLORE] Evaluating drawer candidate: {drawer_name}")
        drawer_pos = positions[drawer_name]

        # 1) move – we use a conservative approach: keep the TCP at 10 cm above.
        # Because we do not know the drawer’s local orientation we approach
        # straight from world Z axis.
        hover = np.array(drawer_pos) + np.array([0.0, 0.0, 0.10])
        _safe_call(move, env, task, target_pos=hover, max_steps=120)

        # 2) rotate gripper by 90 deg around its local Z so that fingers are
        # facing the handle horizontally (this matches the PDDL predicate
        #   (rotated ?g ninety_deg)
        target_quat = np.array([0, 0, np.sin(np.pi / 4), np.cos(np.pi / 4)])
        _safe_call(rotate, env, task, target_quat=target_quat)

        # 3) descend to the actual handle centre (heuristic 2 cm above to avoid
        # collision with table).
        grasp_pose = np.array(drawer_pos) + np.array([0.0, 0.0, 0.02])
        _safe_call(move, env, task, target_pos=grasp_pose, max_steps=60)

        # 4) try to pick the handle.  (The generic ‘pick’ skill closes fingers
        # around the nearest object under the TCP.)
        _safe_call(pick, env, task)

        # 5) pull backwards in world Y (positive or negative depends on setup,
        # assume negative for a typical front-facing robot mounting).
        pull_distance = np.array([0.0, -0.10, 0.0])
        pull_target = grasp_pose + pull_distance
        _safe_call(move, env, task, target_pos=pull_target, max_steps=100)

        # 6) release
        _safe_call(place, env, task)

        # 7) decide whether it opened
        new_positions = get_object_positions()
        new_handle_pos = new_positions.get(drawer_name, drawer_pos)
        delta_y = new_handle_pos[1] - drawer_pos[1]

        open_threshold = 0.025  # 2.5 cm is enough to say “opened”
        is_locked = abs(delta_y) < open_threshold
        knowledge[drawer_name] = dict(is_locked=is_locked)

        print(
            f"[EXPLORE] Drawer {drawer_name}: Δy = {delta_y:.3f}  →  "
            f"{'LOCKED' if is_locked else 'UNLOCKED'}"
        )

        # Move up again to a neutral configuration
        _safe_call(move, env, task, target_pos=hover, max_steps=60)

    return knowledge


def disposal_phase(env, task, positions):
    """
    A tiny placeholder routine that simply goes through every object whose name
    contains ‘trash’ or ‘garbage’, picks it, and drops it into a canonical
    disposal location if present.  This is not part of the exploration goal
    but demonstrates calling the remaining skills (‘pick’ and ‘place’).
    """
    disposal_candidates = [
        k for k in positions.keys() if ("trash" in k or "garbage" in k)
    ]
    if len(disposal_candidates) == 0:
        print("[DISPOSAL] No trash-like object found – skipping disposal.")
        return

    # Pick a disposal bin if one exists
    disposal_bins = [k for k in positions.keys() if "bin" in k or "disposal" in k]
    bin_target_pos = positions[disposal_bins[0]] if disposal_bins else None

    for obj_key in disposal_candidates:
        print(f"[DISPOSAL] Handling object {obj_key}")
        obj_pos = positions[obj_key]
        hover = np.array(obj_pos) + np.array([0.0, 0.0, 0.12])
        _safe_call(move, env, task, target_pos=hover, max_steps=120)
        _safe_call(move, env, task, target_pos=obj_pos, max_steps=60)
        _safe_call(pick, env, task)

        if bin_target_pos is not None:
            bin_hover = np.array(bin_target_pos) + np.array([0.0, 0.0, 0.20])
            _safe_call(move, env, task, target_pos=bin_hover, max_steps=120)
            _safe_call(place, env, task)
        else:
            # place exactly where we grabbed it from to stay safe (noop)
            _safe_call(place, env, task)


def run_skeleton_task():
    """
    Generic skeleton for running any task in your simulation.  It now
    incorporates a minimal ‘exploration phase’ aimed at discovering whether a
    drawer is locked (missing predicate), followed by a very small disposal
    demo that uses the remaining predefined skills.
    """
    print("========  Starting Skeleton Task  ========")

    env, task = setup_environment()
    try:
        # Each RLBench task uses task.reset() to place the world in its initial
        # configuration and to provide the textual description.
        descriptions, obs = task.reset()

        # Optional video recording
        init_video_writers(obs)
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # Query initial object positions from the helper module.
        positions = get_object_positions()
        if not positions:
            print(
                "[WARNING] The helper module returned an empty dictionary.  The "
                "rest of the script will still run but skills that rely on "
                "absolute poses may fail."
            )

        # ------------------------------------------------------------------
        # 1) Exploration ‑ physically test drawers to infer (is-locked ?d)
        # ------------------------------------------------------------------
        drawer_knowledge = exploration_phase(env, task, positions)
        if drawer_knowledge:
            for d, info in drawer_knowledge.items():
                print(f"[KNOWLEDGE] Drawer {d}: {info}")

        # ------------------------------------------------------------------
        # 2) (Optional) Disposal demo – use pick & place on trash objects
        # ------------------------------------------------------------------
        disposal_phase(env, task, positions)

        print("========  Skeleton Task Completed  ========")

    finally:
        # Even when an exception bubbles up we gently close the simulation so
        # that the Docker container keeps running for subsequent test cases.
        shutdown_environment(env)


if __name__ == "__main__":
    run_skeleton_task()