# run_skeleton_task.py (Filled-in Skeleton)

import time
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 * deliberately, so ALL predefined skills (move, pick, place, rotate, pull, …)
# that exist inside skill_code become directly accessible.
from skill_code import *

from video import init_video_writers, recording_step, recording_get_observation

# Returns a dictionary {object_name : np.ndarray([x,y,z])}
from object_positions import get_object_positions


def _safe_skill_call(skill_fn, *args, **kwargs):
    """
    Wrapper that safely executes any low-level skill.
    In case a skill raises an exception we simply print the error message and
    continue the exploration instead of crashing the whole run.
    """
    try:
        return skill_fn(*args, **kwargs)
    except Exception as e:
        # A lot of “domain errors” (e.g. trying to grasp the drawer body
        # instead of the handle) throw generic python errors – we catch them
        # here so that we can still continue our exploration.
        print(f"[WARNING] Skill {skill_fn.__name__} failed with: {e}")
        return None, 0.0, False


def exploration_phase(env, task, positions):
    """
    Minimalistic ‘exploration’ routine whose only purpose is to gather
    information about every reachable object.  We iterate over the known
    object dictionary, navigate to each object, try to pick it up, then put it
    back roughly at its original pose.  While doing so we implicitly execute
    the PDDL actions that would mark predicates such as
      identified / temperature-known / weight-known / durability-known /
      lock-known
    The concrete simulator obviously does not expose these logical predicates
    directly – nevertheless, the exploration phase will make sure that every
    possible skill is triggered at least once so that a later symbolic planner
    can “discover” which predicate had been missing.
    """
    print("==========  EXPLORATION  ==========")
    for obj_name, obj_pos in positions.items():

        print(f"\n[EXPLORE] -> Object “{obj_name}” @ {np.round(obj_pos, 3)}")

        # 1) Navigate close to the object (move skill).
        _safe_skill_call(
            move,
            env,
            task,
            target_pos=obj_pos,
            approach_distance=0.20,
            threshold=0.01,
            approach_axis='z',
            timeout=10.0
        )

        # 2) Try picking it up.
        obs, reward, done = _safe_skill_call(
            pick,
            env,
            task,
            target_pos=obj_pos,
            approach_distance=0.15,
            threshold=0.01,
            approach_axis='z',
            timeout=10.0
        )

        # If the task terminates while we are still exploring, we stop.
        if done:
            print("[EXPLORE] RLBench task ended – aborting exploration early.")
            return

        # 3) Immediately place it back (slightly above original pose so that
        #    we do not collide with the table).
        #    We offset in +z by 3 cm.
        if obs is not None:
            place_target = obj_pos + np.array([0.0, 0.0, 0.03])
            _safe_skill_call(
                place,
                env,
                task,
                target_pos=place_target,
                approach_distance=0.10,
                threshold=0.01,
                approach_axis='-z',
                timeout=10.0
            )

        # 4) Special-case: If the current object looks like a drawer handle,
        #    we also attempt the pull skill once in order to discover the
        #    “lock-known” predicate referred to in the feedback.
        #
        #    A very naive heuristic: any object that contains the substring
        #    “drawer” and “handle”.
        lowered = obj_name.lower()
        if ("drawer" in lowered) and ("handle" in lowered):
            print("[EXPLORE]   Attempting to rotate & pull the drawer handle.")
            # The exact pre-grasp pose for the drawer handle is unknown; we
            # therefore reuse the original position that has already been
            # reached during ‘pick’.
            _safe_skill_call(
                rotate,
                env,
                task,
                angle_from="zero_deg",   # These symbolic parameters must match
                angle_to="ninety_deg"    # the ones used inside the pre-defined
            )                             # rotate() skill implementation.

            _safe_skill_call(
                pull,
                env,
                task
            )


def run_skeleton_task():
    """Generic skeleton for running any task in the RLBench simulation."""
    print("===== Starting Skeleton Task =====")

    env, task = setup_environment()

    try:
        # ------------------------------------------------------------------
        # 1) Reset task & initialise video capturing
        # ------------------------------------------------------------------
        descriptions, obs = task.reset()
        init_video_writers(obs)

        # Wrap step / get_observation so that everything is recorded.
        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)

        # ------------------------------------------------------------------
        # 2) Gather positions of every relevant object in the scene
        # ------------------------------------------------------------------
        positions = get_object_positions()
        print(f"[INFO] Scene contains {len(positions)} objects:",
              ", ".join(positions.keys()))

        # ------------------------------------------------------------------
        # 3) –––––– Exploration Phase ––––––
        #    Before we attempt any goal-directed behaviour we first execute an
        #    exploration routine that will touch *every* object once.  This is
        #    necessary because the feedback we received pointed out that our
        #    symbolic knowledge base was missing a predicate (e.g. lock-known).
        # ------------------------------------------------------------------
        exploration_phase(env, task, positions)

        # ------------------------------------------------------------------
        # 4) Here one could add the *actual* high-level task sequence that
        #    solves the original benchmark (opening drawer, throwing rubbish,
        #    …).  Those steps are intentionally left out because the current
        #    assignment focuses on demonstrating how to trigger every low-level
        #    skill at least once so that the planner can infer the missing
        #    predicates from observations made during exploration.
        # ------------------------------------------------------------------

        print("[MAIN] Finished exploration. No goal-directed plan executed "
              "because this file only demonstrates predicate discovery.")
        time.sleep(1.0)   # Give the simulator a moment before shutting down.

    finally:
        # We *always* shut down the environment, even if an exception happened.
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()