import numpy as np
from pyrep.objects.shape import Shape
from pyrep.objects.proximity_sensor import ProximitySensor     # noqa: F401  (kept for completeness)

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


# ------------------------------------------------------------------
# Helper utilities
# ------------------------------------------------------------------
def fetch_position(obj_name, cached_positions):
    """
    Return the (x, y, z) world-position of an object.

    Priority:
      1) use the pre-cached dictionary returned by object_positions.get_object_positions()
      2) fall back to querying the simulator directly via pyrep Shape

    Raises
    ------
    RuntimeError
        If the object cannot be located.
    """
    if obj_name in cached_positions and cached_positions[obj_name] is not None:
        return np.asarray(cached_positions[obj_name], dtype=np.float32)

    try:
        return np.asarray(Shape(obj_name).get_position(), dtype=np.float32)
    except Exception as exc:
        raise RuntimeError(f"[Task] Cannot find a valid position for object “{obj_name}”.") from exc


def alias(name):
    """
    Map the high-level identifiers used in the oracle specification
    to concrete RLBench shape names that exist inside the scene.

    If the supplied name already corresponds to a valid shape, it is
    returned unchanged.
    """
    mapping = {
        # drawer helper points
        "side-pos-bottom":   "bottom_side_pos",
        "anchor-pos-bottom": "bottom_anchor_pos",
        "side-pos-middle":   "middle_side_pos",
        "anchor-pos-middle": "middle_anchor_pos",
        "side-pos-top":      "top_side_pos",
        "anchor-pos-top":    "top_anchor_pos",

        # objects
        "rubbish": "item3",     # treat “item3” as the piece of trash
        "trash":   "item3",
        "bin":     "bin",

        # shorthand that appeared in the formal spec (step-4 uses “bottom”)
        "bottom":  "bottom_anchor_pos",
    }
    return mapping.get(name, name)


# ------------------------------------------------------------------
# Main task runner
# ------------------------------------------------------------------
def run_task_open_drawer_and_dispose_trash():
    """
    Realises the oracle plan required by the formal specification:

        1) rotate  – orient gripper 90° about global Z
        2) move    – go to “side-pos-bottom” (front of bottom drawer)
        3) move    – go to “anchor-pos-bottom” (drawer handle)
        4) pick    – grasp the drawer handle
        5) pull    – pull the drawer open by 0.25 m along +X
        6) pick    – grasp the rubbish lying on the table
        7) place   – drop the rubbish into the bin

    An auxiliary “place” is inserted after pulling to release the
    drawer handle so that the gripper becomes free to pick the trash.
    """
    print("\n===================  TASK START  ===================")

    # --------------------------------------------------
    # Environment initialisation
    # --------------------------------------------------
    env, task = setup_environment()
    try:
        # Reset the RLBench task to its pristine initial state
        descriptions, obs = task.reset()

        # Optional: enable RGB-D recording utilities (will do nothing if
        # the backend module ‘video’ is a stub)
        init_video_writers(obs)
        task.step            = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # Cache frequently-used object positions (cheap dictionary lookup
        # instead of repeated pyrep Shape queries)
        cached_positions = get_object_positions()

        # ==============================================================
        # STEP-1  ROTATE gripper 90° about global Z-axis
        # ==============================================================
        print("\n[Plan] Step-1  rotate gripper 90° about Z-axis")
        from scipy.spatial.transform import Rotation as R
        target_quat = R.from_euler('z', 90, degrees=True).as_quat()   # xyzw order
        rotate(env, task, target_quat)

        # ==============================================================
        # STEP-2  MOVE to “side-pos-bottom”
        # ==============================================================
        side_pos_name = alias("side-pos-bottom")
        side_pos      = fetch_position(side_pos_name, cached_positions)
        print(f"\n[Plan] Step-2  move to {side_pos_name} @ {side_pos}")
        move(env, task, side_pos)

        # ==============================================================
        # STEP-3  MOVE to “anchor-pos-bottom” (drawer handle)
        # ==============================================================
        anchor_pos_name = alias("anchor-pos-bottom")
        anchor_pos      = fetch_position(anchor_pos_name, cached_positions)
        print(f"\n[Plan] Step-3  move to {anchor_pos_name} @ {anchor_pos}")
        move(env, task, anchor_pos)

        # ==============================================================
        # STEP-4  PICK (grasp) the drawer handle
        # ==============================================================
        print(f"\n[Plan] Step-4  pick / grasp drawer handle at {anchor_pos_name}")
        pick(env, task, target_pos=anchor_pos, approach_distance=0.10, approach_axis='y')

        # ==============================================================
        # STEP-5  PULL the drawer 0.25 m along +X
        # ==============================================================
        print("\n[Plan] Step-5  pull drawer by 0.25 m along +X")
        pull(env, task, pull_distance=0.25, pull_axis='x')

        # --------------------------------------------------------------
        # AUX   release the handle → open gripper so we can grab trash
        # --------------------------------------------------------------
        print("\n[Plan] Auxiliary  release drawer handle (open gripper)")
        place(env, task, target_pos=anchor_pos, approach_distance=0.05, approach_axis='y')

        # ==============================================================
        # STEP-6  PICK the rubbish lying on the table
        # ==============================================================
        rubbish_name = alias("rubbish")
        rubbish_pos  = fetch_position(rubbish_name, cached_positions)
        print(f"\n[Plan] Step-6  pick up object “{rubbish_name}” @ {rubbish_pos}")
        pick(env, task, target_pos=rubbish_pos, approach_distance=0.15, approach_axis='z')

        # ==============================================================
        # STEP-7  PLACE the rubbish into the bin
        # ==============================================================
        bin_pos = fetch_position(alias("bin"), cached_positions)
        print(f"\n[Plan] Step-7  place rubbish into bin @ {bin_pos}")
        place(env, task, target_pos=bin_pos, approach_distance=0.15, approach_axis='z')

        print("\n===================  TASK SUCCESS  ===================")

    finally:
        # Always shut down the simulator, even if an error occurs
        shutdown_environment(env)


# ------------------------------------------------------------------
# Stand-alone execution
# ------------------------------------------------------------------
if __name__ == "__main__":
    run_task_open_drawer_and_dispose_trash()