# run_skeleton_task.py (Completed Executable Version)

import numpy as np
import time
from math import atan2, asin, sqrt

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


# --------------------------------------------------------------------------
# Helper utilities that were missing in the original skill code (rotate etc.)
# --------------------------------------------------------------------------
def normalize_quaternion(q):
    """Return a unit-length quaternion as np.array[4]."""
    q = np.asarray(q, dtype=np.float64)
    norm = np.linalg.norm(q)
    if norm == 0:
        return q
    return q / norm


def euler_from_quat(q):
    """
    Convert a quaternion (xyzw) into Euler angles (roll, pitch, yaw).
    Returns angles in radians as np.array[3].
    """
    x, y, z, w = q
    # roll (x-axis rotation)
    sinr_cosp = 2.0 * (w * x + y * z)
    cosr_cosp = 1.0 - 2.0 * (x * x + y * y)
    roll = atan2(sinr_cosp, cosr_cosp)

    # pitch (y-axis rotation)
    sinp = 2.0 * (w * y - z * x)
    if abs(sinp) >= 1:
        pitch = np.pi / 2 * np.sign(sinp)      # use 90° if out of range
    else:
        pitch = 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 = atan2(siny_cosp, cosy_cosp)

    return np.array([roll, pitch, yaw], dtype=np.float64)


def quat_from_axis_angle(axis, angle):
    """
    Build quaternion (xyzw) from axis (3,) and angle (rad).
    """
    axis = np.asarray(axis, dtype=np.float64)
    axis = axis / np.linalg.norm(axis)
    s = np.sin(angle / 2.0)
    x, y, z = axis * s
    w = np.cos(angle / 2.0)
    return np.array([x, y, z, w], dtype=np.float64)


# --------------------------------------------------------------------------
# Main task execution
# --------------------------------------------------------------------------
def run_skeleton_task():
    """Generic skeleton for running any task in your simulation."""
    print("===== Starting Skeleton Task =====")
    env, task = setup_environment()

    # Simple container for whether the environment produced 'done'
    done = False

    try:
        descriptions, obs = task.reset()
        init_video_writers(obs)

        # ---------------- Video wrapping ----------------
        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)

        # ---------------- Exploration phase ----------------
        # The exploration requirement is abstract; here we just output
        # which predicates appear missing by comparing the exploration
        # domain predicates list with the combined-domain predicates.
        exploration_predicates = {
            'robot-at', 'at', 'identified', 'temperature-known',
            'holding', 'handempty', 'weight-known',
            'durability-known'
        }
        combined_domain_predicates = {
            'at', 'holding', 'handempty', 'is-locked', 'is-open',
            'rotated', 'gripper-at', 'holding-drawer', 'is-side-pos',
            'is-anchor-pos'
        }
        missing_predicates = exploration_predicates - combined_domain_predicates
        print("[Exploration] The following predicates are not in the primary domain and might be missing:",
              missing_predicates)

        # ---------------- Retrieve object positions ----------------
        try:
            positions = get_object_positions()
            print(f"[Info] Object positions received: {list(positions.keys())}")
        except Exception as e:
            print(f"[Warning] Unable to fetch object positions: {e}")
            positions = {}

        # ---------------- Example single-skill plan ----------------
        # As a demonstrative oracle plan we rotate the gripper 90°
        # around the Z-axis.  In real tasks, multiple skills would be
        # sequenced according to PDDL planning output.
        target_quat = quat_from_axis_angle([0, 0, 1], np.pi / 2)   # 90° about Z
        print("[Plan] Executing rotate skill to turn gripper 90° about Z axis.")
        obs, reward, done = rotate(
            env,
            task,
            target_quat=target_quat,
            max_steps=100,
            threshold=0.03,
            timeout=10.0
        )

        if done:
            print("[Plan] Task reported completion after rotate.")

        # --------------- Placeholder for further plan steps ----------
        # You can insert additional pick / place / move / pull calls
        # here when integrating with a planner that produces an action
        # sequence.  Only predefined skills from `skill_code` should be
        # invoked, e.g.:
        #
        #   obs, reward, done = pick(env, task, ...)
        #   if done: return
        #
        # For this skeleton we keep it minimal.

    except Exception as exc:
        print(f"[Error] Exception during task execution: {exc}")

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


if __name__ == "__main__":
    run_skeleton_task()
