# run_skeleton_task.py  (Completed)

import time
import random
import numpy as np
from pyrep.objects.shape import Shape
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment

# We import * because every primitive‐skill’s API (move, pick, place, rotate, pull …)
# is already implemented in skill_code and we must not re-implement them here.
from skill_code import *

from video import init_video_writers, recording_step, recording_get_observation

# Helper that returns the (x, y, z) position of every relevant object in the scene.
from object_positions import get_object_positions


# -----------------------------------------------------------
# Helper utilities – these do NOT add new skills, they only
# provide local book-keeping / reasoning functionality.
# -----------------------------------------------------------
ANCHOR_KEYWORDS = ["anchor", "handle", "knob"]


def _find_anchor_in_names(names):
    """Heuristic:  look for an object name that obviously represents a drawer
    handle / anchor position."""
    for key in names:
        lowered = key.lower()
        for kw in ANCHOR_KEYWORDS:
            if kw in lowered:
                return key
    return None


def _explore_for_anchor(env, task, positions):
    """
    Basic exploration phase that tries to discover the object/location that
    corresponds to the missing predicate  (is-anchor-pos ?p ?d).

    Strategy:
      1.  First, look for intuitive keywords in the object name.
      2.  If that fails, iterate through every object position, approach it,
          try a light pick action – if the pick succeeds *and* a subsequent
          pull succeeds, we assume we found the anchor.
    """
    object_names = list(positions.keys())

    # --- 1) Name-based quick guess -----------------------------------------
    probable_anchor_name = _find_anchor_in_names(object_names)
    if probable_anchor_name:
        print(f"[Exploration] Anchor candidate by name: {probable_anchor_name}")
        return probable_anchor_name

    # --- 2) Behaviour-based probing ----------------------------------------
    # NOTE:  We do *not* create any new skills, we only call already existing
    #        primitives inside a try/except so we can gracefully back-off if
    #        an action fails for that object.
    print("[Exploration] No obvious anchor name found –  probing objects …")

    for obj_name in object_names:
        pos = positions[obj_name]
        print(f"    · Probing object: {obj_name}")

        # Try to *approach* the object.  We assume `move` takes an (x, y, z)
        # target position in world coordinates.
        try:
            move(env, task, target_pos=pos, max_steps=150)
        except Exception as e:
            print(f"        [move] Could not reach {obj_name}  – {e}")
            continue

        # Try a gentle pick – if the pick raises we consider it ‘non-pickable’.
        try:
            pick(env, task, target_obj=obj_name, target_pos=pos,
                 approach_distance=0.10, max_steps=150)
        except Exception:
            # Not pick-able – continue with next object
            continue

        # If we succeeded in picking, try to PULL a little.  A non-drawer
        # object should not succeed the pull or might throw an exception.
        pulled = False
        try:
            pull(env, task, target_obj=obj_name, pull_distance=0.05)
            pulled = True
        except Exception:
            pass

        # Put the object back if it is obviously the wrong one.
        try:
            place(env, task, target_pos=pos)
        except Exception:
            # If placing fails, just ignore – we will reset at the end anyway.
            pass

        if pulled:
            print(f"[Exploration]  SUCCESS – '{obj_name}' behaves like a handle.")
            return obj_name

    raise RuntimeError("Could not identify anchor position – exploration failed.")


# -----------------------------------------------------------
# Main routine
# -----------------------------------------------------------
def run_skeleton_task():
    """Generic skeleton for running the combined-drawer task while
    automatically discovering the anchor position (missing predicate
    is-anchor-pos) via exploration."""
    print("===== Starting Skeleton Task =====")

    # === Environment Setup ==================================================
    env, task = setup_environment()

    try:
        # Reset to a known initial state
        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)

        # === Retrieve Object Positions ======================================
        positions = get_object_positions()     # { name : (x, y, z) }

        # ----------------------------------------------------------
        # 1)  Exploration:  identify which position satisfies the
        #     predicate  (is-anchor-pos ?p ?drawer)
        # ----------------------------------------------------------
        anchor_obj_name = _explore_for_anchor(env, task, positions)
        anchor_pos = positions[anchor_obj_name]
        print(f"[Plan] Using '{anchor_obj_name}' as drawer handle / anchor.")

        # ----------------------------------------------------------
        # 2)  Oracle-style plan execution
        #
        #     High-level plan:
        #        a)  Move to the anchor position.
        #        b)  Rotate gripper to be perpendicular.
        #        c)  Pick / grasp the handle.
        #        d)  Pull to open drawer.
        #        e)  (Optional) Place an item or just finish.
        # ----------------------------------------------------------

        # (a) Move to anchor
        move(env, task, target_pos=anchor_pos, max_steps=180)

        # (b) Rotate the wrist so we can grip the handle sideways
        try:
            rotate(env, task, target_angle="ninety_deg")
        except Exception:
            # Some tasks may already start in correct orientation – ignore.
            print("[Plan] rotate() raised but continuing – assuming orientation OK.")

        # (c) Pick / grasp the handle
        pick(env, task, target_obj=anchor_obj_name, target_pos=anchor_pos,
             approach_distance=0.08, max_steps=120)

        # (d) Pull to open
        pull(env, task, target_obj=anchor_obj_name, pull_distance=0.15)

        # ----------------------------------------------------------
        # At this point the drawer should be open.  If the task
        # requires additional behaviour (e.g. placing trash into
        # the drawer), you can extend the script right here WITHOUT
        # defining new skills – only call existing ones.
        # ----------------------------------------------------------

        print("[Plan] Drawer opened successfully – task complete!")

    finally:
        # =========================================================
        # Clean shutdown of the simulator
        # =========================================================
        shutdown_environment(env)

    print("===== End of Skeleton Task =====")


# Direct execution entry-point
if __name__ == "__main__":
    run_skeleton_task()