# run_skeleton_task.py (Completed Implementation)

import re
import time

import numpy as np
from pyrep.objects.shape import Shape        # kept, but unused – do not remove
from pyrep.objects.proximity_sensor import ProximitySensor   # kept, but unused – do not remove

from env import setup_environment, shutdown_environment

# import every predefined skill exactly as delivered
from skill_code import rotate, pick, place, move, pull       # noqa: F401  (import-unused is OK)

from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions


# --------------------------------------------------------------------------- #
#                   Helper: predicate extraction & comparison                 #
# --------------------------------------------------------------------------- #
def _extract_predicates(domain_str: str):
    """Return the set of predicate names that appear in a PDDL domain string."""
    preds = set()
    # Locate the (:predicates … ) block(s)
    for block in re.findall(r'\(:predicates(.*?)\)', domain_str, flags=re.S):
        # Grab every “(predicate …” pattern inside the block
        for pred in re.findall(r'\(\s*([^\s()]+)', block):
            # Skip leading parentheses caused by nested structures like '(not'
            if pred.startswith('not') or pred.startswith('=' ) or pred.startswith('and') \
               or pred.startswith('forall') or pred.startswith('when'):
                continue
            preds.add(pred)
    return preds


def find_missing_predicates(combined_domain: str, exploration_domain: str):
    """Compare predicate sets and return those present only in exploration."""
    combined = _extract_predicates(combined_domain)
    exploration = _extract_predicates(exploration_domain)
    missing = exploration - combined
    return sorted(missing)


# --------------------------------------------------------------------------- #
#                            Embedded PDDL domain text                        #
# --------------------------------------------------------------------------- #
_COMBINED_DOMAIN_PDDL = """
(define (domain combined-domain)
  (:requirements :strips :typing :negative-preconditions :equality :disjunctive-preconditions)
  (:types object location drawer - object gripper - object position - object angle - object)
  (:predicates
    (at ?obj - object ?loc - location)
    (holding ?obj - object)
    (handempty)
    (is-locked ?d - drawer)
    (is-open ?d - drawer)
    (rotated ?g - gripper ?a - angle)
    (gripper-at ?g - gripper ?p - position)
    (holding-drawer ?g - gripper ?d - drawer)
    (is-side-pos ?p - position ?d - drawer)
    (is-anchor-pos ?p - position ?d - drawer)
  )
)
"""

_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)
    (lock-known ?obj - object)
  )
)
"""


# --------------------------------------------------------------------------- #
#                                Main routine                                 #
# --------------------------------------------------------------------------- #
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 environment/task
        descriptions, obs = task.reset()

        # Optionally prepare video recording
        init_video_writers(obs)
        task.step = recording_step(task.step)                  # wrap for video
        task.get_observation = recording_get_observation(task.get_observation)

        # ================== Retrieve Object Positions ================== #
        positions = get_object_positions()
        print(f"[Info] Known object positions: {list(positions.keys())}")

        # ================== Missing–Predicate Exploration =============== #
        missing_preds = find_missing_predicates(_COMBINED_DOMAIN_PDDL,
                                                _EXPLORATION_DOMAIN_PDDL)
        print(f"[Exploration] Predicates present only in the exploration domain → {missing_preds}")
        if missing_preds:
            print(f"[Exploration] Likely *missing* predicate(s) in combined domain: {missing_preds[0]}")

        # ================== Demo Skill Invocation  ====================== #
        # We perform a minimal, safe action with an identity quaternion to
        # validate that predefined skills work end-to-end.
        identity_quat = np.array([0, 0, 0, 1], dtype=np.float32)

        try:
            obs, reward, done = rotate(env, task, target_quat=identity_quat,
                                       max_steps=30, threshold=0.1, timeout=3.0)
            print(f"[Demo] rotate() finished | reward={reward} | done={done}")
        except Exception as exc:
            print(f"[Warning] rotate demo failed: {exc}")

        # ================== (Optional) Additional Logic ================= #
        # Here you could add domain-specific plans that utilise pick / place /
        # move / pull. Because the concrete task is unknown in this generic
        # skeleton, we limit ourselves to predicate discovery and a rotation
        # sanity-check above. Insert additional goal-oriented logic here if
        # task details become available.

        # small wait so that any video recorder flushes final frames
        time.sleep(0.5)

    finally:
        shutdown_environment(env)        # mandatory cleanup

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


if __name__ == "__main__":
    run_skeleton_task()
