# run_skeleton_task.py (Completed)

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

# NOTE: every skill *already* exists in skill_code and must **not** be re-defined.
# We simply import everything that is available there.
from skill_code import *

from video import init_video_writers, recording_step, recording_get_observation

# get_object_positions() should return a dict:  { 'name': (x, y, z), ... }
from object_positions import get_object_positions


def _deduce_missing_predicate(feedback: str) -> str:
    """
    A tiny helper that ‘parses’ planner feedback strings and extracts the
    missing predicate name.  The current feedback mechanism simply returns
    the missing symbol, so a strip/clean is enough.
    """
    if not feedback:
        return ""
    # The planner returns e.g.  «handempty» (possibly with newline / spaces)
    pred = feedback.strip().lower()
    # Remove eventual parentheses or other junk
    return pred.replace("(", "").replace(")", "")


def _ensure_handempty(env, task, current_obs):
    """
    Convenience wrapper that tries to guarantee ‘handempty’ before
    executing a new pick.  If the gripper is currently holding something,
    we gently place it straight down at the current tool centre point.
    
    This uses only the predefined ‘place’ skill.
    """
    # Many tasks expose “gripper_openness”; if negative we cannot know the
    # state, so we conservatively assume the hand is *not* empty.
    openness = getattr(current_obs, "gripper_openness", -1.0)
    # We treat “fully open” (>= 0.04 by convention) as empty
    HAND_OPEN_THRESHOLD = 0.04
    is_empty = openness >= HAND_OPEN_THRESHOLD

    if is_empty:
        return current_obs, False  # Nothing to do

    # ------------------------------------------------------------------
    # We do hold something → we will perform a safety “place” straight
    # below the gripper (drop-in-place).  For that we build a tiny target
    # point that is 4 cm below the current TCP.
    # ------------------------------------------------------------------
    tcp_pos = current_obs.gripper_pose[:3]
    drop_pos = tcp_pos.copy()
    drop_pos[2] -= 0.04  # drop 4 cm

    print(f"[Exploration] handempty missing → executing emergency place at {drop_pos}")

    try:
        obs, reward, done = place(
            env,
            task,
            target_pos=drop_pos,
            approach_distance=0.05,
            max_steps=60,
            threshold=0.01,
            approach_axis="z",
            timeout=5.0,
        )
    except Exception as exc:
        # Even if the place fails, we log and move on; the mission may still
        # succeed once the gripper is force-opened by the simulator.
        print(f"[Exploration] Warning: emergency place failed – {exc}")
        obs = task.get_observation()
        reward, done = 0.0, False

    return obs, done


def _find_target_container(position_dict):
    """
    A naive heuristic that attempts to discover the ‘goal / disposal /
    target’ location among all known object positions.  Conventionally
    many RLBench scenes contain an object whose name includes one of:
    ‘bin’, ‘basket’, ‘target’, or ‘drawer’.  We simply search for those.
    """
    KEYWORDS = ("bin", "basket", "target", "drawer")
    for name, pos in position_dict.items():
        lname = name.lower()
        if any(k in lname for k in KEYWORDS):
            return name, pos
    # Fallback: just return the first element (not ideal but prevents crash)
    if position_dict:
        first_name = next(iter(position_dict))
        return first_name, position_dict[first_name]
    return None, None


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

    # === Environment Setup ===
    env, task = setup_environment()
    try:
        # Reset the task to its initial state
        descriptions, obs = task.reset()

        # (Optional) Initialize video writers for capturing your simulation
        init_video_writers(obs)

        # Wrap the task steps for recording
        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)

        # === Retrieve Object Positions ===
        positions = get_object_positions()
        print(f"[Info] Discovered {len(positions)} scene objects via get_object_positions().")

        # Decide which object is the “goal container” for placement
        target_container_name, target_container_pos = _find_target_container(positions)
        if target_container_name is None:
            raise RuntimeError("Could not determine a target container within the scene!")

        print(f"[Planner] Chosen target container: {target_container_name} @ {target_container_pos}")

        # ---------------  EXPLORATION PHASE  ----------------
        # We received planner feedback that the predicate ‘handempty’ was
        # missing.  Therefore we execute a *short exploratory routine* that
        # tries to guarantee the pre-conditions required by any upcoming
        # ‘pick’.
        missing_predicate = _deduce_missing_predicate("handempty")
        if missing_predicate == "handempty":
            obs, premature_done = _ensure_handempty(env, task, obs)
            if premature_done:
                print("[Exploration] Task ended early during handempty handling.")
                return

        # ---------------  MAIN TASK LOGIC  -----------------
        #
        # The plan is deliberately *very* generic: iterate over every
        # object that is *not* the target container, pick it up, and place
        # it into the container.  This respects the PDDL domain (pick +
        # place) and satisfies the “handempty” pre-condition between picks.
        #
        for obj_name, obj_pos in positions.items():
            if obj_name == target_container_name:
                continue  # skip the container itself

            print(f"[Task] Processing '{obj_name}' at {obj_pos}")

            # Ensure we start each pick with an empty hand
            obs, premature_done = _ensure_handempty(env, task, obs)
            if premature_done:
                print("[Task] Environment signalled DONE while clearing hand.")
                break

            # ---------- MOVE & PICK ----------
            try:
                obs, reward, done = move(
                    env,
                    task,
                    target_pos=obj_pos,
                    max_steps=150,
                    threshold=0.02,
                    approach_distance=0.15,
                    approach_axis="z",
                    timeout=7.0,
                )
                if done:
                    print("[Task] Episode finished during move.")
                    break

                obs, reward, done = pick(
                    env,
                    task,
                    target_pos=obj_pos,
                    approach_distance=0.05,
                    max_steps=100,
                    threshold=0.01,
                    approach_axis="z",
                    timeout=7.0,
                )
                if done:
                    print("[Task] Episode finished during pick.")
                    break
            except Exception as exc:
                print(f"[Task] Warning: failed to move/pick '{obj_name}' – {exc}")
                continue  # skip to next object

            # ---------- MOVE TO CONTAINER & PLACE ----------
            try:
                obs, reward, done = move(
                    env,
                    task,
                    target_pos=target_container_pos,
                    max_steps=150,
                    threshold=0.02,
                    approach_distance=0.18,
                    approach_axis="z",
                    timeout=7.0,
                )
                if done:
                    print("[Task] Episode finished while moving to container.")
                    break

                # Small z-offset so that the object is released *inside* the bin
                drop_pos = np.array(target_container_pos)
                drop_pos[2] += 0.05

                obs, reward, done = place(
                    env,
                    task,
                    target_pos=drop_pos,
                    approach_distance=0.05,
                    max_steps=100,
                    threshold=0.01,
                    approach_axis="z",
                    timeout=7.0,
                )
                if done:
                    print("[Task] Episode finished during place.")
                    break
            except Exception as exc:
                print(f"[Task] Warning: failed to move/place into container – {exc}")
                continue

        # --------------- OPTIONAL: DRAWER HANDLING ---------------
        #
        # If the scene *also* contains a drawer handle, we show how to pull
        # it open using the existing ‘rotate’ + ‘pull’ skills.  This block
        # is optional and will silently skip when the handle cannot be
        # located.
        drawer_handle_name = None
        for name in positions:
            if "handle" in name.lower() or "drawer" in name.lower():
                drawer_handle_name = name
                break

        if drawer_handle_name:
            print(f"[Drawer] Found potential handle '{drawer_handle_name}'; attempting to open it.")

            # Get current orientation, then rotate by 90 deg about tool-Z.
            current_quat = obs.gripper_pose[3:7]
            # A +90° rotation around z-axis in xyzw
            ninety_deg_quat = np.array([0.0, 0.0, np.sin(np.pi / 4), np.cos(np.pi / 4)])

            obs, reward, done = rotate(
                env,
                task,
                target_quat=ninety_deg_quat,
                max_steps=120,
                threshold=0.05,
                timeout=8.0,
            )
            if not done:
                try:
                    obs, reward, done = pull(
                        env,
                        task,
                        force=1.0,          # exemplar argument – signature may differ
                        max_steps=120,
                        timeout=8.0,
                    )
                except Exception as exc:
                    print(f"[Drawer] Pull failed or skill signature mismatch – {exc}")

        # ----------------------------------------------------------
        print("===== Task execution loop completed! =====")

    finally:
        # Always ensure the environment is properly shut down
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()