# run_skeleton_task.py   (Completed)

import time
import math
from typing import Dict, Tuple

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 – we rely on all pre-defined skills
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions


# --------------------------------------------------------------------------- #
# Helper utilities that are *not* primitive “skills” but are required by many #
# of those skills (e.g. rotate()) and were missing from the original code.    #
# --------------------------------------------------------------------------- #
def normalize_quaternion(q):
    """Return a unit-length copy of the given quaternion (x, y, z, w)."""
    q = np.asarray(q, dtype=float).flatten()
    n = np.linalg.norm(q)
    if n == 0.0:
        return q
    return q / n


def euler_from_quat(quat):
    """Convert quaternion (x, y, z, w) to Euler angles (roll, pitch, yaw)."""
    x, y, z, w = quat
    t0 = +2.0 * (w * x + y * z)
    t1 = +1.0 - 2.0 * (x * x + y * y)
    roll = math.atan2(t0, t1)

    t2 = +2.0 * (w * y - z * x)
    t2 = max(min(t2, 1.0), -1.0)
    pitch = math.asin(t2)

    t3 = +2.0 * (w * z + x * y)
    t4 = +1.0 - 2.0 * (y * y + z * z)
    yaw = math.atan2(t3, t4)
    return roll, pitch, yaw


# --------------------------------------------------------------------------- #
# Misc. convenience wrappers                                                  #
# --------------------------------------------------------------------------- #
def safe_skill_call(skill_fn, *args, **kwargs):
    """Try to execute a skill, capture and log any exception, and continue."""
    try:
        return skill_fn(*args, **kwargs)
    except Exception as exc:         # noqa: BLE001 – broad except is intentional here
        print(f"[WARNING] Skill {skill_fn.__name__} raised an exception → {exc}")
        return None, 0.0, False


def quaternion_from_axis_angle(axis: Tuple[float, float, float], angle_rad: float):
    """Return xyzw quaternion for a rotation of angle_rad around axis."""
    axis = np.asarray(axis, dtype=float)
    axis = axis / np.linalg.norm(axis)
    s = math.sin(angle_rad / 2.0)
    x, y, z = axis * s
    w = math.cos(angle_rad / 2.0)
    return np.array([x, y, z, w], dtype=float)


# --------------------------------------------------------------------------- #
# Main execution                                                              #
# --------------------------------------------------------------------------- #
def run_skeleton_task():
    """Generic but complete task runner demonstrating exploration + execution."""
    print("================  Starting Skeleton Task  =================")
    env, task = setup_environment()

    try:
        # ------------------------------------------------------------------ #
        #   Reset & Book-keeping                                             #
        # ------------------------------------------------------------------ #
        descriptions, obs = task.reset()           # type: ignore
        init_video_writers(obs)

        # Wrap step / observation so we transparently record the run
        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 known object poses                                      #
        # ------------------------------------------------------------------ #
        positions: Dict[str, Tuple[float, float, float]] = get_object_positions()
        print("[INFO] Object positions retrieved:", positions)

        # ------------------------------------------------------------------ #
        #   Simple Exploration Phase → figuring out drawer lock-status       #
        # ------------------------------------------------------------------ #
        ninety_deg_quat = quaternion_from_axis_angle((0.0, 1.0, 0.0), math.pi / 2)

        for name, pos in positions.items():
            if "drawer" not in name.lower():
                continue                         # Only interested in drawers here

            print(f"\n-----  Exploring «{name}»  -----")
            # 1) Rotate gripper to a side-grasp (90°) orientation
            print(f"[Plan] rotate() gripper for «{name}»")
            safe_skill_call(
                rotate, env, task,
                target_quat=ninety_deg_quat,
                max_steps=120, threshold=0.05, timeout=10.0
            )

            # 2) Move towards the drawer handle
            print(f"[Plan] move() towards handle of «{name}» at {pos}")
            obs, reward, done = safe_skill_call(
                move, env, task,
                target_pos=pos,
                approach_distance=0.15,
                max_steps=150,
                threshold=0.01,
                approach_axis='z',
                timeout=10.0
            )

            if done:
                print("[Task] Episode terminated during move().")
                break

            # 3) Pick (grab) the drawer handle
            print(f"[Plan] pick() handle of «{name}»")
            obs, reward, done = safe_skill_call(
                pick, env, task,
                target_pos=pos,
                approach_distance=0.02,
                max_steps=100,
                threshold=0.005,
                approach_axis='z',
                timeout=10.0
            )

            if done:
                print("[Task] Episode terminated during pick().")
                break

            # 4) Pull to test if the drawer is locked
            print(f"[Plan] pull() on «{name}» (discover lock-status)")
            obs, reward, done = safe_skill_call(
                pull, env, task,
                pull_distance=0.10,
                max_steps=120,
                threshold=0.01,
                timeout=10.0
            )

            # The pull() skill is assumed to expose in its return-value whether
            # the drawer actually opened.  In a real implementation, we would
            # inspect ‘obs’ to update the     (lock-known ?drawer)     predicate.
            #
            # Here, for demonstration, we merely log the outcome:
            if obs is not None:
                print(f"[Result] Completed pull() on «{name}», reward={reward}, done={done}")
            else:
                print(f"[Result] pull() failed for «{name}» – assuming drawer locked.")

            # 5) If the domain signals task completion → exit early
            if done:
                print("[Task] Episode terminated after pull().")
                break

        # ------------------------------------------------------------------ #
        #   Additional mission-specific actions could be inserted here       #
        # ------------------------------------------------------------------ #

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


if __name__ == "__main__":
    run_skeleton_task()