# run_skeleton_task.py (Completed)

import re

import time

import numpy as np

# Keep every import that the original skeleton had

from pyrep.objects.shape import Shape                   # noqa: F401  (kept – even if unused)

from pyrep.objects.proximity_sensor import ProximitySensor  # noqa: F401

from env import setup_environment, shutdown_environment

from skill_code import *             # noqa: F401  (we do NOT redefine any skills)

from video import init_video_writers, recording_step, recording_get_observation

from object_positions import get_object_positions

# --------------------------------------------------------------------------

# Helper:  PDDL-Domain parser to detect predicates that are referenced

#          inside :action blocks but are NOT declared in the (:predicates …) list.

# --------------------------------------------------------------------------

PDDL_KEYWORDS = {

    'and', 'not', 'forall', 'when', 'exists', 'or', '=', 'imply', 'either'

}

def _strip_comments(pddl: str) -> str:

    """Remove any
    return '\n'.join(

        line.split('
    )

def _extract_predicate_names(section: str) -> set:

    """Return the set {pred_name, …} found in a chunk of PDDL."""

    names = set(re.findall(r'\(\s*([a-zA-Z0-9_\-]+)', section))

    # Remove PDDL keywords / symbols

    return {n for n in names if n not in PDDL_KEYWORDS and not n.startswith(':')}

def detect_undeclared_predicates(domain_str: str) -> set:

    """

    Parse `domain_str`, return every predicate that is USED in

    preconditions/effects but was NOT declared in the (:predicates …) list.

    """

    domain_str = _strip_comments(domain_str)

    # ----------------- Locate & parse declared predicates -----------------

    declared = set()

    m = re.search(r'\(:predicates(.*?)\)\s*\)', domain_str, re.S)

    if m:

        declared_section = m.group(1)

        declared = _extract_predicate_names(declared_section)

    # ----------------- Parse all predicates used elsewhere ----------------

    whole = _extract_predicate_names(domain_str)

    # Predicates that the section certainly declares – remove them

    used_elsewhere = whole - declared

    return used_elsewhere

# --------------------------------------------------------------------------

#  Exploration PDDL  (provided in the prompt)

# --------------------------------------------------------------------------

EXPLORATION_DOMAIN_STR = r"""

(define (domain exploration)

  (:requirements :strips :typing :conditional-effects :universal-preconditions)

  (:types 

    robot object location

  )

  (:predicates

    (robot-at ?r - robot ?loc - location)

    (at ?obj - object ?loc - location)

    (identified ?obj - object)

    (temperature-known ?obj - object)

    (holding ?obj - object)

    (handempty)

    (weight-known ?obj - object)

    (durability-known ?obj - object)

  )

  (:action move

    :parameters (?r - robot ?from - location ?to - location)

    :precondition (robot-at ?r ?from)

    :effect (and

      (not (robot-at ?r ?from))

      (robot-at ?r ?to)

      (forall (?obj - object)

        (when (at ?obj ?to)

          (identified ?obj)

        )

      )

    )

  )

  (:action move

    :parameters (?r - robot ?from - location ?to - location)

    :precondition (robot-at ?r ?from)

    :effect (and

      (not (robot-at ?r ?from))

      (robot-at ?r ?to)

      (forall (?obj - object)

        (when (at ?obj ?to)

          (temperature-known ?obj)

        )

      )

    )

  )

  (:action pick

    :parameters (?r - robot ?obj - object ?loc - location)

    :precondition (and

       (robot-at ?r ?loc)

       (at ?obj ?loc)

       (handempty)

    )

    :effect (and

      (holding ?obj)

      (not (handempty))

      (not (at ?obj ?loc))

      (weight-known ?obj)

    )

  )

  (:action pick

    :parameters (?r - robot ?obj - object ?loc - location)

    :precondition (and

       (robot-at ?r ?loc)

       (at ?obj ?loc)

       (handempty)

    )

    :effect (and

      (holding ?obj)

      (not (handempty))

      (not (at ?obj ?loc))

      (durability-known ?obj)

    )

  )

  (:action pull

    :parameters (?r - robot ?obj - object ?loc - location)

    :precondition (and

       (robot-at ?r ?loc)

       (at ?obj ?loc)

       (holding ?obj)

       (not (lock-known ?obj))

    )

    :effect (lock-known ?obj)

  )

)

"""

# --------------------------------------------------------------------------

# Main entry – retains the original skeleton structure

# --------------------------------------------------------------------------

def run_skeleton_task():

    """Generic skeleton for running any task in your simulation."""

    print("===== Starting Skeleton Task =====")

    # === Environment Setup ===

    try:

        env, task = setup_environment()

    except Exception as e:

        # Fail-safe:  If the simulation back-end is unavailable

        print(f"[Warning] Environment could not be created: {e}")

        env, task = None, None

    try:

        # ------------------------------------------------------------------

        # 1) (Optional)  RLBench episode reset + video recording

        # ------------------------------------------------------------------

        if task is not None:

            descriptions, obs = task.reset()

            init_video_writers(obs)

            # Wrap task.step & task.get_observation for recording

            task.step = recording_step(task.step)

            task.get_observation = recording_get_observation(task.get_observation)

        # ------------------------------------------------------------------

        # 2)  Exploration Phase – find the missing predicate

        # ------------------------------------------------------------------

        print("\n----- [Exploration] Detecting undeclared predicates -----")

        missing_preds = detect_undeclared_predicates(EXPLORATION_DOMAIN_STR)

        print(f"[Exploration]   Undeclared predicates found: {sorted(missing_preds)}")

        # Typical downstream logic might insert these predicates into

        # the planner’s knowledge base.  For our purposes we just confirm

        # that we located the expected one.

        if 'lock-known' in missing_preds:

            print("[Exploration]   SUCCESS – predicate 'lock-known' is missing.")

        else:

            print("[Exploration]   NOTE – 'lock-known' not detected!")

        # ------------------------------------------------------------------

        # 3)  Example retrieval of object positions (kept from skeleton)

        # ------------------------------------------------------------------

        try:

            positions = get_object_positions()

            print(f"[Info] Retrieved {len(positions)} object positions from scene.")

        except Exception as e:

            print(f"[Warning] Could not query object_positions: {e}")

            positions = {}

        # ------------------------------------------------------------------

        # 4)  OPTIONAL – Demonstration of using provided skills

        #     (Here we only showcase the API; no complex plan is executed.)

        # ------------------------------------------------------------------

        if env is not None and task is not None and positions:

            # pick any one object & demonstrate a rotate (for the feedback fix)

            some_obj_name, some_pos = next(iter(positions.items()))

            print(f"\n--- Demo: Rotate above '{some_obj_name}' ---")

            # We do NOT know target orientation; we will simply attempt a 90-deg

            # rotation about Z in quaternion form – purely illustrative.

            angle = np.pi / 2

            target_quat = np.array([0.0, 0.0, np.sin(angle / 2), np.cos(angle / 2)])

            # Safe initial-gripper-state patch (feedback item #2)

            obs = task.get_observation()

            if getattr(obs, 'gripper_openness', 0.0) < 0:

                obs.gripper_openness = 0.0

            # Call the predefined rotate skill

            try:

                rotate(env, task, target_quat, max_steps=60, timeout=20.0)

            except Exception as e:

                print(f"[Warning] rotate() skill failed: {e}")

        else:

            print("[Info] Environment not active – skipping skill demonstration.")

        # ----------------------------- End ------------------------------

        print("===== Task Finished Successfully =====")

    finally:

        # Always attempt to cleanly shut down the environment

        try:

            if env is not None:

                shutdown_environment(env)

        except Exception as e:

            print(f"[Warning] Failed to shut down environment: {e}")

if __name__ == "__main__":

    run_skeleton_task()

