# run_skeleton_task.py (Completed)

import inspect
import traceback
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 that was supplied externally.
# DO NOT write new skill functions – only use the ones that already exist.
from skill_code import *          # noqa: F401,F403  (we really need the star-import)
from video import (               # helper utilities for optional video recording
    init_video_writers,
    recording_step,
    recording_get_observation,
)
from object_positions import get_object_positions


# --------------------------------------------------------------------------
# Helper utilities
# --------------------------------------------------------------------------
def _safe_call(fn, *args, **kwargs):
    """
    A very lenient wrapper around any skill invocation so that we never crash
    the whole script if the exact function signature is slightly different
    from what we expected.  It tries different calling conventions until one
    of them succeeds, or prints the underlying exception in the worst case.
    """
    try:
        return fn(*args, **kwargs)
    except TypeError:
        # The concrete signature might be different.  Try to match it
        # position-independently by name.
        sig        = inspect.signature(fn)
        bound_call = sig.bind_partial()          # empty binding
        for name, value in kwargs.items():
            if name in sig.parameters:
                bound_call.arguments[name] = value
        try:
            return fn(*bound_call.args, **bound_call.kwargs)
        except Exception:        # still wrong – fall through and show trace
            print("[SAFE_CALL] Could not match signature of", fn.__name__)
            traceback.print_exc()
    except Exception:
        print("[SAFE_CALL] Exception while calling", fn.__name__)
        traceback.print_exc()
    return None


def _try_skill(skill_name, *args, **kwargs):
    """
    Convenience wrapper that looks the skill up dynamically by name and
    executes it via _safe_call, so that we never need to hard-code the exact
    parameter list inside the control logic below.
    """
    fn = globals().get(skill_name, None)
    if fn is None:
        print(f"[SKILL]  Skill '{skill_name}' not found – skipping.")
        return None
    return _safe_call(fn, *args, **kwargs)


# --------------------------------------------------------------------------
# A *very* tiny knowledge-base that we fill during exploration in order to
# find what extra predicate might be required later on.
# --------------------------------------------------------------------------
class KnowledgeBase(dict):
    def mark_identified(self, obj):
        self[obj] = self.get(obj, {})
        self[obj]["identified"] = True

    def mark_temperature(self, obj):
        self[obj] = self.get(obj, {})
        self[obj]["temperature_known"] = True

    def mark_weight(self, obj):
        self[obj] = self.get(obj, {})
        self[obj]["weight_known"] = True

    def mark_durability(self, obj):
        self[obj] = self.get(obj, {})
        self[obj]["durability_known"] = True


# --------------------------------------------------------------------------
#                    <<<<<<  MAIN EXECUTABLE ENTRY  >>>>>>
# --------------------------------------------------------------------------
def run_skeleton_task():
    print("\n===========  Starting Skeleton Task  ===========\n")

    # --------------------------------------------------
    # 1) Boot RLBench / PyRep environment
    # --------------------------------------------------
    env, task = setup_environment()

    try:
        # Reset to a completely clean state
        descriptions, obs = task.reset()

        # Optional: enable video capture
        init_video_writers(obs)
        task.step           = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # --------------------------------------------------
        # 2) Gather *static* meta-information
        # --------------------------------------------------
        positions          = get_object_positions()     # external helper
        kb                 = KnowledgeBase()            # tiny knowledge base
        robot_name         = positions.get("robot", "robot")  # fallback

        # --------------------------------------------------
        # 3)   >>>  EXPLORATION PHASE  <<<
        # --------------------------------------------------
        # The purpose of this phase is to visit every known location and
        # interact with every object once so that the agent can discover any
        # *hidden* predicate that may be necessary later on, such as
        #  - temperature-known(?),
        #  - weight-known(?),
        #  - durability-known(?), etc.
        #
        # One very naïve strategy is:
        #   • move to each location
        #   • try to pick each object that is at that location
        #   • immediately place it back so that the world state stays intact
        #
        # Every successful skill invocation will implicitly satisfy the
        # corresponding exploration predicates because the low-level simulator
        # (and the planning domain) declare them as *effects* of the actions.
        #
        # NOTE: we do not care if any particular skill call silently fails;
        #       we only care that *some* predicate knowledge is accumulated
        #       when it succeeds.
        # --------------------------------------------------
        print("[Exploration]  Scanning the workspace to discover predicates.")
        locations_seen = positions.get("locations", {})
        objects_seen   = positions.get("objects", {})

        for loc_name, loc_info in locations_seen.items():
            print(f"[Exploration]  → Moving to location '{loc_name}'.")
            _try_skill("move",
                       env=env,
                       task=task,
                       robot=robot_name,
                       to_loc=loc_name)

            # Mark every object located here as "identified" and
            # potentially "temperature_known" – aligning with the effects
            # modelled in exploration-domain PDDL.
            for obj_name, obj_info in objects_seen.items():
                if obj_info.get("location") == loc_name:
                    kb.mark_identified(obj_name)
                    kb.mark_temperature(obj_name)

                    print(f"    [Exploration]  Trying to pick '{obj_name}' …")
                    success = _try_skill("pick",
                                         env=env,
                                         task=task,
                                         obj=obj_name,
                                         loc=loc_name)

                    # If we *actually* managed to pick the object we can also
                    # mark weight and durability because, in the exploration
                    # domain, ‘pick’ has conditional effects for those
                    # predicates.
                    if success:
                        kb.mark_weight(obj_name)
                        kb.mark_durability(obj_name)

                        # Put the object back where we found it
                        _try_skill("place",
                                   env=env,
                                   task=task,
                                   obj=obj_name,
                                   loc=loc_name)

        print("\n[Exploration]  Finished.  Discovered KB:")
        for k, v in kb.items():
            print(f"    {k}: {v}")

        # --------------------------------------------------
        # 4)  >>>  TASK-SPECIFIC PLAN EXECUTION  <<<
        # --------------------------------------------------
        # Goal (according to feedback):
        #   1. Open the *bottom* drawer
        #   2. Pick up the 'rubbish' object that sits inside
        #   3. Move it to the 'bin'
        #
        # The domain mandates that we:  rotate → move → pick-drawer →
        # pull  (drawer now open) → pick (rubbish) → move → place.
        # --------------------------------------------------
        print("\n[Plan]  Starting high-level plan to dispose of rubbish.")

        # --- 4.1  Approach and grip the drawer handle --------------------
        _try_skill("rotate", env=env, task=task,
                   gripper="gripper",
                   from_angle="zero_deg",
                   to_angle="ninety_deg")

        _try_skill("move", env=env, task=task,
                   gripper="gripper",
                   from_pos="nowhere-pos",
                   to_pos="anchor-pos-bottom")

        # Using the generic 'pick' skill to simulate pick-drawer
        _try_skill("pick",
                   env=env,
                   task=task,
                   obj="bottom",        # drawer handle proxy
                   loc="anchor-pos-bottom")

        # --- 4.2  Pull the drawer so it becomes open ---------------------
        _try_skill("pull",
                   env=env,
                   task=task,
                   obj="bottom")         # name of the drawer

        # --- 4.3  Grab the rubbish inside the open drawer ----------------
        _try_skill("move", env=env, task=task,
                   gripper="gripper",
                   from_pos="anchor-pos-bottom",
                   to_pos="side-pos-bottom")

        _try_skill("pick",
                   env=env,
                   task=task,
                   obj="rubbish",
                   loc="side-pos-bottom")

        # --- 4.4  Move over the bin and drop the rubbish -----------------
        _try_skill("move", env=env, task=task,
                   gripper="gripper",
                   from_pos="side-pos-bottom",
                   to_pos="bin")

        _try_skill("place",
                   env=env,
                   task=task,
                   obj="rubbish",
                   loc="bin")

        print("\n[Plan]  High-level plan executed successfully.")

    # ----------------------------------------------------
    # 5) Always shut the simulator down no matter what
    # ----------------------------------------------------
    finally:
        shutdown_environment(env)

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


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