# run_skeleton_task.py (Completed)

import time
import traceback
from typing import Dict, Tuple

import numpy as np
from pyrep.objects.shape import Shape      # skeleton import – do not remove
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment

# bring every low-level skill that was shipped with the simulator into scope
from skill_code import *                   # noqa: F403, F401

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


# --------------------------------------------------------------------------- #
# --------------------------- Helper / Guard Routines ----------------------- #
# --------------------------------------------------------------------------- #
def _log_exception(err: Exception, msg: str = "") -> None:
    """Utility to give very compact but useful error logs and keep simulation
    running whenever reasonable."""
    print(f"[Error] {msg}\n{err}")
    traceback.print_exc()


def _force_calibrated() -> bool:
    """
    The real force-sensor interface is not exposed in the public template.
    We therefore return True but keep the function for completeness so that we
    can extend it later without changing the main control flow.
    """
    return True


def _drawer_is_open(task) -> bool:
    """
    Very light-weight safety guard.  
    Because we do not have a symbolic ‘is_open’ query on the actual RLBench
    task, we approximate by checking a (possibly) existing drawer handle shape
    and seeing if it moved far enough.  If anything goes wrong we fall back
    to ‘unknown / closed’ and let the plan deal with it.
    """
    try:
        # Adopt a widely used naming convention in RLBench assets
        handle = Shape("drawer_handle")
        # The drawer starts roughly at x ≈ 0.12 (depends on task).  If it is
        # beyond that we treat it as *open*.
        handle_pos = handle.get_position()
        if handle_pos[1] > 0.12:
            return True
    except Exception:
        pass
    return False


def _object_is_known(obj_name: str, positions: Dict[str, Tuple[float]]) -> bool:
    """Item-existence guard used before pick() call."""
    return obj_name in positions


# --------------------------------------------------------------------------- #
# ---------------------------    Main Controller    ------------------------- #
# --------------------------------------------------------------------------- #
def run_skeleton_task() -> None:
    """
    Generic controller that (1) runs an exploration loop to gather all missing
    symbolic knowledge that might be required by the high-level planner
    (temperature-known, durability-known, lock-known …) and (2) shows an
    example of how the primitive skills can be chained together.
    """
    print("===== Starting Skeleton Task =====")

    # === Environment Setup ===
    env, task = setup_environment()

    try:
        # ------------------------------------------------------------------ #
        # ------------------ Initialisation / Recording -------------------- #
        # ------------------------------------------------------------------ #
        descriptions, obs = task.reset()
        init_video_writers(obs)

        # Make every environment step automatically go through the recording
        # wrappers.  Their implementations are part of the public template.
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # ------------------------------------------------------------------ #
        # ------------------------ Exploration Phase ----------------------- #
        # ------------------------------------------------------------------ #
        """
        The symbolic exploration described in the additional PDDL domain is
        achieved by *physically* moving the robot to every location that might
        contain an object, doing a soft touch (pick attempt) if the gripper is
        empty, and – if we encounter a drawer – running the pull() skill once.
        The loop yields:
            identified          – by being at the same location
            temperature-known   – idem               (second move variant)
            weight-known        – by pick()
            durability-known    – by pick()
            lock-known          – by pull()
        """

        positions = get_object_positions()     # type: Dict[str, Tuple[float]]
        if not positions:
            print("[Exploration] No positions returned – skipping exploration.")
        else:
            # A minimalistic radius (in metre) that we keep above any object
            APPROACH_HEIGHT = 0.15

            for obj_name, obj_pos in positions.items():
                try:
                    print(f"[Exploration] -> approaching {obj_name} @ {obj_pos}")

                    # ------------------------------------------------------ #
                    # 1) MOVE close to the object – gathers
                    #    identified & temperature-known predicates.
                    # ------------------------------------------------------ #
                    try:
                        # skill signatures: (env, task, target_pos, **kwargs)
                        # ‘move’ keeps the gripper 15 cm above the target
                        obj_pos_above = (obj_pos[0], obj_pos[1], obj_pos[2] + APPROACH_HEIGHT)
                        move(                           # noqa: F405 (imported *)
                            env,
                            task,
                            target_pos=obj_pos_above,
                            max_steps=150,
                            threshold=0.02,
                            timeout=6.0
                        )
                    except Exception as err:
                        _log_exception(err, f"while moving above {obj_name}")
                        continue

                    # ------------------------------------------------------ #
                    # 2) Try to PICK the object – gathers weight/durability.
                    # ------------------------------------------------------ #
                    if _object_is_known(obj_name, positions):
                        # We descend slightly towards the object then call pick().
                        real_target = (obj_pos[0], obj_pos[1], obj_pos[2] + 0.01)
                        pick(                           # noqa: F405
                            env,
                            task,
                            target_pos=real_target,
                            approach_distance=APPROACH_HEIGHT,
                            max_steps=120,
                            threshold=0.01,
                            approach_axis='z',
                            timeout=8.0
                        )
                        # The template pick() function will *not* drop the
                        # object afterwards, so we put it back immediately –
                        # exploration should leave the scene unchanged.
                        place(                          # noqa: F405
                            env,
                            task,
                            target_pos=obj_pos_above,
                            approach_distance=APPROACH_HEIGHT,
                            max_steps=120,
                            threshold=0.01,
                            approach_axis='z',
                            timeout=8.0
                        )

                    # ------------------------------------------------------ #
                    # 3) If the object is likely a drawer handle run pull().
                    # ------------------------------------------------------ #
                    if "drawer" in obj_name or "handle" in obj_name:
                        try:
                            if _drawer_is_open(task):
                                print("[Exploration] Drawer already open – skipping pull.")
                            else:
                                # Basic force safety: skip if the gripper
                                # calibration is unverified.
                                if not _force_calibrated():
                                    print("[Exploration] Gripper not calibrated – abort pull.")
                                else:
                                    pull(env, task)      # noqa: F405
                                    print("[Exploration] pull() executed on drawer.")
                        except Exception as err:
                            _log_exception(err, "during drawer pull")

                except Exception as err_outer:
                    _log_exception(err_outer, "outer exploration loop")

        # ------------------------------------------------------------------ #
        # ---------------------------  Main Plan --------------------------- #
        # ------------------------------------------------------------------ #
        """
        A *very* small demonstration plan: 
        1) Rotate the gripper 90° (to ‘ninety_deg’ which the PDDL expects).
        2) Move to the drawer handle (we re-use the position dict).
        3) Pull the drawer open. 
        """

        # Step 1 – rotate
        ninety_deg_quat = np.array([0., 0., np.sin(np.pi/4), np.cos(np.pi/4)])  # (x, y, z, w)
        try:
            rotate(                        # noqa: F405
                env,
                task,
                target_quat=ninety_deg_quat,
                max_steps=100,
                threshold=0.04,
                timeout=8.0
            )
        except Exception as err:
            _log_exception(err, "rotating the gripper")

        # Step 2 – locate a drawer handle again
        drawer_handle_name = None
        drawer_handle_pos = None
        for name, pos in positions.items():
            if "drawer" in name or "handle" in name:
                drawer_handle_name = name
                drawer_handle_pos = pos
                break

        if drawer_handle_name and drawer_handle_pos:
            print(f"[Plan] Approach drawer handle '{drawer_handle_name}'.")

            # Move above it
            try:
                move(                      # noqa: F405
                    env,
                    task,
                    target_pos=(drawer_handle_pos[0],
                                drawer_handle_pos[1],
                                drawer_handle_pos[2] + 0.12),
                    max_steps=150,
                    threshold=0.02,
                    timeout=6.0
                )
            except Exception as err:
                _log_exception(err, "approaching drawer handle")

            # Pull
            try:
                if not _drawer_is_open(task):
                    pull(env, task)        # noqa: F405
                    print("[Plan] Drawer pull executed.")
                else:
                    print("[Plan] Drawer already open – no pull needed.")
            except Exception as err:
                _log_exception(err, "executing pull()")

        else:
            print("[Plan] No drawer handle found in the position list – plan ends.")

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


if __name__ == "__main__":
    run_skeleton_task()