# run_skeleton_task.py (Completed)

import time
import numpy as np

from env import setup_environment, shutdown_environment

# NOTE: the next import brings *all* officially provided skill functions
# (pick, place, move, rotate, pull …).  DO NOT re-declare any of them.
from skill_code import *

from video import (                     # utilities for optional video capture
    init_video_writers,
    recording_step,
    recording_get_observation,
)

# Helper that gives us every named object’s 3-D position in the scene
from object_positions import get_object_positions


def run_skeleton_task() -> None:
    """
    Generic driver that ❶ explores the scene so that every “unknown” predicate
    (identified / temperature-known / weight-known / durability-known /
    lock-known …) can become TRUE and ❷ then executes a minimal drawer-opening
    scenario that demonstrates usage of the officially provided skills:
    move → rotate → pull.
    """
    print("\n==========  STARTING  GENERIC  TASK  ==========\n")

    # -------------  Environment initialisation -------------
    env, task = setup_environment()
    try:
        # ── reset the RLBench task
        descriptions, obs = task.reset()

        # ── optional video writers
        init_video_writers(obs)

        # ── wrap step / observation so that every call is automatically recorded
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # -------------  GLOBAL SCENE INFORMATION -------------
        # Dictionary: {object_name: np.ndarray([x, y, z])}
        positions = get_object_positions()
        if not positions:
            print("[INIT]  get_object_positions() returned an empty dict –"
                  " nothing to do.")
            return

        # -----------------------------------------------------
        # 1)  EXPLORATION PHASE  –  find *all* “unknown” objects.
        #     (this mimics the ‘move’ / ‘pick’ variations in the
        #      exploration-PDDL that turn on identified/temperature/weight …)
        # -----------------------------------------------------
        print("[EXPLORATION]  Visiting every location so that identification-"
              "related predicates become true.")
        robot_home = obs.gripper_pose[:3]     # remember where we started

        for obj_name, obj_pos in positions.items():
            print(f"[EXPLORATION]  → moving close to  {obj_name}  at {obj_pos}")

            # the skill API always follows  (env, task, …other-kwargs)
            try:
                obs, reward, done = move(
                    env,
                    task,
                    target_pos=obj_pos,           #   <-- xyz world position
                    approach_distance=0.15,
                    max_steps=120,
                    threshold=0.01,
                    approach_axis="z",
                    timeout=10.0,
                )
            except Exception as exc:
                print(f"[WARN]  move() failed for {obj_name}: {exc}")
                continue                            # try the remaining objects

            if done:                               # the task finished early
                print("[INFO]  Task ended during exploration.")
                return

            # OPTIONAL: pick the object once so that *weight-known* /
            #           *durability-known* become true (cf. exploration PDDL)
            try:
                obs, reward, done = pick(
                    env,
                    task,
                    target_pos=obj_pos,
                    approach_distance=0.05,
                    grasp_distance=0.01,
                    max_steps=80,
                    timeout=8.0,
                )
                # immediately place it back where we found it
                if not done:
                    obs, reward, done = place(
                        env,
                        task,
                        target_pos=obj_pos,
                        release_height=0.05,
                        max_steps=80,
                        timeout=8.0,
                    )
            except Exception as exc:
                # not every object is graspable – that is fine
                print(f"[NOTE]  Could not pick/place {obj_name}: {exc}")

            if done:
                print("[INFO]  Task ended during exploration.")
                return

        # -----------------------------------------------------
        # 2)  TASK-SPECIFIC PLAN  –  open a drawer.
        # -----------------------------------------------------
        # heuristically pick the first object whose name implies a drawer handle
        drawer_handle_name = None
        for name in positions.keys():
            if "drawer" in name or "handle" in name:
                drawer_handle_name = name
                break

        if drawer_handle_name is None:
            print("[PLAN]  No drawer handle detected – skipping the "
                  "drawer-opening part.")
            return

        handle_pos = positions[drawer_handle_name]
        print(f"[PLAN]  Opening Drawer – using handle object "
              f"'{drawer_handle_name}' at {handle_pos}")

        # (i)  move the robot close to the drawer handle
        obs, reward, done = move(
            env,
            task,
            target_pos=handle_pos,
            approach_distance=0.10,
            max_steps=120,
            threshold=0.01,
            approach_axis="y",       # come in sideways – tends to work better
            timeout=10.0,
        )
        if done:
            print("[PLAN]  Task finished during move-to-handle.")
            return

        # (ii) rotate the wrist so that the gripper is aligned with the handle
        #      (‘ninety_deg’ in PDDL ≅ world   z-axis facing lateral direction)
        #      Here we simply use a quaternion that rotates the gripper 90°
        #      about its local x-axis – *modify as needed for your scene*.
        ninety_deg_quat = np.array([0.7071, 0.0, 0.0, 0.7071])
        obs, reward, done = rotate(
            env,
            task,
            target_quat=ninety_deg_quat,
            max_steps=100,
            threshold=0.05,
            timeout=8.0,
        )
        if done:
            print("[PLAN]  Task finished during rotate.")
            return

        # (iii) pull – attempt to slide the drawer open
        obs, reward, done = pull(
            env,
            task,
            pull_distance=0.20,      # how far we wish to pull
            max_steps=120,
            threshold=0.01,
            timeout=10.0,
        )
        if done:
            print("[PLAN]  Drawer opened successfully – task complete!")
        else:
            print("[PLAN]  Finished pull() but task not marked as done.")

        # -----------------------------------------------------
        # 3)  Return to home position (optional clean-up)
        # -----------------------------------------------------
        try:
            move(
                env,
                task,
                target_pos=robot_home,
                approach_distance=0.10,
                max_steps=120,
                threshold=0.01,
                approach_axis="z",
                timeout=10.0,
            )
        except Exception:
            pass                                          # not mission-critical

    finally:
        # Always shut down the RLBench instance – prevents orphaned processes
        shutdown_environment(env)
        print("\n==========  END  OF  TASK  ==========\n")


if __name__ == "__main__":
    run_skeleton_task()