# run_skeleton_task.py (Completed Version)

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 *      # DO NOT redefine any primitives
from video import init_video_writers, recording_step, recording_get_observation

from object_positions import get_object_positions


# ----------------------------------------------------------
# Helper utilities (safe to add, they are NOT new skills)
# ----------------------------------------------------------
def quat_from_euler(roll: float, pitch: float, yaw: float) -> np.ndarray:
    """Convert Euler angles (XYZ convention) to quaternion (xyzw)."""
    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)

    w = cr * cp * cy + sr * sp * sy
    x = sr * cp * cy - cr * sp * sy
    y = cr * sp * cy + sr * cp * sy
    z = cr * cp * sy - sr * sp * cy
    return np.array([x, y, z, w], dtype=np.float32)


def run_skeleton_task():
    """Generic skeleton with an added exploration phase for the `rotated` predicate."""
    print("===== Starting Skeleton Task =====")

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

        # Wrap the task for video‐recording
        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)

        # --------------------------------------------------
        # Retrieve any object/pose information needed
        # --------------------------------------------------
        positions = get_object_positions()
        # (Optional) print to see what is available
        print("[Info] Known object positions:", positions)

        # --------------------------------------------------
        # Exploration Phase – identify the missing predicate
        # --------------------------------------------------
        #
        # Feedback told us the missing predicate is `rotated'.
        # We therefore explore the effect of different
        # end-effector orientations on the environment to
        # verify when the predicate is satisfied.  We only use
        # the predefined `rotate` skill here.
        #
        print("\n===== Exploration Phase: Checking `rotated` predicate =====")

        # A small set of commonly used angles.  The keys match
        # the symbolic names that appear in the PDDL, e.g.,
        # `ninety_deg`, `zero_deg`, etc.
        ANGLES_DEG = {
            "zero_deg":    (0.0, 0.0,   0.0),
            "ninety_deg":  (0.0, 0.0,  90.0),
            "minus_90_deg":(0.0, 0.0, -90.0),
            "one_eighty":  (0.0, 0.0, 180.0),
        }

        # Convert to quaternions once
        angle_quats = {name: quat_from_euler(
            math.radians(r), math.radians(p), math.radians(y))
            for name, (r, p, y) in ANGLES_DEG.items()
        }

        # Cycle through each candidate orientation
        for angle_name, quat in angle_quats.items():
            print(f"\n--- Trying orientation `{angle_name}` ---")
            try:
                obs, reward, done = rotate(
                    env,
                    task,
                    target_quat=quat,
                    max_steps=150,
                    threshold=0.03,
                    timeout=8.0
                )
                # If the task ended, exit early
                if done:
                    print("[Exploration] Task unexpectedly finished while rotating.")
                    break

                # After rotation, we simply print the observed EE
                # orientation so the user can confirm the predicate
                current_quat = obs.gripper_pose[3:7]
                print(f"[Exploration] Current gripper quaternion: {current_quat}")

                # Optional: short pause for clarity in video
                for _ in range(10):
                    # keep the arm still for a few steps
                    noop = np.zeros(env.action_shape)
                    noop[-1] = getattr(obs, 'gripper_openness', -1.0)
                    obs, reward, done = task.step(noop)
                    if done:
                        break

            except Exception as exc:
                print(f"[Warning] Rotation `{angle_name}` failed with exception: {exc}")

        print("\n===== Exploration Completed =====")
        print("You can now decide which of the above orientations "+
              "corresponds to the PDDL predicate `(rotated ?g ninety_deg)`.")

        # --------------------------------------------------
        # Main Task Logic (placeholder)
        # --------------------------------------------------
        #
        # Insert oracle plan here once the correct orientation
        # satisfying the `rotated` predicate is confirmed.  The
        # plan will consist of calls to the existing skills
        # (move, pick, place, pull, …).  For now we end the
        # demo after exploration so the script remains safe on
        # environments where the full task is unknown.
        #
        print("\n[Task] No further high-level plan provided.  "+
              "Stopping after exploration phase.")

    finally:
        # --------------------------------------------------
        # Always shut down the environment
        # --------------------------------------------------
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()