# run_skeleton_task.py  (Completed Version – uses ONLY the pre-defined skills)

import time
import traceback
import numpy as np

from env import setup_environment, shutdown_environment
from skill_code import move, pick, place, rotate, pull          # ONLY the allowed skills
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions


# ------------------------------------------------------------------
# Helper utilities – kept light-weight, do NOT re-implement any skill
# ------------------------------------------------------------------
def quat_from_euler(roll: float, pitch: float, yaw: float):
    """Small local replacement for transforms3d / scipy.  Returns xyzw order."""
    cr, sr = np.cos(roll * .5), np.sin(roll * .5)
    cp, sp = np.cos(pitch * .5), np.sin(pitch * .5)
    cy, sy = np.cos(yaw * .5), np.sin(yaw * .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)


def safe_call(fn, *args, **kwargs):
    """
    Light-weight wrapper that executes a skill, prints any exception,
    and keeps the main control-flow alive.
    """
    try:
        return fn(*args, **kwargs)
    except Exception as e:
        print(f"[safe_call] WARNING – skill <{fn.__name__}> raised an exception:")
        traceback.print_exc()
        # Return sensible defaults to keep the rest of the logic alive
        return None, 0.0, False


# ------------------------------------------------------------------
# Missing-predicate exploration
# ------------------------------------------------------------------
def exploration_phase(env, task, positions: dict):
    """
    A quick exploratory routine whose purpose is to discover whether the
    drawer is locked (corresponds to missing predicate ‘is-locked’
    / ‘lock-known’ in the PDDL feedback).

    Strategy:
        • Try to grasp a drawer-handle-looking object.
        • Attempt pull().  If it fails immediately we hypothesise that the
          drawer is locked and log this information.
    """
    drawer_handle_key = next((k for k in positions if 'handle' in k.lower()
                              or ('drawer' in k.lower() and 'handle' in k.lower())),
                             None)
    if drawer_handle_key is None:
        print("[exploration] No obvious drawer handle found – skipping exploration.")
        return

    handle_pos = positions[drawer_handle_key]
    print(f"[exploration] Candidate drawer handle found: {drawer_handle_key} @ {handle_pos}")

    # 1) Move to the handle
    print("[exploration] Moving towards the drawer handle …")
    safe_call(move, env, task, target_pos=handle_pos,
              approach_distance=0.15, max_steps=120,
              threshold=0.01, approach_axis='z', timeout=8.0)

    # 2) Pick the handle (simulates pick-drawer)
    print("[exploration] Grasping the handle …")
    safe_call(pick, env, task, target_pos=handle_pos,
              approach_distance=0.07, max_steps=120,
              threshold=0.01, approach_axis='z', timeout=8.0)

    # 3) Try a light pull
    print("[exploration] Attempting pull() to see if the drawer is locked …")
    obs, reward, done = safe_call(pull, env, task)

    # Depending on the result we set an internal flag
    if done is False:
        print("[exploration] Drawer pull did not terminate the task – likely LOCKED.")
        task.drawer_locked = True          # Store hypothesis on the task object
    else:
        print("[exploration] Drawer opened – it was UNLOCKED.")
        task.drawer_locked = False


# ------------------------------------------------------------------
# Main task runner
# ------------------------------------------------------------------
def run_skeleton_task():
    print("==========  STARTING SKELETON TASK  ==========")

    # ---------------------------------------------------
    # 0)  Environment initialisation
    # ---------------------------------------------------
    env, task = setup_environment()
    try:
        descriptions, obs = task.reset()
        init_video_writers(obs)

        # Patch the task object so that every step is recorded automatically
        original_step_fn = task.step
        task.step = recording_step(original_step_fn)
        original_get_obs = task.get_observation
        task.get_observation = recording_get_observation(original_get_obs)

        # ---------------------------------------------------
        # 1)  Acquire object/world information
        # ---------------------------------------------------
        positions = get_object_positions()
        print(f"[main] Retrieved {len(positions)} object positions from helper module.")

        # ---------------------------------------------------
        # 2)  Quick exploration to deduce hidden predicates
        # ---------------------------------------------------
        exploration_phase(env, task, positions)

        # ---------------------------------------------------
        # 3)  High-level oracle plan
        #      (Very generic – adjust to your concrete RLBench scene)
        # ---------------------------------------------------
        #
        # We try to:
        #   a) Open a drawer (if not locked)
        #   b) Pick a rubbish object
        #   c) Put it inside the opened drawer
        #
        # ---------------------------------------------------

        # a)  Try to open (again) only if earlier exploration either succeeded
        #     or determined the drawer was unlocked.
        drawer_handle_key = next((k for k in positions if 'handle' in k.lower()
                                  or ('drawer' in k.lower() and 'handle' in k.lower())),
                                 None)
        if drawer_handle_key:
            handle_pos = positions[drawer_handle_key]

            # Align gripper – coarse 90-deg rotation around its local Z
            target_quat = quat_from_euler(0.0, 0.0, np.deg2rad(90))
            print("[main] Rotating gripper for a better pulling posture …")
            safe_call(rotate, env, task, target_quat,
                      max_steps=100, threshold=0.04, timeout=8.0)

            # Move & pick the handle if not already holding it
            print("[main] Moving to the drawer handle …")
            safe_call(move, env, task, target_pos=handle_pos,
                      approach_distance=0.15, max_steps=120,
                      threshold=0.01, approach_axis='z', timeout=8.0)

            print("[main] Grasping the handle …")
            safe_call(pick, env, task, target_pos=handle_pos,
                      approach_distance=0.07, max_steps=120,
                      threshold=0.01, approach_axis='z', timeout=8.0)

            # If the drawer is believed not locked → pull
            if getattr(task, 'drawer_locked', False) is False:
                print("[main] Pulling the drawer open …")
                safe_call(pull, env, task)
            else:
                print("[main] Drawer is believed locked – skipping pull to avoid failure.")

        else:
            print("[main] No drawer handle found – skipping drawer manipulation.")

        # ---------------------------------------------------
        #  b)  PICK the rubbish / trash object
        # ---------------------------------------------------
        trash_key = next((k for k in positions if
                          any(word in k.lower() for word in ['trash', 'rubbish', 'garbage', 'object'])),
                         None)
        if trash_key is None:
            print("[main] No obvious rubbish object found – task ends here.")
            return

        trash_pos = positions[trash_key]
        print(f"[main] Rubbish candidate: {trash_key} @ {trash_pos}")

        print("[main] Moving to rubbish …")
        safe_call(move, env, task, target_pos=trash_pos,
                  approach_distance=0.15, max_steps=120,
                  threshold=0.01, approach_axis='z', timeout=8.0)

        print("[main] Grasping rubbish …")
        safe_call(pick, env, task, target_pos=trash_pos,
                  approach_distance=0.05, max_steps=120,
                  threshold=0.01, approach_axis='z', timeout=8.0)

        # ---------------------------------------------------
        #  c)  PLACE inside the drawer (or any receptacle)
        # ---------------------------------------------------
        receptacle_key = next((k for k in positions if
                               any(sub in k.lower() for sub in ['drawer_inside', 'drawer', 'bin', 'receptacle'])),
                              None)
        if receptacle_key is None:
            print("[main] No receptacle location known – dropping at current spot.")
            target_place_pos = trash_pos + np.array([0.0, 0.0, 0.10])   # lift a bit and drop
        else:
            target_place_pos = positions[receptacle_key]
            print(f"[main] Receptacle candidate: {receptacle_key} @ {target_place_pos}")

        print("[main] Moving to placement location …")
        safe_call(move, env, task, target_pos=target_place_pos,
                  approach_distance=0.15, max_steps=120,
                  threshold=0.01, approach_axis='z', timeout=8.0)

        print("[main] Placing rubbish …")
        safe_call(place, env, task, target_pos=target_place_pos,
                  approach_distance=0.05, max_steps=120,
                  threshold=0.01, approach_axis='z', timeout=8.0)

        print("==========  TASK FINISHED SUCCESSFULLY  ==========")

    finally:
        shutdown_environment(env)
        print("==========  ENVIRONMENT SHUTDOWN COMPLETE  ==========")


if __name__ == "__main__":
    run_skeleton_task()