# --------------------------------------------------------------------------- #
#  run_combined_task.py                                                       #
#                                                                             #
#  Concrete implementation that exactly follows the oracle plan:              #
#    1) rotate  →  2) move-to-side  →  3) move-to-anchor  →                   #
#    4) grasp & pull the drawer  →  5) release handle  →                      #
#    6-9) transfer two tomatoes onto the plate.                               #
#                                                                             #
#  Only the predefined low-level skills contained in `skill_code` are used.   #
# --------------------------------------------------------------------------- #

import numpy as np
from scipy.spatial.transform import Rotation as R
from pyrep.objects.shape import Shape                 # kept for completeness
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment

# import all predefined low-level skills (rotate, move, pick, place, pull)
from skill_code import *            # noqa: F403,F401  (wild-card – primitives only)

from video import (
    init_video_writers,
    recording_step,
    recording_get_observation,
)

from object_positions import get_object_positions


# --------------------------------------------------------------------------- #
#                               Helper utilities                              #
# --------------------------------------------------------------------------- #
def safe_get_pos(pos_dict, key, z_offset=0.0):
    """Return a copy of the requested position with an optional z-offset."""
    if key not in pos_dict:
        raise KeyError(f"[Task] Missing position for object “{key}”.")
    p = np.array(pos_dict[key], dtype=np.float32).copy()
    p[2] += z_offset
    return p


def quat_about_z(rad):
    """Return (x,y,z,w) quaternion representing a pure Z-axis rotation."""
    return R.from_euler("z", rad).as_quat().astype(np.float32)


# --------------------------------------------------------------------------- #
#                               Main Task Logic                               #
# --------------------------------------------------------------------------- #
def run_combined_task() -> None:
    print("\n==========  Combined Drawer-&-Tomato Task  ==========")

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

        # optional: activate video recording
        init_video_writers(obs)
        task.step            = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # ---------------  Query all relevant positions  --------------- #
        pos = get_object_positions()

        # drawer landmarks (slightly lifted to avoid collisions)
        side_pos   = safe_get_pos(pos, "bottom_side_pos",   0.04)
        anchor_pos = safe_get_pos(pos, "bottom_anchor_pos", 0.04)

        # tomatoes and plate (lifted a little to hover)
        tomato1 = safe_get_pos(pos, "item1",  0.04)
        tomato2 = safe_get_pos(pos, "item2",  0.04)
        plate   = safe_get_pos(pos, "plate",  0.06)

        # keep track of task termination flags
        done   = False
        reward = 0.0

        # ------------------------------------------------------------------ #
        #                           Oracle Plan                              #
        # ------------------------------------------------------------------ #

        # Step-1: rotate gripper 90° about +Z
        target_quat = quat_about_z(np.pi / 2.0)            # 90 deg
        obs, reward, done = rotate(env, task, target_quat)  # noqa: F405
        if done:
            return

        # Step-2: move to drawer side position
        obs, reward, done = move(env, task, side_pos)       # noqa: F405
        if done:
            return

        # Step-3: move to drawer anchor (handle) position
        obs, reward, done = move(env, task, anchor_pos)     # noqa: F405
        if done:
            return

        # Step-4: grasp the drawer handle (simulate pick-drawer)
        #         → we approach a little from +Y direction and close gripper
        obs, reward, done = pick(                           # noqa: F405
            env,
            task,
            target_pos=anchor_pos,
            approach_distance=0.05,
            approach_axis="-y"
        )
        if done:
            return

        # Step-5: pull drawer straight along +X (≈15 cm)
        obs, reward, done = pull(                           # noqa: F405
            env,
            task,
            pull_distance=0.15,
            pull_axis="x"
        )
        if done:
            return

        # --------------------  release the handle  -------------------- #
        # open gripper in-place so we can manipulate tomatoes
        curr_p = task.get_observation().gripper_pose[:3]
        obs, reward, done = place(                          # noqa: F405
            env,
            task,
            target_pos=curr_p,
            approach_distance=0.02,
            approach_axis="z"
        )
        if done:
            return

        # --------------------  Tomato transfer loop  ------------------ #
        hover = 0.10      # hover 10 cm above target

        for tomato_pos in [tomato1, tomato2]:
            # approach tomato from above, then pick
            obs, reward, done = move(env, task, tomato_pos + np.array([0, 0, hover]))  # noqa: F405
            if done:
                return
            obs, reward, done = pick(                        # noqa: F405
                env,
                task,
                target_pos=tomato_pos,
                approach_distance=hover,
                approach_axis="-z"
            )
            if done:
                return

            # bring tomato above plate and place it
            obs, reward, done = move(env, task, plate + np.array([0, 0, hover]))  # noqa: F405
            if done:
                return
            obs, reward, done = place(                       # noqa: F405
                env,
                task,
                target_pos=plate,
                approach_distance=hover,
                approach_axis="-z"
            )
            if done:
                return

        print("\n[Task] Success – Drawer opened and both tomatoes placed on plate.")

    finally:
        shutdown_environment(env)
        print("==========  End of Task  ==========\n")


# --------------------------------------------------------------------------- #
#                              Script Entrypoint                              #
# --------------------------------------------------------------------------- #
if __name__ == "__main__":
    run_combined_task()