# run_skeleton_task.py (Completed)

import time
import inspect
import numpy as np

from pyrep.objects.shape import Shape          #  ❗  DO NOT DELETE  – used by some skills internally
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment
from skill_code import *                       #  <-  we only CALL the already-defined skills
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions


# ---------------------------------------------------------------
# Helper utilities
# ---------------------------------------------------------------

def _call_skill(func, *args, **kw):
    """
    Little wrapper that makes it easier to call skills whose
    signatures differ slightly (some use target_pos, some use x, y, z…).

    We look at the function signature and keep only the parameters
    that actually exist.  This prevents crashes when the skeleton
    guesses an argument name that is not present in the concrete
    implementation delivered with the environment.
    """
    sig        = inspect.signature(func)
    final_kw   = {k: v for k, v in kw.items() if k in sig.parameters}
    try:
        return func(*args, **final_kw)
    except Exception as exc:               # graceful degradation
        print(f"[WARN] Skill {func.__name__} raised {exc}")
        return None, 0.0, True             # pretend episode ended


def _safe_get(pos_dict, key, default=(0, 0, 0)):
    """Utility to fetch an object position with a readable warning."""
    if key in pos_dict:
        return pos_dict[key]
    print(f"[WARN] Position for «{key}» not provided – using {default}")
    return default


# ---------------------------------------------------------------
# Main routine
# ---------------------------------------------------------------

def run_skeleton_task():
    """
    Pipeline:
        1.  Reset environment & wrap for recording.
        2.  Collect object positions (drawer-handle, trash, bin).
        3.  ── Exploration Phase ──
            Move next to the handle and ‘pull’ once in order to
            learn the missing predicate lock-known / is-locked.
        4.  If unlocked → open drawer, pick trash, move to bin,
            place trash.
        5.  Shut down cleanly.
    """
    print("===== Starting Skeleton Task =====")

    env, task = setup_environment()
    try:
        # -------------------------------------------------------
        #  Environment / video initialisation
        # -------------------------------------------------------
        descriptions, obs = task.reset()
        init_video_writers(obs)

        task.step           = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # -------------------------------------------------------
        #  Object locations (comes from helper module)
        # -------------------------------------------------------
        pos = get_object_positions()
        handle_pos = _safe_get(pos, 'drawer_handle')
        trash_pos  = _safe_get(pos, 'trash')
        bin_pos    = _safe_get(pos, 'bin')

        print("[INFO] Known object positions:")
        for k, v in pos.items():
            print(f"       {k:15}: {np.round(v, 3)}")

        # =======================================================
        #  1) ── EXPLORE: Approach handle & detect lock status ──
        # =======================================================
        print("\n[Phase-1] Approach drawer handle to identify lock status.")
        obs, reward, done = _call_skill(move, env, task,
                                        target_pos=handle_pos,
                                        approach_distance=0.10,
                                        max_steps=150,
                                        threshold=0.01,
                                        timeout=10.0)
        if done:
            print("[ABORT] Episode ended during move-to-handle.")
            return

        print("[Phase-1] Pull once to learn «lock-known» predicate.")
        obs, reward, done = _call_skill(pull, env, task)
        if done:
            print("[ABORT] Episode ended during exploratory pull.")
            return

        # =======================================================
        #  2) ── OPEN DRAWER (if not already open) ──
        # =======================================================
        #  A second pull usually opens the drawer after learning
        #  the lock state; many benchmark tasks are modelled that
        #  way.  We simply try again.
        print("\n[Phase-2] Pull again – attempt to actually open drawer.")
        obs, reward, done = _call_skill(pull, env, task)
        if done:
            print("[ABORT] Episode ended while opening drawer.")
            return

        # Optional: rotate gripper to ninety degrees if the domain
        # requires ‘side-grasp’.  We fetch quat either from domain
        # knowledge or simply use a Z-axis half-turn.
        target_quat = np.array([0.0, 0.0, np.sin(np.pi/4), np.cos(np.pi/4)])
        print("[Phase-2] Rotate gripper for side approach (≈90°).")
        obs, reward, done = _call_skill(rotate, env, task,
                                        target_quat=target_quat,
                                        threshold=0.05)
        if done:
            print("[ABORT] Episode ended during rotate.")
            return

        # =======================================================
        #  3) ── PICK the TRASH from the opened drawer ──
        # =======================================================
        print("\n[Phase-3] Pick the trash.")
        obs, reward, done = _call_skill(pick, env, task,
                                        target_pos=trash_pos,
                                        approach_distance=0.08,
                                        max_steps=150,
                                        threshold=0.01,
                                        timeout=10.0)
        if done:
            print("[ABORT] Episode ended during pick.")
            return

        # =======================================================
        #  4) ── MOVE to bin & PLACE the trash ──
        # =======================================================
        print("\n[Phase-4] Move to bin.")
        obs, reward, done = _call_skill(move, env, task,
                                        target_pos=bin_pos,
                                        approach_distance=0.10,
                                        max_steps=150,
                                        threshold=0.01,
                                        timeout=10.0)
        if done:
            print("[ABORT] Episode ended while moving to bin.")
            return

        print("[Phase-4] Place trash into bin.")
        obs, reward, done = _call_skill(place, env, task,
                                        target_pos=bin_pos,
                                        approach_distance=0.05,
                                        max_steps=100,
                                        threshold=0.01,
                                        timeout=10.0)
        if done:
            print("[COMPLETE] Task finished after placing trash!")
            return

        # =======================================================
        #  5) ── (Optional) close the drawer again ──
        # =======================================================
        # Many benchmarks do not require re-closing the drawer.
        # If your specific domain does, you might push the handle
        # back in here.  We keep it optional for robustness.

    finally:
        shutdown_environment(env)
        print("===== End of Skeleton Task =====")


if __name__ == "__main__":
    run_skeleton_task()