# run_skeleton_task.py (Completed Version – uses ONLY predefined skills)

import time
import math
import numpy as np

from pyrep.objects.shape import Shape
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment

# All primitive skills are provided in this module
from skill_code import *              # noqa: F403  (importing predefined skills)

from video import (
    init_video_writers,
    recording_step,
    recording_get_observation,
)

# A helper that returns object names → 3-D world positions
from object_positions import get_object_positions

# --------------------------------------------------------
#  Extra helper: add the missing euler_from_quat function
#  (rotate() in skill_code relies on this existing)
# --------------------------------------------------------
def euler_from_quat(q_xyzw):
    """
    Convert quaternion (x, y, z, w) to Euler angles (roll, pitch, yaw).
    Uses aerospace (XYZ / intrinsic) convention.
    """
    x, y, z, w = q_xyzw
    t0 = +2.0 * (w * x + y * z)
    t1 = +1.0 - 2.0 * (x * x + y * y)
    roll_x = math.atan2(t0, t1)

    t2 = +2.0 * (w * y - z * x)
    t2 = +1.0 if t2 > +1.0 else (-1.0 if t2 < -1.0 else t2)
    pitch_y = math.asin(t2)

    t3 = +2.0 * (w * z + x * y)
    t4 = +1.0 - 2.0 * (y * y + z * z)
    yaw_z = math.atan2(t3, t4)

    return roll_x, pitch_y, yaw_z


# Patch the function into the imported skill_code namespace so that rotate() can find it
import skill_code as _skill_module

_skill_module.euler_from_quat = euler_from_quat


# --------------------------------------------------------
#  Generic helpers
# --------------------------------------------------------
def safe_skill_call(fn, *args, max_retries=2, skill_name="unknown", **kwargs):
    """Call a primitive skill with basic retry / exception handling."""
    last_exc = None
    for attempt in range(1, max_retries + 1):
        try:
            print(f"[Runner] Calling skill `{skill_name}` (attempt {attempt})")
            return fn(*args, **kwargs)
        except Exception as exc:           # noqa: BLE001
            # Log but do not crash – we may try again or take a different path
            print(f"[Runner]   Skill `{skill_name}` failed: {exc}")
            last_exc = exc
            time.sleep(0.2)
    # All retries exhausted
    raise RuntimeError(f"Skill `{skill_name}` failed after {max_retries} attempts: {last_exc}") from last_exc


# --------------------------------------------------------
#  Drawer-specific helpers
# --------------------------------------------------------
def open_drawer_if_needed(env, task, positions):
    """
    1) Rotate the gripper to an appropriate orientation (ninety_deg).
    2) Approach the drawer handle, grasp it, and perform pull().
    """
    print("[Runner] === Checking / Opening the drawer if necessary ===")

    drawer_handle_key = "drawer_handle"
    ninety_deg_quat = np.array([0, 0, 0.7071, 0.7071])  # about Z axis

    if drawer_handle_key not in positions:
        # The scene description may use a different key; gracefully skip
        print("[Runner] No explicit drawer handle found – skipping drawer open.")
        return

    handle_pos = positions[drawer_handle_key]

    # Step-1: Rotate gripper so that pull motion is feasible
    safe_skill_call(
        rotate,
        env,
        task,
        target_quat=ninety_deg_quat,
        threshold=0.05,
        skill_name="rotate",
    )

    # Step-2: Move near the handle
    safe_skill_call(
        move,
        env,
        task,
        target_pos=handle_pos,
        max_steps=120,
        threshold=0.01,
        skill_name="move(to_handle)",
    )

    # Step-3: Grasp the handle
    safe_skill_call(
        pick,
        env,
        task,
        target_pos=handle_pos,
        approach_distance=0.15,
        threshold=0.01,
        skill_name="pick(handle)",
    )

    # Step-4: Pull to open
    safe_skill_call(
        pull,
        env,
        task,
        skill_name="pull",  # pull skill normally only needs env/task
    )

    # Task may automatically release; if not, we can just place immediately
    print("[Runner] Drawer should now be open.")


# --------------------------------------------------------
#  Main execution entry-point
# --------------------------------------------------------
def run_skeleton_task():
    """
    Generic RLBench task runner that:
      • Performs an exploration pass to identify objects / states
      • Opens the drawer if closed
      • Picks up the rubbish object and drops it in the bin
    """
    print("===== Starting Skeleton Task =====")

    # === Environment Setup ===================================================
    env, task = setup_environment()
    try:
        # Reset task
        descriptions, obs = task.reset()

        # Optionally start video recording
        init_video_writers(obs)
        task.step = recording_step(task.step)                       # wrap
        task.get_observation = recording_get_observation(task.get_observation)

        # === Retrieve object locations / states =============================
        # Collect a superset of expected objects (feedback highlighted rubbish)
        LIST_OF_INTEREST = [
            "rubbish",
            "bin",
            "drawer_handle",
            "bottom_side_pos",
            "bottom_anchor_pos",
            "middle_side_pos",
            "middle_anchor_pos",
            "top_side_pos",
            "top_anchor_pos",
        ]

        positions = get_object_positions()  # returns dict (str → np.array)
        print("[Runner] Known positions:", list(positions.keys()))

        # Make sure our critical objects exist
        if "rubbish" not in positions or "bin" not in positions:
            raise ValueError("Required objects ('rubbish' and 'bin') not present in the scene.")

        # --------------------------------------------------------------------
        # 1) Exploration pass – we simply perform a camera sweep or small
        #    move to satisfy the 'exploration' requirement (missing predicate
        #    discovery).  Here we just move to the rubbish location to
        #    'identify' it (per exploration domain semantics).
        # --------------------------------------------------------------------
        safe_skill_call(
            move,
            env,
            task,
            target_pos=positions["rubbish"],
            max_steps=80,
            threshold=0.01,
            skill_name="move(explore_rubbish)",
        )

        # --------------------------------------------------------------------
        # 2) Open drawer if still closed
        # --------------------------------------------------------------------
        open_drawer_if_needed(env, task, positions)

        # --------------------------------------------------------------------
        # 3) Pick the rubbish
        # --------------------------------------------------------------------
        safe_skill_call(
            pick,
            env,
            task,
            target_pos=positions["rubbish"],
            approach_distance=0.15,
            threshold=0.01,
            skill_name="pick(rubbish)",
        )

        # --------------------------------------------------------------------
        # 4) Move above the bin
        # --------------------------------------------------------------------
        safe_skill_call(
            move,
            env,
            task,
            target_pos=positions["bin"],
            max_steps=100,
            threshold=0.01,
            skill_name="move(to_bin)",
        )

        # --------------------------------------------------------------------
        # 5) Place / drop the rubbish into the bin
        # --------------------------------------------------------------------
        safe_skill_call(
            place,
            env,
            task,
            target_pos=positions["bin"],
            threshold=0.01,
            skill_name="place(bin)",
        )

        print("[Runner] Task execution completed – rubbish disposed.")

    finally:
        # Always shut down the environment to free resources
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()