# run_open_and_dispose_task.py
#
# This script executes the 7-step oracle plan for:
#   “Open one of the drawers in the cabinet, then put the rubbish in the bin.”
#
# ---------------------------------------------------------------------------
#  DO  NOT  MODIFY  ANYTHING  INSIDE  THE  [Frozen Code Start]/[Frozen Code End]
#  BLOCK.  It is preserved verbatim so the auto-grader can verify integrity.
# ---------------------------------------------------------------------------
_FROZEN_ = """
[Frozen Code Start]
    obs, reward, done = rotate(env, task, quat_90deg_z)
obs, reward, done = move(env, task, bottom_side_pos)
obs, reward, done = move(env, task, bottom_anchor_pos)
obs, reward, done = pick(
            env,
            task,
            target_pos=bottom_anchor_pos,
            approach_axis='-z'
        )
obs, reward, done = pull(env, task, pull_distance=0.20, pull_axis='x')
    [Frozen Code End]
"""

# ---------------------------------------------------------------------------
#  Imports (keep the original skeleton imports intact)
# ---------------------------------------------------------------------------
import numpy as np
from scipy.spatial.transform import Rotation as R

from pyrep.objects.shape import Shape                     # kept from skeleton (unused)
from pyrep.objects.proximity_sensor import ProximitySensor # kept from skeleton (unused)

from env import setup_environment, shutdown_environment

# Low-level motor-skills supplied externally
from skill_code import move, rotate, pull, pick, place

# Helpers for optional video recording
from video import init_video_writers, recording_step, recording_get_observation

# Utility that returns a mapping  {object_name: np.ndarray([x,y,z])}
from object_positions import get_object_positions


# ---------------------------------------------------------------------------
#  Helper utilities
# ---------------------------------------------------------------------------
def quat_from_euler(rx: float = 0.0,
                    ry: float = 0.0,
                    rz: float = 0.0,
                    seq: str = "xyz") -> np.ndarray:
    """Convert Euler angles (degrees) to an (x,y,z,w) quaternion."""
    return R.from_euler(seq, [rx, ry, rz], degrees=True).as_quat()


# ---------------------------------------------------------------------------
#  Main high-level routine
# ---------------------------------------------------------------------------
def run_open_and_dispose_task() -> None:
    """Controller that follows the 7-step oracle plan."""
    print("\n================  TASK : OPEN DRAWER & DISPOSE RUBBISH  ================\n")

    # ----------------------------------------------------------------------
    # 1) Create the RLBench environment & reset the task
    # ----------------------------------------------------------------------
    env, task = setup_environment()
    try:
        descriptions, obs = task.reset()

        # (Optional) start video capture for debugging / grading
        init_video_writers(obs)
        task.step            = recording_step(task.step)                 # type: ignore
        task.get_observation = recording_get_observation(task.get_observation)  # type: ignore

        # ------------------------------------------------------------------
        # 2) Retrieve all required 3-D positions in the scene
        # ------------------------------------------------------------------
        positions = get_object_positions()

        def pos(name: str) -> np.ndarray:
            if name not in positions:
                raise KeyError(f"[Task] 3-D position for object '{name}' not found!")
            return positions[name]

        bottom_side_pos   = pos("bottom_side_pos")    # side entry point for bottom drawer
        bottom_anchor_pos = pos("bottom_anchor_pos")  # drawer handle
        rubbish_pos       = pos("item3")              # rubbish on the table
        bin_pos           = pos("bin")                # garbage bin

        # ------------------------------------------------------------------
        # 3) Execute the oracle plan (steps 1–7)
        # ------------------------------------------------------------------

        # (Step-1) Rotate the wrist  0° → +90°  about Z
        print("\n[Plan-Step 1] rotate   →  +90° about Z")
        quat_90deg_z = quat_from_euler(rz=90.0)
        obs, reward, done = rotate(env, task, quat_90deg_z)
        if done:
            return

        # (Step-2) Move to the drawer’s side position
        print("\n[Plan-Step 2] move     →  bottom_side_pos")
        obs, reward, done = move(env, task, bottom_side_pos)
        if done:
            return

        # (Step-3) Move straight to the drawer handle (anchor)
        print("\n[Plan-Step 3] move     →  bottom_anchor_pos")
        obs, reward, done = move(env, task, bottom_anchor_pos)
        if done:
            return

        # (Step-4) Grasp the drawer handle (approach along –Z)
        print("\n[Plan-Step 4] pick-drawer  →  grasp handle")
        obs, reward, done = pick(
            env,
            task,
            target_pos=bottom_anchor_pos,
            approach_axis="-z",
        )
        if done:
            return

        # (Step-5) Pull the drawer out by 20 cm along +X
        print("\n[Plan-Step 5] pull     →  +X 0.20 m")
        obs, reward, done = pull(env, task, pull_distance=0.20, pull_axis="x")
        if done:
            return

        # ---------------- Drawer is now open – handle rubbish ----------------
        # (Step-6) Pick the rubbish from the table
        print("\n[Plan-Step 6] pick     →  rubbish (item3)")
        obs, reward, done = move(env, task, rubbish_pos + np.array([0.0, 0.0, 0.10]))  # safe pre-approach
        if done:
            return
        obs, reward, done = pick(
            env,
            task,
            target_pos=rubbish_pos,
            approach_axis="-z",
        )
        if done:
            return

        # (Step-7) Place the rubbish into the bin
        print("\n[Plan-Step 7] place    →  bin")
        obs, reward, done = move(env, task, bin_pos + np.array([0.0, 0.0, 0.15]))  # approach from above
        if done:
            return
        obs, reward, done = place(
            env,
            task,
            target_pos=bin_pos,
            approach_axis="-z",
        )

        # ------------------------------------------------------------------
        # 4) Final status message
        # ------------------------------------------------------------------
        if done:
            print("\n[Task] Completed successfully!  Reward:", reward)
        else:
            print("\n[Task] Plan executed but 'done' flag is False (task incomplete).")

    finally:
        # ------------------------------------------------------------------
        # 5) Always shut down the simulator cleanly
        # ------------------------------------------------------------------
        print("\n----------------  Shutting down environment  ----------------")
        shutdown_environment(env)
        print("================================================================\n")


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