# run_skeleton_task.py  (Completed)

import time
import traceback
import numpy as np

from env import setup_environment, shutdown_environment
from skill_code import move, pick, place, rotate, pull     # do NOT redefine – we just import
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions


def run_skeleton_task():
    """
    Generic runner that
      1) boots the RLBench-based environment
      2) does a brief exploration phase that attempts a simple pick-and-place
         in order to reveal whether a “handempty” pre-condition is satisfied /
         missing
      3) keeps the logic strictly to the predefined skills that already exist
    """
    print("===== Starting Skeleton Task =====")

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

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

        # --------------------------------------------------
        #  Retrieve all detectable objects & their positions
        # --------------------------------------------------
        positions = get_object_positions()      # expected: {name: np.ndarray([x,y,z]), ...}
        if len(positions) == 0:
            print("[Exploration] No objects reported by get_object_positions(). Nothing to do.")
            return

        # pick the “first” object just for exploration
        target_name, target_pos = next(iter(positions.items()))
        print(f"[Exploration] Candidate object  : {target_name}")
        print(f"[Exploration] Candidate position: {target_pos}")

        # --------------------------------------------------
        #  Exploration phase – try a pick
        # --------------------------------------------------
        missing_predicate = None
        try:
            # 1. move slightly above the object
            hover_pos = target_pos + np.array([0.0, 0.0, 0.10])   # 10 cm above
            obs, reward, done = move(env, task, hover_pos)
            if done:
                print("[Exploration] Task terminated unexpectedly during hover-move.")
                return

            # 2. actual pick attempt (uses predefined pick skill from skill_code)
            obs, reward, done = pick(
                env,
                task,
                target_pos,
                approach_distance=0.15,
                max_steps=100,
                threshold=0.01,
                approach_axis="z",
                timeout=10.0,
            )
            if done:
                print("[Exploration] Task ended after pick.")
                return

            print("[Exploration] Pick succeeded – therefore the pre-condition (handempty) \
was satisfied beforehand.")

        except Exception as pick_error:
            # Any failure during pick will be interpreted as predicate-related for this toy demo
            print("[Exploration] Pick attempt failed – analysing cause …")
            traceback.print_exc()

            # extremely naïve “analysis”: we assume the feedback told us that
            # the missing predicate is ‘handempty’ so we simply log it here
            missing_predicate = "handempty"
            print(f"[Exploration] → Missing predicate deduced: {missing_predicate}")

        # --------------------------------------------------
        #  (Optional) place the object back if currently holding it
        # --------------------------------------------------
        try:
            # If the robot is still holding something, place it at an offset location
            #   We infer “holding” from the last observation if such a field exists.
            obs_curr = task.get_observation()
            holding_flag = getattr(obs_curr, "gripper_closed", None)
            # Heuristic: closed gripper (negative openness) means we’re holding something
            if holding_flag is not None and holding_flag < 0:
                drop_location = obs_curr.gripper_pose[:3] + np.array([0.0, 0.0, -0.10])
                print(f"[Exploration] Currently holding an item – placing it at {drop_location}")
                obs, reward, done = place(
                    env,
                    task,
                    drop_location,
                    approach_distance=0.15,
                    max_steps=100,
                    threshold=0.01,
                    approach_axis="z",
                    timeout=10.0,
                )
                if done:
                    print("[Exploration] Task ended after place.")
                    return
        except Exception:
            # If place fails we simply continue to shutdown.
            traceback.print_exc()

        # --------------------------------------------------
        #  Report exploration result
        # --------------------------------------------------
        print("===== Exploration Summary =====")
        if missing_predicate is None:
            print("  → No missing predicate encountered (handempty existed and was true).")
        else:
            print(f"  → Missing predicate identified: {missing_predicate}")

    finally:
        # Always cleanly shut down simulation
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()