# run_skeleton_task.py
#
# Concrete implementation that executes the oracle plan required by the
# specification:
#
#   1) rotate – turn the wrist 90 deg
#   2) move   – drawer side position
#   3) move   – drawer anchor / handle (slightly above)
#   4) pull   – open the drawer                    (Frozen Region – DO NOT EDIT)
#
#   5) move   – descend and grasp the drawer handle ( “pick-drawer” analogue )
#   6) move   – retreat to a safe waypoint
#   7) pick   – grasp the rubbish on the table
#   8) move   – bring rubbish above the bin
#   9) place  – drop rubbish into the bin
#
# Only skills declared in the specification are ever invoked.

import numpy as np
from scipy.spatial.transform import Rotation as R

from env import setup_environment, shutdown_environment
from skill_code import rotate, move, pull, pick, place          # predefined skills
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions


# ---------------------------------------------------------------------------
# Helper utilities
# ---------------------------------------------------------------------------
def safe_pose_lookup(name: str, cache: dict) -> np.ndarray:
    """Return the (x, y, z) position of *name*.

    1) Primary source: cached dictionary from get_object_positions().
    2) Fallback:  query simulator directly via PyRep Shape.
    """
    if name in cache:
        return np.asarray(cache[name], dtype=float)

    from pyrep.objects.shape import Shape   # local import (fallback only)
    return np.asarray(Shape(name).get_position(), dtype=float)


def add_z_offset(pos: np.ndarray, dz: float = 0.10) -> np.ndarray:
    """Lift *pos* by *dz* along +Z (useful for collision-free overhead moves)."""
    lifted = np.asarray(pos, dtype=float).copy()
    lifted[2] += dz
    return lifted


# ---------------------------------------------------------------------------
# Main task runner
# ---------------------------------------------------------------------------
def run_skeleton_task() -> None:
    print("===== Starting Skeleton Task =====")

    # ===== Environment Setup =================================================
    env, task = setup_environment()
    try:
        _, 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()

        # Drawer related keys (we use the *bottom* drawer)
        drawer_side_key   = "bottom_side_pos"
        drawer_anchor_key = "bottom_anchor_pos"

        # A high safe waypoint already present in the scene
        safe_waypoint_key = "waypoint1"

        # Rubbish object (may be “rubbish” or one of the generic items)
        rubbish_key = (
            "rubbish"
            if "rubbish" in positions
            else ("item3" if "item3" in positions else "item1")
        )

        # Bin
        bin_key = "bin"

        # Look-up all poses ----------------------------------------------------
        drawer_side_pos   = safe_pose_lookup(drawer_side_key,   positions)
        drawer_anchor_pos = safe_pose_lookup(drawer_anchor_key, positions)
        safe_waypoint_pos = safe_pose_lookup(safe_waypoint_key, positions)
        rubbish_pos       = safe_pose_lookup(rubbish_key,       positions)
        bin_pos           = safe_pose_lookup(bin_key,           positions)

        # Tweaked overhead poses ----------------------------------------------
        drawer_side_pos_up   = add_z_offset(drawer_side_pos,   0.08)
        drawer_anchor_pos_up = add_z_offset(drawer_anchor_pos, 0.02)
        bin_pos_up           = add_z_offset(bin_pos,           0.10)

        # ===== Oracle-Plan Execution =========================================
        done   = False
        reward = 0.0

        # ---------------------------------------------------------------------
        # Step-1 … Step-4 are FIXED (Frozen region – DO NOT EDIT)
        # ---------------------------------------------------------------------
        target_quat = R.from_euler("xyz", [0.0, 0.0, np.pi/2]).as_quat()

        # [Frozen Code Start]
        obs, reward, done = rotate(env, task, target_quat)
        obs, reward, done = move(env, task, drawer_side_pos_up)
        obs, reward, done = move(env, task, drawer_anchor_pos_up)
        obs, reward, done = pull(env, task, pull_distance=0.15, pull_axis="x")
        # [Frozen Code End]

        if done:
            print("[Plan] Task reported done during frozen section – aborting.")
            return

        # ---------------------------------------------------------------------
        # Step-5  “pick-drawer” (descend, close gripper on the handle)
        #
        #        We descend 4 cm from the current ‘anchor_pos_up’ to reach the
        #        handle, then call *pick* which internally closes the gripper.
        # ---------------------------------------------------------------------
        pick_handle_target = drawer_anchor_pos         # precise handle centre
        obs, reward, done = pick(
            env,
            task,
            target_pos       = pick_handle_target,
            approach_distance=0.04,     # only a small descent – already nearby
            approach_axis    ="z",
        )
        if done:
            print("[Plan] Task ended after grasping drawer handle.")
            return

        # ---------------------------------------------------------------------
        # Step-6  MOVE to a safe waypoint (retreat before table operations)
        # ---------------------------------------------------------------------
        obs, reward, done = move(env, task, add_z_offset(safe_waypoint_pos, 0.12))
        if done:
            print("[Plan] Task ended after retreat move.")
            return

        # ---------------------------------------------------------------------
        # Step-7  PICK the rubbish (approach vertically from above)
        # ---------------------------------------------------------------------
        obs, reward, done = pick(
            env,
            task,
            target_pos       = rubbish_pos,
            approach_distance=0.15,
            approach_axis    ="z",
        )
        if done:
            print("[Plan] Task ended after picking rubbish.")
            return

        # ---------------------------------------------------------------------
        # Step-8  MOVE above the bin (hover pose)
        # ---------------------------------------------------------------------
        obs, reward, done = move(env, task, bin_pos_up)
        if done:
            print("[Plan] Task ended while moving to the bin.")
            return

        # ---------------------------------------------------------------------
        # Step-9  PLACE the rubbish into the bin
        # ---------------------------------------------------------------------
        obs, reward, done = place(
            env,
            task,
            target_pos       = bin_pos,
            approach_distance=0.15,
            approach_axis    ="z",
        )

        # ---------------- Final outcome --------------------------------------
        if done:
            print(f"[Plan] Task completed successfully! Final reward: {reward}")
        else:
            print("[Plan] Plan executed but environment still running (done=False).")

    finally:
        # Always shut down the simulation, even if an exception occurred
        shutdown_environment(env)

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


# ---------------------------------------------------------------------------
# Entry-point guard
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    run_skeleton_task()