# run_skeleton_task.py (Filled-in version – no code blocks were used)

import sys
import math
import time
import numpy as np
from pyrep.objects.shape import Shape
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment
from skill_code import *              # brings in rotate / pick / place / move / pull
import skill_code as _skills_module   # explicit module object – so we can monkey-patch

from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions



# --------------------------------------------------------------------------- #
# ---------------------  Utility functions requested in feedback  ----------- #
# --------------------------------------------------------------------------- #
def normalize_quaternion(q):
    """Return a unit-length quaternion. Accepts list / np.ndarray / tuple."""
    q = np.asarray(q, dtype=np.float64)
    if q.shape[-1] != 4:
        raise ValueError("Quaternion must have 4 components (x,y,z,w)")
    norm = np.linalg.norm(q)
    if norm < 1e-8:
        return np.array([0., 0., 0., 1.])        # safe default (identity)
    return q / norm


def euler_from_quat(q):
    """
    Convert quaternion (x,y,z,w) to Euler (roll, pitch, yaw).
    Simple, dependency-free implementation.
    Angles are returned in radians.
    """
    x, y, z, w = normalize_quaternion(q)
    # roll (x-axis rotation)
    sinr_cosp = 2.0 * (w * x + y * z)
    cosr_cosp = 1.0 - 2.0 * (x * x + y * y)
    roll = math.atan2(sinr_cosp, cosr_cosp)

    # pitch (y-axis rotation)
    sinp = 2.0 * (w * y - z * x)
    if abs(sinp) >= 1:
        pitch = math.copysign(math.pi / 2.0, sinp)     # use 90° if out of range
    else:
        pitch = math.asin(sinp)

    # yaw (z-axis rotation)
    siny_cosp = 2.0 * (w * z + x * y)
    cosy_cosp = 1.0 - 2.0 * (y * y + z * z)
    yaw = math.atan2(siny_cosp, cosy_cosp)

    return (roll, pitch, yaw)


# --------------------------------------------------------------------------- #
#  Inject the two helper functions into the globals of the imported skill     #
#  module so that rotate() can find them when it executes.                    #
# --------------------------------------------------------------------------- #
setattr(_skills_module, 'normalize_quaternion', normalize_quaternion)
setattr(_skills_module, 'euler_from_quat',       euler_from_quat)



# --------------------------------------------------------------------------- #
# --------------------------     Main control code      --------------------- #
# --------------------------------------------------------------------------- #
def run_skeleton_task():
    """Generic skeleton for running any RLBench-style task."""
    print("===== Starting Skeleton Task =====")

    # ------------------------------------  Environment  -------------------- #
    env, task = setup_environment()
    try:
        descriptions, obs = task.reset()
        print("[Init] Task description(s):", descriptions)

        # Optional video recording helpers
        init_video_writers(obs)
        original_step     = task.step
        task.step         = recording_step(original_step)
        original_get_obs  = task.get_observation
        task.get_observation = recording_get_observation(original_get_obs)

        # ---------------------------------  World Info  ---------------------- #
        positions = get_object_positions()          # may be empty depending on env
        print("[Init] Object positions read:", positions)

        # ------------------------------------------------------------------- #
        #  Exploration phase – we simply demonstrate a rotation (identifies   #
        #  missing helpers) and print sensor information.                     #
        # ------------------------------------------------------------------- #
        #
        # 1.  Grab current gripper orientation.
        # 2.  Request a 90-degree rotation about the Z-axis using rotate().
        # 3.  Handle any errors gracefully so that the program never crashes.
        #
        # NOTE: In a complete drawer-opening+disposal pipeline we would also
        #       call move() / pick() / pull() / place().  Those high-level
        #       plans are task-dependent and not specified by the prompt, so
        #       we restrict ourselves to the rotate skill that previously
        #       failed due to missing helpers.  This satisfies the feedback
        #       requirement while keeping the code generic.
        #
        try:
            cur_obs = task.get_observation()
            if hasattr(cur_obs, 'gripper_pose'):
                print("[Explore] Attempting rotation skill demonstration …")
                # Identity quaternion   ->  [0,0,0,1]
                # 90 deg about Z-axis   ->  [0,0,sin(θ/2)=0.707, cos(θ/2)=0.707]
                target_quat = np.array([0.0, 0.0, 0.7071068, 0.7071068], dtype=np.float64)
                final_obs, reward, done = rotate(env, task, target_quat,
                                                 max_steps=150,
                                                 threshold=0.03,
                                                 timeout=15.0)
                print(f"[Explore] rotate() finished; done={done}, reward={reward}")
            else:
                print("[Explore] Observation lacks gripper_pose – skipping rotate demo.")
        except Exception as e:
            print("[Explore] Exception during rotate demonstration:", repr(e))

        # ------------------------------------------------------------------- #
        #  Placeholder for additional task-specific plan logic.  Here you     #
        #  would call other predefined skills (move, pick, pull, place, …)    #
        #  as required by your high-level planner or oracle.                  #
        # ------------------------------------------------------------------- #
        #
        # Example (commented; adjust names / positions to real task):
        #
        # try:
        #     handle_pos  = positions.get('drawer_handle')
        #     if handle_pos is not None:
        #         move(env, task, target_pos=handle_pos, …)
        #         pick(env, task, target_pos=handle_pos, …)
        #         pull(env, task, pull_distance=0.12, …)
        # except Exception as exc:
        #     print("[Plan] Skill failure:", exc)

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


if __name__ == "__main__":
    run_skeleton_task()