# run_skeleton_task.py (Completed)

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

from env import setup_environment, shutdown_environment

# Import every primitive that is already defined for us
from skill_code import *

from video import init_video_writers, recording_step, recording_get_observation

from object_positions import get_object_positions


def _has_predicate(obs, name: str) -> bool:
    """
    Very light-weight predicate search.
    The observation object that comes from RLBench is usually either a dict or a
    tuple containing low-dimensional strings that include the symbolic state
    the task wrapper keeps.  Whatever its exact structure is, converting the
    whole thing into a single long string and looking for a “(predicate …”
    token is a robust, domain independent way of checking whether the
    predicate appears.
    """
    try:
        s = str(obs)
    except Exception:
        # Fallback: if obs is not stringifiable in a normal way just do a repr
        s = repr(obs)
    target1 = f"({name} "
    target2 = f"({name})"
    return (target1 in s) or (target2 in s)


def _find_grippers_in_positions(positions: dict):
    """
    Heuristic helper that extracts every key that contains 'gripper'
    (case-insensitive) from the position dictionary returned by
    get_object_positions().  If none exist we return an empty list and let the
    caller decide how to proceed.
    """
    grippers = []
    for k in positions.keys():
        if "gripper" in k.lower():
            grippers.append(k)
    return grippers


def run_skeleton_task():
    """
    Generic skeleton for running any task in your simulation.
    The code has been completed so that it now:
      1) Detects that the predicate ‘rotated’ is missing from the initial
         state description.
      2) Performs a short exploration phase that tries to execute the built-in
         ‘rotate’ primitive on every detected gripper in order to introduce the
         ‘rotated’ predicate into the symbolic state.
    NOTE:  The *actual* downstream task logic (e.g. opening a drawer, moving
    objects, etc.) is intentionally left minimal because the focal point for
    this milestone is the exploration necessary to discover the missing
    predicate.  After the exploration succeeds you could plug your domain
    specific plan right after the comment that says “# === Your domain plan
    goes here ===”.
    """
    print("===== Starting Skeleton Task =====")

    # === Environment Setup ===
    env, task = setup_environment()
    try:
        # Reset the task to its initial state and obtain the very first
        # low-dimensional observation
        descriptions, obs = task.reset()

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

        # ------------------------------------------------------------
        #               1)  Retrieve object positions
        # ------------------------------------------------------------
        try:
            positions = get_object_positions()
        except Exception as e:
            # If, for any reason, the helper is unavailable we fall back to an
            # empty dictionary.  The rest of the code can still operate by
            # trying generic names such as “gripper”.
            print("[Warning] get_object_positions() failed –", e)
            positions = {}

        # ------------------------------------------------------------
        #     2)  MISSING-PREDICATE EXPLORATION (“rotated”)
        # ------------------------------------------------------------
        print("[Exploration] Checking whether predicate ‘rotated’ is present in the initial observation …")
        if not _has_predicate(obs, "rotated"):
            print("[Exploration] ‘rotated’ not found.  Starting exploration using the built-in rotate skill.")
            discovered = False

            # Heuristically collect plausible gripper names
            candidate_grippers = _find_grippers_in_positions(positions)
            if not candidate_grippers:
                # Ensure at least *one* name is attempted – many RLBench
                # environments simply call the active end-effector “gripper”
                candidate_grippers = ["gripper"]

            # Typical angle constants used in the drawer domain
            from_angles = ["zero_deg", "0_deg", "zero", "0"]
            to_angle = "ninety_deg"

            for g in candidate_grippers:
                if discovered:
                    break

                for a in from_angles:
                    try:
                        # Attempt the rotation.  The rotate primitive’s exact
                        # signature may differ depending on the underlying
                        # implementation, so we try the most common pattern
                        # (env, task, gripper_name, from_angle, to_angle)
                        print(f"[Exploration] Trying rotate({g}, {a} -> {to_angle})")
                        rotate(env, task, g, a, to_angle)     # noqa – we do not care about its return value
                    except Exception as exc:
                        # Keep iterating through other combinations before we
                        # give up; print but do not crash the entire run
                        print(f"[Exploration]   rotation attempt failed ({exc}).")
                        continue

                    # Immediately fetch a fresh observation
                    new_obs = task.get_observation()
                    if _has_predicate(new_obs, "rotated"):
                        print(f"[Exploration] Predicate ‘rotated’ discovered after rotating gripper '{g}'.")
                        discovered = True
                        obs = new_obs     # Keep the updated observation
                        break   # leave the angle loop

            if not discovered:
                print("[Exploration] WARNING: Unable to discover predicate ‘rotated’.  "
                      "Subsequent actions that rely on it may fail.")
        else:
            print("[Exploration] ‘rotated’ is already present – no exploration needed.")

        # ------------------------------------------------------------
        #               3)  YOUR DOMAIN SPECIFIC PLAN
        # ------------------------------------------------------------
        # As soon as the rotated predicate is guaranteed to be available you can
        # write a high level plan – e.g. open the drawer and dispose of an
        # object – by calling further primitives like move, pick, pull, place,
        # etc.  The following placeholder illustrates the pattern without
        # assuming any details about the concrete environment.
        #
        # Example (commented so it does not break if the referenced objects are
        # not present in the current benchmark):
        #
        # try:
        #     print("[Plan] Moving close to ‘drawer_handle’ …")
        #     move(env, task,
        #          target_pos=positions.get('drawer_handle', (0, 0, 0)),
        #          approach_distance=0.10,
        #          max_steps=150,
        #          threshold=0.01,
        #          approach_axis='z',
        #          timeout=10.0)
        #
        #     print("[Plan] Pulling the drawer open …")
        #     pull(env, task, 'drawer_handle_location')
        #
        # except Exception as plan_exc:
        #     print("[Plan] Encountered an error in the high-level plan –")
        #     traceback.print_exc()

        # Nothing else to do for this submission
        print("[Task] Execution finished (demo ends here).")

    except Exception:
        print("[Error] Unhandled exception during run:")
        traceback.print_exc()
    finally:
        # Always shut the environment down, even when errors occur
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()
