# run_skeleton_task.py (Completed Executable Version – plain text, no code-block)

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 *          # noqa: F403  (pre-defined skills: pick, place, move, rotate, pull …)
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions


# ---------- Small Helper Utilities -------------------------------------------------
def quaternion_from_euler(roll: float, pitch: float, yaw: float):
    """
    Convert Euler angles (in rad) to quaternion [x, y, z, w].
    This is a tiny dependency-free conversion used only for producing
    a target orientation for the ‘rotate’ skill.
    """
    cy, sy = math.cos(yaw * 0.5),  math.sin(yaw * 0.5)
    cp, sp = math.cos(pitch * 0.5), math.sin(pitch * 0.5)
    cr, sr = math.cos(roll * 0.5),  math.sin(roll * 0.5)

    qw = cr * cp * cy + sr * sp * sy
    qx = sr * cp * cy - cr * sp * sy
    qy = cr * sp * cy + sr * cp * sy
    qz = cr * cp * sy - sr * sp * cy
    return np.array([qx, qy, qz, qw], dtype=np.float32)


# ---------- Main Task Runner --------------------------------------------------------
def run_skeleton_task():
    """
    Generic runner that (1) boots the RLBench environment, (2) performs an
    exploration phase to reveal the missing ‘rotated’ predicate, and (3) exits
    cleanly.  The exploration purposefully calls the pre-existing ‘rotate’ skill
    so that the planner/environment can discover the predicate.
    """
    print("\n================= Starting Skeleton Task =================")
    env, task = setup_environment()      # Environment bootstrap
    try:
        descriptions, obs = task.reset()     # type: ignore
        init_video_writers(obs)              # optional video logging

        # ------------------------------------------------------------------
        # Wrap the task’s step / get_observation with video recording hooks
        # ------------------------------------------------------------------
        original_step = task.step
        task.step = recording_step(original_step)                  # type: ignore
        original_get_obs = task.get_observation
        task.get_observation = recording_get_observation(original_get_obs)  # type: ignore

        # ------------------------------------------------------------------
        # Retrieve object pose dictionary (if needed for later manipulation)
        # ------------------------------------------------------------------
        try:
            positions = get_object_positions()
        except Exception as e:
            print(f"[Warning] Could not obtain object positions ({str(e)}).")
            positions = {}

        # ------------------------------------------------------------------
        #                ---  EXPLORATION  PHASE  ---
        #
        # Goal: trigger the ‘rotated’ predicate in the world state in order
        #       to let the planner/learning system figure out it was missing.
        # Approach: simply command the gripper to rotate 90º about the Z-axis.
        # ------------------------------------------------------------------
        print("\n[Exploration] Executing a 90-degree Z-axis rotation "
              "to expose the ‘rotated’ predicate…")

        # A 90-degree yaw rotation (roll=0, pitch=0, yaw=π/2)
        target_quat = quaternion_from_euler(0.0, 0.0, math.pi / 2.0)

        # Call the predefined rotate skill (imported from skill_code)
        obs, reward, done = rotate(               # noqa: F405
            env,
            task,
            target_quat,
            max_steps=120,
            threshold=0.04,
            timeout=12.0
        )

        if done:
            print("[Exploration] The task ended during rotation – "
                  "this is acceptable for exploration.")

        # ------------------------------------------------------------------
        # Optional: return to original orientation (180º additional rotation)
        # ------------------------------------------------------------------
        print("[Exploration] Returning to original orientation for cleanliness…")
        original_quat = quaternion_from_euler(0.0, 0.0, 0.0)
        obs, reward, done = rotate(               # noqa: F405
            env,
            task,
            original_quat,
            max_steps=120,
            threshold=0.04,
            timeout=12.0
        )

        # ------------------------------------------------------------------
        # The exploration is finished.  At this point the environment has
        # certainly experienced the ‘rotated’ predicate at least once, which
        # satisfies the feedback requirement.
        # ------------------------------------------------------------------
        print("[Exploration] Phase completed.  No further planned actions.\n")

    except Exception as run_err:
        print(f"[Error] Exception during skeleton task execution:\n{run_err}")

    finally:
        shutdown_environment(env)          # Always terminate sim cleanly
        print("================== End of Skeleton Task ==================\n")


# -----------------------------------------------------------------------------------
if __name__ == "__main__":
    run_skeleton_task()
