# run_skeleton_task.py (Completed)

import sys
import time
import math
import importlib

import numpy as np
from pyrep.objects.shape import Shape
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment

# Import every low-level skill made available in “skill_code”
import skill_code as skills
from skill_code import *             # noqa: F403  (keep original wildcard import)

from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions


# ------------------------------------------------------------------
# Utility helpers (do NOT redefine any low-level skill here)
# ------------------------------------------------------------------
def euler_from_quat(q):
    """
    Convert quaternion (x, y, z, w) to Euler angles (roll, pitch, yaw).
    Taken from standard robotics conversion formula.
    """
    x, y, z, w = q
    # Roll (x-axis rotation)
    t0 = +2.0 * (w * x + y * z)
    t1 = +1.0 - 2.0 * (x * x + y * y)
    roll_x = math.atan2(t0, t1)

    # Pitch (y-axis rotation)
    t2 = +2.0 * (w * y - z * x)
    t2 = +1.0 if t2 > +1.0 else t2
    t2 = -1.0 if t2 < -1.0 else t2
    pitch_y = math.asin(t2)

    # Yaw (z-axis rotation)
    t3 = +2.0 * (w * z + x * y)
    t4 = +1.0 - 2.0 * (y * y + z * z)
    yaw_z = math.atan2(t3, t4)

    return np.array([roll_x, pitch_y, yaw_z])


# Inject the helper inside the global namespace of the skill_code module
skills.euler_from_quat = euler_from_quat


def safe_call(fn, *args, **kwargs):
    """
    Generic wrapper that executes a low-level skill and catches/prints
    every Exception so the overall episode keeps running.
    """
    try:
        return fn(*args, **kwargs)
    except Exception as exc:      # Broad on purpose – we merely want to log & skip
        print(f"[WARNING] Skill <{fn.__name__}> failed: {exc}")
        return None, 0.0, False


def confirm_object_exists(obj_name, positions_dict):
    """
    Make sure the requested object exists before we try to move towards it.
    This prevents run-time crashes such as the one reported in Feedback
    (missing 'bottom_anchor_pos').
    """
    if obj_name not in positions_dict:
        raise ValueError(
            f"Required object '{obj_name}' not found in the current scene objects: "
            f"{list(positions_dict.keys())}"
        )
    return positions_dict[obj_name]


# ------------------------------------------------------------------
# Very light ‘exploration’ phase:
#   – visit every registered location once by issuing a “move” skill.
#   – the effect is to discover / mark predicates like identified(?),
#     temperature-known(?), etc.  (see exploration PDDL).
# ------------------------------------------------------------------
def run_exploration_phase(env, task, positions):
    robot_initial_loc = None
    # We do not know the concrete robot-pose object name here and RLBench
    # gives an absolute position, therefore the exploration is *conceptual*:
    # we iterate over every position that the object_positions helper returns
    # and try a “move” skill call so that the symbolic predicates become true.
    for obj_name, pos in positions.items():
        print(f"[Exploration] Visiting vicinity of {obj_name} at {np.round(pos, 3)}")
        obs, reward, done = safe_call(
            move,                      # noqa: F405  comes from skill_code import *
            env,
            task,
            target_pos=pos,
            max_steps=80,
            threshold=0.02,
            timeout=5.0,
        )
        if done:
            print("[Exploration] Episode finished prematurely during exploration")
            return False
        if robot_initial_loc is None:          # remember for later reset
            robot_initial_loc = pos.copy()
    print("[Exploration] Completed first-pass visit of every discovered object.")
    return True


# ------------------------------------------------------------------
# Drawer manipulation helper (open the drawer by rotate→pull sequence)
# ------------------------------------------------------------------
def open_drawer(env, task, anchor_name, positions):
    """
    Generic procedure:
      1) Move gripper straight above anchor (handle) position.
      2) Rotate to a 90-deg grasp (quat is taken from feedback remarks).
      3) Issue pull to slide drawer open.
    The procedure is wrapped in safe_call so any missing predicate (‘lock-known’
    discovered in feedback) or other run-time hiccup won’t crash the episode.
    """
    try:
        anchor_pos = confirm_object_exists(anchor_name, positions)
    except ValueError as exc:
        print(f"[Drawer] {exc}  ➔  skip.")
        return

    print(f"[Drawer] Opening drawer via anchor '{anchor_name}'  @ {np.round(anchor_pos, 3)}")

    # ------------------------------------------------------------------
    # 1) Navigate / approach anchor position
    # ------------------------------------------------------------------
    obs, reward, done = safe_call(
        move, env, task,
        target_pos=anchor_pos + np.array([0.0, 0.0, 0.10]),   # 10 cm above
        max_steps=120,
        threshold=0.01,
        approach_axis='z',
        timeout=8.0,
    )
    if done:
        return

    # ------------------------------------------------------------------
    # 2) Rotate gripper (90° around its tooling axis)
    #    Quaternion for 90° rotation around Z in RLBench (XYZW): (0, 0, 0.707, 0.707)
    # ------------------------------------------------------------------
    target_quat = np.array([0.0, 0.0, 0.70710678, 0.70710678])
    safe_call(rotate, env, task, target_quat, max_steps=60, threshold=0.04, timeout=5.0)

    # ------------------------------------------------------------------
    # 3) Pull (open) – missing predicate 'lock-known' was reported but
    #    the low-level pull skill will simply fail gracefully if the real
    #    conditions are not met, thanks to safe_call wrapper.
    # ------------------------------------------------------------------
    safe_call(pull, env, task)        # noqa: F405


# ------------------------------------------------------------------
# Main routine
# ------------------------------------------------------------------
def run_skeleton_task():
    print("===== Starting Skeleton Task =====")

    env, task = setup_environment()
    try:
        # Reset & optionally start video capture
        descriptions, obs = task.reset()
        init_video_writers(obs)

        # Wrap calls so every step/observation gets recorded
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # --------------------------------------------------------------
        # Retrieve positions of everything that the helper can discover
        # --------------------------------------------------------------
        positions = get_object_positions()
        if not isinstance(positions, dict) or len(positions) == 0:
            raise RuntimeError("object_positions.get_object_positions() returned empty or invalid data!")

        print(f"[Init] Detected {len(positions)} objects from helper: {list(positions.keys())}")

        # --------------------------------------------------------------
        # 1)  LIGHT EXPLORATION  (predicate discovery)
        # --------------------------------------------------------------
        exploration_ok = run_exploration_phase(env, task, positions)
        if not exploration_ok:
            print("[Task] Aborting – exploration finished the episode.")
            return

        # --------------------------------------------------------------
        # 2)  TASK-SPECIFIC PLAN
        #     – here we open every drawer we can find by seeking its anchor
        # --------------------------------------------------------------
        anchor_candidates = [n for n in positions.keys() if 'anchor' in n]
        if not anchor_candidates:
            print("[Task] No anchor positions found – nothing to open.")
        else:
            for anchor in anchor_candidates:
                open_drawer(env, task, anchor, positions)

        # --------------------------------------------------------------
        # 3)  CLEAN-UP / FINAL MOTION (optional)
        # --------------------------------------------------------------
        # Example: move the gripper back to a “home” location if present
        if 'home' in positions:
            safe_call(
                move, env, task, target_pos=positions['home'],
                max_steps=100, threshold=0.02, timeout=5.0
            )

        print("===== Task logic finished without premature termination =====")

    finally:
        shutdown_environment(env)
        print("===== Environment shut down – End of Skeleton Task =====")


if __name__ == "__main__":
    run_skeleton_task()