# 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

# Every low-level motor primitive is already implemented inside skill_code
# (pick, place, move, rotate, pull …)
from skill_code import *          # noqa: F401, F403 – we really want the star import here

from video import (
    init_video_writers,
    recording_step,
    recording_get_observation,
)

# Utility that returns a dictionary  {object_name: (x, y, z)}
from object_positions import get_object_positions


# ------------------------------------------------------------
# Helper functions – purely for high-level flow and logging
# ------------------------------------------------------------
def run_exploration_phase(env, task):
    """
    Very small exploration routine whose sole purpose is to exercise the
    ‘rotate’ primitive so that the planner / learning system can observe
    the (missing) predicate:  (rotated ?g ?a).

    Nothing in here is task specific – we just iterate over a few angles
    and grippers, catching exceptions so that we never break the overall
    run if a particular combination is unsupported in the current scene.
    """
    print("\n===== Exploration Phase : Looking for missing predicate ‘rotated’ =====")
    # Typical gripper naming convention in RLBench
    candidate_grippers = ["left_gripper", "right_gripper", "gripper"]
    candidate_angles   = [
        ("zero_deg", "ninety_deg"),
        ("ninety_deg", "one_eighty_deg"),
        ("one_eighty_deg", "two_seventy_deg"),
        ("two_seventy_deg", "zero_deg"),
    ]

    for g in candidate_grippers:
        for frm, to in candidate_angles:
            try:
                print(f"[Exploration]  rotate( g={g}, from={frm}, to={to} )")
                # We do not know the exact signature of the rotate skill
                # because it is provided externally.  The canonical
                # implementation in our benchmarks looks like:
                #
                #     rotate(env, task, gripper, from_angle, to_angle,
                #            max_steps=,   timeout= )
                #
                # but if the signature is different the called function
                # will raise, which we silently ignore (we only need the
                # *attempt* to be executed so that the system can
                # observe the corresponding predicate appearance).
                #
                rotate(
                    env,
                    task,
                    gripper=g,
                    from_angle=frm,
                    to_angle=to,
                    max_steps=30,
                    timeout=4.0,
                )
            except Exception as exc:
                # We deliberately swallow every error – this exploration
                # should never compromise the remainder of the episode.
                print(f"    ↳ rotate failed ({type(exc).__name__}: {exc}) – ignored.")
    print("===== End Exploration Phase =====\n")


def try_pick_and_place(env, task, obj_name, obj_pos, target_pos):
    """
    Very resilient high-level helper that first attempts to *move* close to
    an object, then *pick* it, and finally *place* it at a given target
    pose/location.  Any error is logged and the function simply returns.
    """
    try:
        print(f"[Task] Moving near {obj_name} at {np.round(obj_pos, 3)}")
        move(
            env,
            task,
            target_pos=obj_pos,
            approach_distance=0.20,
            max_steps=80,
            threshold=0.01,
            timeout=6.0,
        )

        print(f"[Task] Picking {obj_name}")
        pick(
            env,
            task,
            target_pos=obj_pos,
            approach_distance=0.12,
            max_steps=120,
            threshold=0.01,
            approach_axis="z",
            timeout=8.0,
        )

        print(f"[Task] Placing {obj_name} at {np.round(target_pos, 3)}")
        place(
            env,
            task,
            target_pos=target_pos,
            approach_distance=0.15,
            max_steps=120,
            threshold=0.01,
            approach_axis="z",
            timeout=8.0,
        )

        print(f"[Task] {obj_name} handled successfully.\n")

    except Exception as exc:
        print(f"[Warning] Could not complete pick-and-place for {obj_name}: {exc}")
        traceback.print_exc()
        print("Continuing with the next object …\n")


# ------------------------------------------------------------
# Main entry point
# ------------------------------------------------------------
def run_skeleton_task():
    print("===== Starting Skeleton Task =====")

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

    try:
        # Reset task to its initial state
        descriptions, obs = task.reset()

        # Optional video recording setup
        init_video_writers(obs)
        task.step           = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # --------------------------------------------------------
        # 1)  EXPLORATION (discover missing predicate "rotated")
        # --------------------------------------------------------
        run_exploration_phase(env, task)

        # --------------------------------------------------------
        # 2)  Retrieve object poses from helper – we rely on the fact
        #     that object_positions returns at least one “target” object
        #     and (optionally) a disposal / goal area named “bin”.
        # --------------------------------------------------------
        positions = get_object_positions()
        print(f"[Info] Found {len(positions)} objects via ‘get_object_positions’.")

        # Heuristic: any name containing “bin” or “goal” is considered the
        # destination for our pick-and-place routine.
        bin_candidates = [k for k in positions if ("bin" in k.lower() or "goal" in k.lower())]
        if bin_candidates:
            goal_name = bin_candidates[0]
            goal_pos  = positions[goal_name]
            print(f"[Info] Using {goal_name} as drop-off location ({np.round(goal_pos, 3)})")
        else:
            # Fallback: put items 10 cm above the first object’s XY with a
            # constant Z offset – good enough for most toy tasks.
            first_key = list(positions.keys())[0]
            goal_pos  = np.array(positions[first_key]) + np.array([0.0, 0.0, 0.10])
            print(f"[Info] No explicit bin found – synthesised drop position {np.round(goal_pos, 3)}")

        # Any object that is *not* the bin/goal will be picked and placed.
        manipulatable_objects = [k for k in positions if k not in bin_candidates]

        # --------------------------------------------------------
        # 3)  Main high-level plan: pick every object once
        # --------------------------------------------------------
        for obj_name in manipulatable_objects:
            obj_pos = positions[obj_name]
            try_pick_and_place(env, task, obj_name, obj_pos, goal_pos)

    finally:
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()