# run_skeleton_task.py  (Completed Executable Code)

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

# NOTE: every primitive / low-level skill is **already** implemented in skill_code.
#       We simply import * and **call** them when needed.
from skill_code import *          # noqa: F401, F403  (keep full wildcard – we do not redefine skill fns)

from video import (
    init_video_writers,
    recording_step,
    recording_get_observation,
)

# Utility that the framework provides for querying object poses, etc.
from object_positions import get_object_positions


# ----------------------------------------------------------
#  Helper utilities (NOT new “skills”, just math wrappers)
# ----------------------------------------------------------
def quat_from_euler(roll: float, pitch: float, yaw: float):
    """
    Convert Euler XYZ (rad) to quaternion (xyzw).

    We keep everything in numpy float32 for consistency with RLBench.
    """
    cy = math.cos(yaw * 0.5)
    sy = math.sin(yaw * 0.5)
    cp = math.cos(pitch * 0.5)
    sp = math.sin(pitch * 0.5)
    cr = math.cos(roll * 0.5)
    sr = math.sin(roll * 0.5)

    q_w = cr * cp * cy + sr * sp * sy
    q_x = sr * cp * cy - cr * sp * sy
    q_y = cr * sp * cy + sr * cp * sy
    q_z = cr * cp * sy - sr * sp * cy
    return np.array([q_x, q_y, q_z, q_w], dtype=np.float32)


# ----------------------------------------------------------
#  Main entry-point
# ----------------------------------------------------------
def run_skeleton_task():
    """
    Generic task runner.

    The *feedback* we received told us that the “rotated” predicate was
    previously missing in our execution traces, so here we explicitly add an
    exploration phase that **rotates** the gripper to a variety of target
    orientations.  This guarantees that the predicate becomes satisfied for any
    downstream high-level plan that depends on it (e.g., moving the gripper to
    the side of a drawer before pulling).
    """
    print("\n=================  Starting Skeleton Task  =================")

    # ----------------------------------------------------------
    #  ENVIRONMENT SETUP
    # ----------------------------------------------------------
    env, task = setup_environment()
    try:
        # Reset to obtain the first observation
        descriptions, obs = task.reset()
        print("[Init] Task descriptions:", descriptions)

        # (Optional) initialise optional video capture
        init_video_writers(obs)

        # Wrap task.step / task.get_observation with recorders
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # ------------------------------------------------------
        #  Retrieve relevant object positions (handles, etc.)
        # ------------------------------------------------------
        object_positions = get_object_positions()
        print("[Init] Known object positions:", object_positions)

        # We attempt to find a plausible drawer or handle.  The exact key names
        # depend on the specific RLBench scenario, so we fall back gracefully
        # if they are not present.
        drawer_handle_key = None
        for key in object_positions.keys():
            if "handle" in key.lower() or "drawer" in key.lower():
                drawer_handle_key = key
                break

        if drawer_handle_key is None:
            print(
                "[Warning] Could not automatically locate a 'drawer handle' "
                "in object_positions – proceeding with purely exploratory "
                "rotation to satisfy the (rotated ?) predicate."
            )

        # ------------------------------------------------------
        #  (1)  EXPLORATION PHASE – satisfy “rotated” predicate
        # ------------------------------------------------------
        # Rotate to several canonical orientations:  0°,  90°,  −90°
        # around the tool-Z axis.  We keep roll / pitch to zero so that we
        # do not disrupt the gripper’s approach vector.
        exploration_targets = [
            quat_from_euler(0.0, 0.0, 0.0),                    # identity
            quat_from_euler(0.0, 0.0,  math.pi * 0.5),          # +90°
            quat_from_euler(0.0, 0.0, -math.pi * 0.5),          # -90°
        ]

        rotation_success = False
        for idx, target_quat in enumerate(exploration_targets):
            print(f"\n[Exploration] Trial {idx+1}: rotate gripper "
                  f"to yaw={round(math.degrees(math.atan2(target_quat[2], target_quat[3])), 2)}°")
            try:
                obs, reward, done = rotate(
                    env=env,
                    task=task,
                    target_quat=target_quat,
                    max_steps=120,
                    threshold=0.03,
                    timeout=8.0,
                )
                rotation_success = True
                if done:
                    print("[Exploration] Task signalled completion inside rotation loop.")
                    return
            except Exception as e:
                print(f"[Exploration] Rotation attempt failed due to: {e}")

        if not rotation_success:
            print("[Error] Could not complete any rotation – the predicate "
                  "'rotated' may still be unsatisfied. Continuing anyway.")

        # ------------------------------------------------------
        #  (2)  SIMPLE PLAN  (example: open a drawer if one exists)
        # ------------------------------------------------------
        # Outline:
        #   a)  Approach & pick the drawer handle (if we found one)
        #   b)  Pull / place actions as necessary
        #
        # We keep this HIGH-LEVEL and robust to missing keys / skills.  All
        # individual low-level implementations live in skill_code so we focus
        # only on flow control here.

        if drawer_handle_key is None:
            print(
                "[Plan] No drawer handle detected – skipping pick/pull plan "
                "and finalising task."
            )
        else:
            handle_pos = object_positions[drawer_handle_key]
            print(f"[Plan] Preparing to pick drawer handle at: {handle_pos}")

            # --------------------------------------------------
            #  (a) pick  (signature may vary; handle gracefully)
            # --------------------------------------------------
            try:
                obs, reward, done = pick(
                    env=env,
                    task=task,
                    target_pos=np.asarray(handle_pos, dtype=np.float32),
                    approach_distance=0.15,
                    max_steps=150,
                    threshold=0.01,
                    approach_axis="z",
                    timeout=10.0,
                )
                if done:
                    print("[Plan] Task ended during pick – exiting.")
                    return
            except TypeError:
                # Signature mismatch; try alternative minimal calling
                try:
                    obs, reward, done = pick(env, task, np.asarray(handle_pos))
                except Exception as e:
                    print(f"[Plan] Unable to execute 'pick' for drawer handle: {e}")
                    done = False

            # --------------------------------------------------
            #  (b) pull  (open the drawer)
            # --------------------------------------------------
            try:
                obs, reward, done = pull(env, task)
                if done:
                    print("[Plan] Drawer pull completed – task done.")
                    return
            except Exception as e:
                print(f"[Plan] 'pull' skill failed or is not applicable: {e}")

            # --------------------------------------------------
            #  (c) optional: place / cleanup
            # --------------------------------------------------
            # We release anything we might be holding to restore initial state.
            try:
                obs, reward, done = place(env, task)
            except Exception:
                pass

        # ------------------------------------------------------
        #  Task finished (either success or end of scripted steps)
        # ------------------------------------------------------
        print("\n=================  End of Skeleton Task  =================")

    finally:
        # Always guarantee proper shutdown to free CoppeliaSim resources
        shutdown_environment(env)


# Entry-point guard
if __name__ == "__main__":
    run_skeleton_task()