# run_skeleton_task.py (Completed Skeleton – exploration for the missing predicate
# “is-anchor-pos”)
#
# NOTE:
# • No new low-level skills are introduced – the code only CALLS the skills that
#   are assumed to exist in “skill_code” (move, pick, place, rotate, pull).
# • The focus is to demonstrate an “exploration phase” that tries to figure out
#   which position around a drawer satisfies the (unknown to the robot) predicate
#   is-anchor-pos.  We do this by iteratively moving the gripper to every candidate
#   position returned by get_object_positions() and attempting a “pull”.  If the
#   drawer opens we record that position as the anchor point.
# • The code is intentionally defensive – every skill invocation is wrapped in a
#   try/except so that unexpected errors do not crash the whole run.


import numpy as np
from pyrep.objects.shape import Shape
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment

# We do NOT redefine skills here – they already live in skill_code.
from skill_code import move, pick, place, rotate, pull

from video import init_video_writers, recording_step, recording_get_observation

# Utility that gives us a dictionary {object_or_pos_name: (x,y,z)}
from object_positions import get_object_positions


def run_skeleton_task() -> None:
    """Entry point for the combined-domain Drawer + Disposal task.
    The run consists of:
      1) Environment setup / video wiring.
      2)  Exploration phase whose purpose is to discover which workspace
          pose satisfies the hidden predicate (is-anchor-pos ?p ?d).
      3)  (Optional)  Task logic that would follow once the anchor pose is
          known.  Since an oracle plan is not supplied, after discovery we
          simply print the result and exit.
    """
    print("===== Starting Skeleton Task =====")

    # ---------------------------------------------------------
    # 0.  Environment initialisation
    # ---------------------------------------------------------
    env, task = setup_environment()
    try:
        # Reset the task to the initial demonstration state
        descriptions, obs = task.reset()

        # Initialise optional video capture utilities
        init_video_writers(obs)

        # Wrap task.step / task.get_observation with recorder versions
        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)

        # ---------------------------------------------------------
        # 1.  Retrieve all known positions in the scene
        # ---------------------------------------------------------
        positions = get_object_positions()   # e.g.  {'drawer_anchor': (x,y,z), ...}
        if not positions:
            print("[Warning] object_positions returned an empty dictionary.")
        else:
            print(f"[Info] Retrieved {len(positions)} named positions from scene.")

        # ---------------------------------------------------------
        # 2.  Heuristic filtering – keep only pose names that might be
        #     an anchor for a drawer.  We retain any position whose name
        #     includes the substring “drawer” or “anchor”.
        # ---------------------------------------------------------
        anchor_candidates = {}
        for name, pos in positions.items():
            lowered = name.lower()
            if ("drawer" in lowered) or ("anchor" in lowered):
                anchor_candidates[name] = pos

        if not anchor_candidates:
            # As a fall-back keep *all* positions; maybe they just weren’t
            # named using a helpful convention.
            anchor_candidates = positions.copy()

        print(f"[Exploration] Number of anchor pose hypotheses: "
              f"{len(anchor_candidates)}")

        # ---------------------------------------------------------
        # 3.  Exploration loop – try each candidate to see which one
        #     meets the missing predicate’s semantics.
        #     Strategy:
        #       • move() the gripper to the candidate pose
        #       • rotate() the gripper so that it is aligned for a pull
        #       • attempt pull();  if the drawer opens then by the domain
        #         definition we are *necessarily* at the anchor position
        #         (because pull’s preconditions require holding-drawer,
        #          which in turn needs is-anchor-pos to have been true).
        #       • Record the first pose that causes success as the anchor.
        # ---------------------------------------------------------
        discovered_anchor_name = None
        discovered_anchor_pos = None

        # Generic fixed parameters for motion; values are intentionally
        # conservative so that they work for most RLBench tasks.
        APPROACH_DIST = 0.08      # [m]
        MAX_STEPS     = 150
        THRESH_DIST   = 0.01
        TIMEOUT_S     = 5.0

        for name, pos in anchor_candidates.items():
            print(f"[Exploration] Trying candidate pose “{name}” – {pos}")

            try:
                # 3-a) Move to the pose
                obs, rew, done = move(
                    env,
                    task,
                    target_pos     = pos,
                    approach_distance = APPROACH_DIST,
                    max_steps      = MAX_STEPS,
                    threshold      = THRESH_DIST,
                    approach_axis  = 'z',      # come from above
                    timeout        = TIMEOUT_S
                )
            except Exception as e:
                print(f"    [move] failed for {name}: {e}")
                continue

            # 3-b) Ensure the gripper is rotated 90 deg (domain demands)
            try:
                obs, rew, done = rotate(
                    env,
                    task,
                    angle=90.0      # assume the skill understands degrees
                )
            except Exception as e:
                print(f"    [rotate] failed for {name}: {e}")
                continue

            # 3-c) Try a pull
            drawer_opened = False
            try:
                obs, rew, done = pull(
                    env,
                    task
                )
                # Some implementations return a flag in “done” signalling
                # success – fallback to reward > 0 if “done” is False.
                if done or (rew is not None and rew > 0):
                    drawer_opened = True
            except Exception as e:
                # Pull can raise if pre-conditions not met; that’s expected
                print(f"    [pull] failed for {name}: {e}")

            if drawer_opened:
                discovered_anchor_name = name
                discovered_anchor_pos  = pos
                print(f"[Success] Drawer opened at pose “{name}”. "
                      f"Predicate (is-anchor-pos {name}) discovered!")
                break
            else:
                print(f"[Exploration] Pose “{name}” is *not* the anchor.")

        if discovered_anchor_name is None:
            print("[Result] No anchor pose identified in exploration phase.")
        else:
            print(f"[Result] Anchor pose = “{discovered_anchor_name}” "
                  f"@ {discovered_anchor_pos}")

        # ---------------------------------------------------------
        # 4.  Placeholder for the remainder of the task plan.
        # ---------------------------------------------------------
        #
        # At this point, one would normally:
        #   • pick the drawer handle,
        #   • pull completely,
        #   • pick a target object,
        #   • place it inside the drawer, etc.
        #
        # Those steps are highly task-specific and cannot be filled in
        # without an oracle plan.  We therefore finish after exploration.

    finally:
        # ---------------------------------------------------------
        # 5.  Always make sure we tidy up, even on crash / Ctrl-C
        # ---------------------------------------------------------
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()