# run_skeleton_task.py (Completed)

import numpy as np

from pyrep.objects.shape import Shape

from env import setup_environment, shutdown_environment
from skill_code import *            # use pre-defined primitives only
from video import (init_video_writers,
                   recording_step,
                   recording_get_observation)
from object_positions import get_object_positions


# --------------------------------------------------------------------------
# Exploration: explicitly try to discover whether the predicate “rotated”
# can be achieved in the current world by calling the available rotate skill.
# We simply iterate over every shape that looks like a gripper and ask the
# low-level “rotate” skill to change its configuration from one symbolic
# angle to another.  Any success/failure information is printed so that an
# oracle (or an automatic learner) could infer the truth value of
# (rotated ?g ?a) facts.
# --------------------------------------------------------------------------
def explore_rotated_predicate(env, task) -> None:
    print("[Exploration] ===  Searching for ’rotated’ predicate support  ===")

    # RLBench usually names the end-effector with the word “gripper”
    gripper_shapes = [obj for obj in Shape.get_all()
                      if 'gripper' in obj.get_name().lower()]

    # Symbolic angle names that also appear in the PDDL domain
    candidate_angles = ['zero_deg', 'ninety_deg', 'oneeighty_deg']

    if not gripper_shapes:
        print("[Exploration] No gripper objects discovered – skipping rotation test.")
        return

    for gripper in gripper_shapes:
        g_name = gripper.get_name()
        #  try a round-robin rotation sequence               (0 → 90 → 180 → 0)
        for from_a, to_a in zip(candidate_angles,
                                candidate_angles[1:] + candidate_angles[:1]):
            try:
                print(f"[Exploration] rotate({g_name}, {from_a} → {to_a})")
                # the predefined rotate skill is called with keyword
                # arguments that most RLBench wrappers accept
                rotate(env=env,
                       task=task,
                       gripper_name=g_name,
                       from_angle=from_a,
                       to_angle=to_a)

            except Exception as err:
                # Non-fatal; we only want information
                print(f"[Exploration]    failed: {err}")

    print("[Exploration] ===  Finished rotation exploration  ===")


# --------------------------------------------------------------------------
#  Main routine – keeps the original skeleton structure but now
#  (1) sets up exploration for missing predicates,
#  (2) demonstrates a concise task using only the predefined skills.
# --------------------------------------------------------------------------
def run_skeleton_task():
    print("===== Starting Skeleton Task =====")

    # === Environment Setup ===
    env, task = setup_environment()
    try:
        # reset task
        descriptions, obs = task.reset()

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

        # retrieve current object positions (helper util)
        positions = get_object_positions()

        # --------------------------------------------------
        # 1)  EXPLORATION PHASE  – look for missing predicate
        # --------------------------------------------------
        explore_rotated_predicate(env, task)

        # --------------------------------------------------
        # 2)  DEMONSTRATION / EXECUTION PHASE (very generic)
        #     – pick the first manipulable object we find,
        #       place it slightly to the side,
        #       then pull the first drawer we can locate.
        # --------------------------------------------------

        # --- Choose an object to pick ---
        movable_objs = [name for name in positions
                        if any(k in name.lower()
                               for k in ['cube', 'block', 'object'])]
        if movable_objs:
            obj = movable_objs[0]
            obj_pos = positions[obj]
            print(f"[Task] Picking {obj} at {obj_pos}")

            pick(env=env,
                 task=task,
                 target_pos=obj_pos,
                 approach_distance=0.15,
                 max_steps=120,
                 threshold=0.01,
                 approach_axis='z',
                 timeout=10.0)

            # After picking, place slightly offset
            place_pos = np.array(obj_pos) + np.array([0.15, 0.0, 0.0])
            print(f"[Task] Placing {obj} at {place_pos}")

            place(env=env,
                  task=task,
                  target_pos=place_pos,
                  approach_distance=0.15,
                  max_steps=120,
                  threshold=0.01,
                  approach_axis='z',
                  timeout=10.0)

        # --- Try to pull a drawer if present ---
        drawers = [name for name in positions if 'drawer' in name.lower()]
        if drawers:
            drawer = drawers[0]
            drawer_pos = positions[drawer]
            print(f"[Task] Moving to drawer {drawer} at {drawer_pos}")

            move(env=env,
                 task=task,
                 target_pos=drawer_pos,
                 approach_distance=0.2,
                 max_steps=100,
                 threshold=0.02,
                 approach_axis='z',
                 timeout=10.0)

            print(f"[Task] Pulling drawer {drawer}")
            pull(env=env,
                 task=task,
                 handle_name=drawer,
                 max_steps=120,
                 timeout=10.0)

    finally:
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()