# run_skeleton_task.py (Completed Generic Skeleton – Exploration + Task Plan)

import math
import time
import traceback
from typing import Dict, Tuple

import numpy as np

from env import setup_environment, shutdown_environment
from skill_code import *                 # noqa – we must use the predefined skills exactly as shipped
from video import (                      # helper utilities supplied by the framework
    init_video_writers,
    recording_step,
    recording_get_observation,
)
from object_positions import get_object_positions


# --------------------------------------------------------------------------- #
# Helper utilities (pure-python, do NOT redefine any robot “skills”)          #
# --------------------------------------------------------------------------- #
def euler_to_quat(roll: float, pitch: float, yaw: float) -> np.ndarray:
    """
    Convert intrinsic XYZ (roll-pitch-yaw) Euler angles into a quaternion
    in the (x, y, z, w) format that RLBench expects.

    The math follows the standard aerospace sequence.
    """
    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.asarray([qx, qy, qz, qw], dtype=np.float32)


def find_first_drawer(positions: Dict[str, Tuple[float, float, float]]) -> str:
    """
    Simple heuristic to choose a drawer-like object.  We just pick the first
    object whose name contains 'drawer'.  Returns None if nothing matches.
    """
    for name in positions:
        if "drawer" in name.lower():
            return name
    return None


def print_exception(e: Exception) -> None:
    """Light-weight helper that prints a traceback without killing the program."""
    print("-------------- Exception --------------")
    traceback.print_exc()
    print("---------------------------------------")


# --------------------------------------------------------------------------- #
# Main routine                                                                #
# --------------------------------------------------------------------------- #
def run_skeleton_task() -> None:
    """
    Generic RLBench task runner.

    • Sets up the simulator
    • Carries out an automated “exploration” phase to discover whether the
      environment exposes a “lock-known” predicate (this is only an
      informational step; we do not modify any skills)
    • Executes a minimal plan that (a) approaches the first detected drawer,
      (b) rotates the gripper by 90 deg (aligning with combined-domain.pddl’s
      'rotate → ninety_deg'), and (c) attempts to pull the drawer open.
    """
    print("===== Starting Skeleton Task =====")

    # === Environment Setup =================================================== #
    env, task = setup_environment()
    try:
        # Reset the task to its initial state
        descriptions, obs = task.reset()

        # (Optional) initialise video writers
        init_video_writers(obs)

        # Wrap task.step / task.get_observation so that video frames are dumped
        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 Object Positions (one-shot query) ====================== #
        # The helper returns a dict mapping object-names to Cartesian positions
        positions: Dict[str, Tuple[float, float, float]] = get_object_positions()
        if not positions:
            print("[Warning] get_object_positions() returned an empty dict.")

        # ------------------------------------------------------------------ #
        # 1) Exploration phase: try a dummy pull on each object to infer
        #    whether a *lock-known* predicate exists and is required.
        # ------------------------------------------------------------------ #
        # Our heuristic:
        #     • Attempt pull() on every candidate handle once.
        #     • If a RuntimeError/Assertion complaining about an unmet
        #       precondition “… lock-known …” surfaces, we infer that
        #       the missing predicate is indeed 'lock-known'.
        #
        # NOTE: This is *diagnostic* only – we don’t attempt to patch the
        #       low-level PDDL, we just print the conclusion.
        #
        missing_predicate_inferred = None
        for obj_name in positions.keys():
            try:
                # We intentionally call pull() with a nonsense location; the
                # skill is expected to raise if preconditions are violated.
                print(f"[Exploration] Probing object '{obj_name}' with pull()")
                pull(env, task, obj_name)         # signature is assumed
            except Exception as e:
                # Quick pattern search
                msg = str(e).lower()
                if "lock-known" in msg or "lock_known" in msg:
                    missing_predicate_inferred = "lock-known"
                    print("[Exploration] Detected reference to missing "
                          "'lock-known' predicate in exception message.")
                    break
                # Print but continue
                print_exception(e)

        if missing_predicate_inferred is None:
            print("[Exploration] No evidence of a missing 'lock-known' "
                  "predicate was found (this is OK if the domain already "
                  "includes it).")
        else:
            print(f"[Exploration] => Missing predicate probably: "
                  f"'{missing_predicate_inferred}'")

        # ------------------------------------------------------------------ #
        # 2) Task execution: open the first drawer we can find                #
        # ------------------------------------------------------------------ #
        drawer_name = find_first_drawer(positions)
        if drawer_name is None:
            print("[Task] No drawer-like object found – nothing to do.")
            return

        drawer_pos = positions[drawer_name]
        print(f"[Task] Target drawer = '{drawer_name}', "
              f"approx. position {drawer_pos}")

        # ----- 2-a) Approach the drawer ----------------------------------- #
        try:
            # We assume move() moves the end-effector to the given position
            print(f"[Task] Moving towards drawer '{drawer_name}' …")
            move(
                env,
                task,
                target_pos=np.asarray(drawer_pos, dtype=np.float32),
                approach_distance=0.12,
                max_steps=150,
                threshold=0.015,
                approach_axis="z",
                timeout=10.0,
            )
        except Exception as e:
            print_exception(e)
            print("[Task] move() failed; continuing anyway (may already be "
                  "close enough).")

        # ----- 2-b) Rotate the gripper by 90 deg about the Z-axis ---------- #
        # PDDL 'rotate' action uses the enumerated constant 'ninety_deg'.
        # Here we convert that into a quaternion (roll=0, pitch=0, yaw=+90°).
        print("[Task] Rotating gripper to 'ninety_deg' orientation …")
        target_quat = euler_to_quat(0.0, 0.0, math.pi / 2.0)  # 90 degrees
        try:
            rotate(
                env,
                task,
                target_quat=target_quat,
                max_steps=120,
                threshold=0.04,
                timeout=8.0,
            )
        except Exception as e:
            print_exception(e)
            print("[Task] rotate() encountered an error – attempting to "
                  "continue to pull step anyway.")

        # ----- 2-c) Attempt to pull the drawer open ------------------------ #
        print("[Task] Attempting to pull the drawer open …")
        try:
            obs, reward, done = pull(
                env,
                task,
                object_name=drawer_name,   # assumed kwarg
                max_steps=150,
                timeout=10.0,
            )
            if done:
                print("[Task] Task episode ended after pull().")
        except Exception as e:
            print_exception(e)
            print("[Task] pull() failed – drawer might be locked or the "
                  "predicate precondition was unmet.")

        print("===== Task Execution Complete =====")
    finally:
        # Make sure simulator resources are cleaned up
        shutdown_environment(env)
        print("===== Environment Shut Down =====")


if __name__ == "__main__":
    run_skeleton_task()
