# run_skeleton_task.py (Completed)

import re

import time

import numpy as np

from pyrep.objects.shape import Shape                # Do not remove (required by RLBench)

from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment

from skill_code import *                             # You should NOT redefine primitives

from video import init_video_writers, recording_step, recording_get_observation

from object_positions import get_object_positions    # Helper for querying the scene

EXPLORATION_DOMAIN_PDDL = """

(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)

  )

)

"""

def _strip_pddl_comments(text: str) -> str:

    """Remove semicolon style PDDL comments."""

    return "\n".join(line for line in text.splitlines() if not line.strip().startswith(";"))

def _collect_predicates_from_block(block: str) -> set:

    """Return all predicate names appearing after an opening parenthesis."""

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

    return set(tokens)

def find_missing_predicates(domain_pddl: str) -> set:

    """

    Very small-scale PDDL checker:

    • parse predicates declared in (:predicates …)

    • parse all predicate tokens used inside :precondition / :effect blocks

    • return those that are used but never declared

    """

    text = _strip_pddl_comments(domain_pddl)

    # ---------------- Collect declared predicates ----------------

    declared_predicates = set()

    predicates_block_match = re.search(r'\(:predicates(.*?)\)\s*[\)\(]', text, flags=re.DOTALL)

    if predicates_block_match:

        declared_predicates = _collect_predicates_from_block(predicates_block_match.group(1))

    # ---------------- Collect used predicates -------------------

    used_predicates = set()

    # Consider everything that looks like "(predicate …" inside actions

    for action_block in re.findall(r'\(:action[\s\S]*?\)', text, flags=re.DOTALL):

        used_predicates |= _collect_predicates_from_block(action_block)

    # Filter out PDDL keywords and equality

    pddl_keywords = {"and", "not", "when", "forall", "exists", "=", ">" , "<", ">=", "<=", "increase", "decrease"}

    used_predicates -= pddl_keywords

    # Any used predicate that is not declared is missing.

    missing = used_predicates - declared_predicates

    return missing

def run_skeleton_task():

    """

    Generic skeleton for running any task in your simulation

    (augmented with a quick ‘exploration domain’ sanity-check step

    to identify a missing predicate as required by feedback).

    """

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

    # 1) First, automatically detect missing predicates in the exploration domain

    print("----- Exploration Phase: Checking Exploration Domain for Missing Predicates -----")

    missing_preds = find_missing_predicates(EXPLORATION_DOMAIN_PDDL)

    if missing_preds:

        print(f"[Exploration] Missing predicate(s) detected in the exploration domain: {sorted(missing_preds)}")

    else:

        print("[Exploration] No missing predicates detected!")

    # 2) Set-up the RLBench environment

    env, task = setup_environment()

    try:

        descriptions, obs = task.reset()

        init_video_writers(obs)

        # Wrap step/observation so that any subsequent calls automatically record video

        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)

        # In this minimal example we will not execute an elaborate plan.

        # We simply wait for a short duration to capture an initial frame,

        # demonstrating that the environment is operational.

        time.sleep(1.0)

        # Optionally, you could query the scene and print positions:

        positions = get_object_positions()

        print(f"[Info] Known object positions: {positions.keys()}")

        # Example no-op loop to leave a few frames in the video

        for _ in range(3):

            # RLBench tasks accept a zero action for “do nothing”

            zero_action = np.zeros(env.action_shape)

            obs, reward, done = task.step(zero_action)

            if done:

                break

    finally:

        shutdown_environment(env)

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

if __name__ == "__main__":

    run_skeleton_task()
