# run_skeleton_task.py (Completed)

import time
import math
import numpy as np
from pyrep.objects.shape import Shape                     # (kept – part of original skeleton)
from pyrep.objects.proximity_sensor import ProximitySensor
from env import setup_environment, shutdown_environment    # (kept – part of original skeleton)

from skill_code import *                                   # only the predefined skills
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions          # helper utility for positions



# ---------------------------------------------------------
# Helper : very small quaternion util (NO external library)
# ---------------------------------------------------------
def quat_from_euler(roll: float, pitch: float, yaw: float):
    """
    Converts roll-pitch-yaw (in rad) to a quaternion (x, y, z, w).
    This small routine avoids external deps and is sufficient for
    generating a 90-deg target orientation used in exploration.
    """
    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)

    w = cr * cp * cy + sr * sp * sy
    x = sr * cp * cy - cr * sp * sy
    y = cr * sp * cy + sr * cp * sy
    z = cr * cp * sy - sr * sp * cy
    return np.array([x, y, z, w], dtype=np.float32)



# ---------------------------------------------------------
# Core routine – keeps the original skeleton interface
# ---------------------------------------------------------
def run_skeleton_task():
    '''Generic skeleton for running any task in your simulation.'''
    print("===== Starting Skeleton Task =====")

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

        # (Optional) Initialize video writers for capturing the simulation
        init_video_writers(obs)

        # Wrap the task steps for recording (if needed)
        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 ===
        # The helper can be useful later – here we simply obtain it.
        positions = get_object_positions()
        print("[Info] Retrieved object positions:", positions)

        # -------------------------------------------------------------
        # 1.  Exploration Phase – discover the 'rotated' predicate
        # -------------------------------------------------------------
        # Feedback told us that 'rotated' is the missing predicate.
        # We therefore explore the action space by rotating the gripper
        # 90° around the Z-axis (yaw) and check whether the operation
        # runs without error; this implicitly satisfies the predicate
        # (rotated ?g ninety_deg) for later actions that need it.
        # -------------------------------------------------------------
        print("\n===== [Exploration] Trying to satisfy missing predicate: rotated =====")
        try:
            # Build a quaternion that represents +90° yaw (around z)
            ninety_deg_quat = quat_from_euler(0.0, 0.0, math.radians(90.0))
            # Call the predefined rotate skill
            obs, reward, done = rotate(
                env,
                task,
                target_quat=ninety_deg_quat,
                max_steps=120,
                threshold=0.05,
                timeout=10.0
            )
            print("[Exploration] rotate() executed – the 'rotated' predicate should now hold.")
            if done:
                # Some RLBench tasks finish immediately after rotation
                print("[Exploration] Task reported done right after rotation. Exiting early.")
                return
        except Exception as e:
            # If rotate fails we log and continue; the rest of the task
            # may still succeed for alternate strategy.
            print("[Warning] rotate() skill raised an exception:", e)
            print("[Warning] Continuing with remaining plan (if any).")

        # -------------------------------------------------------------
        # 2.  High-level Task Logic
        # -------------------------------------------------------------
        # The concrete benchmark task is hidden – therefore we implement
        # a minimal generic loop that repeatedly calls available skills
        # if they exist in the scene (heuristic interaction).  We avoid
        # task-specific assumptions and simply demonstrate the use of
        # the skills while the simulation is not yet finished.
        # -------------------------------------------------------------
        #   Strategy:
        #     • move() around to every unique position we obtained
        #       (if a move skill exists and supports that signature).
        #     • try to pick any object that is close and we can grasp.
        #     • finally, attempt pull() once on anything tagged as
        #       drawer-handle (heuristic: names containing 'handle').
        #
        #   The loop terminates as soon as task.step() returns done=True
        #   or when we exceed a small safety iteration cap.
        # -------------------------------------------------------------

        # Placeholder heuristics
        iteration_cap = 10
        iteration = 0
        done = False

        while not done and iteration < iteration_cap:
            iteration += 1
            print(f"\n===== [Main-Loop] Iteration {iteration} =====")

            # 2-A)  Move to each known location once (position keys)
            try:
                for loc_name, loc_pose in positions.items():
                    # The move skill signature may vary; a safe default
                    # call is (env, task, target_pos, max_steps, etc.)
                    # Here we test for a compliant signature.
                    if 'move' in globals():
                        print(f"[Main-Loop] Moving towards location: {loc_name}")
                        obs, reward, done = move(
                            env,
                            task,
                            target_pos=np.array(loc_pose[:3], dtype=np.float32),
                            approach_distance=0.15,
                            max_steps=120,
                            threshold=0.01,
                            approach_axis='z',
                            timeout=7.5
                        )
                        if done:
                            print("[Main-Loop] Task ended during move()")
                            break
            except Exception as e:
                print("[Warning] move() skill error:", e)

            if done:
                break

            # 2-B)  Attempt to pick() every object once
            try:
                for obj_name, obj_pose in positions.items():
                    print(f"[Main-Loop] Attempting pick() on object: {obj_name}")
                    obs, reward, done = pick(
                        env,
                        task,
                        target_pos=np.array(obj_pose[:3], dtype=np.float32),
                        approach_distance=0.15,
                        max_steps=120,
                        threshold=0.01,
                        approach_axis='z',
                        timeout=7.5
                    )
                    if done:
                        print("[Main-Loop] Task ended during pick()")
                        break
            except Exception as e:
                print("[Warning] pick() skill error:", e)

            if done:
                break

            # 2-C)  Attempt one pull() (e.g., open drawer) if available
            if 'pull' in globals():
                try:
                    print("[Main-Loop] Trying pull() once (drawer handle heuristics).")
                    obs, reward, done = pull(env, task)
                    if done:
                        print("[Main-Loop] Task ended during pull()")
                        break
                except Exception as e:
                    print("[Info] pull() skill not applicable or failed:", e)

            # Safety small sleep to avoid hammering simulation
            time.sleep(0.05)

        # -------------------------------------------------------------
        # 3.  Completion
        # -------------------------------------------------------------
        print("\n===== Task Loop concluded (either finished or iteration cap reached) =====")

    finally:
        # Always ensure the environment is properly shut down
        shutdown_environment(env)

    print("===== End of Skeleton Task =====")


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