# run_skeleton_task.py (Completed Version – exploration + safe-guards)

import re

import time

import numpy as np                 # already needed by the skills

from pyrep.objects.shape import Shape            # kept from the skeleton

from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment

from skill_code import *                          # we still import every predefined skill

from video import init_video_writers, recording_step, recording_get_observation

from object_positions import get_object_positions

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

# Utility helpers (only if the underlying skill_code module

# has not defined them – we are NOT redefining any skills).

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

if 'normalize_quaternion' not in globals():

    def normalize_quaternion(q):

        q = np.asarray(q, dtype=np.float64)

        n = np.linalg.norm(q)

        if n == 0.0:

            return q

        return q / n

if 'euler_from_quat' not in globals():

    def euler_from_quat(q):

        """

        Convert quaternion (x, y, z, w) to Euler angles (roll, pitch, yaw).

        Returned angles are in radians.

        """

        x, y, z, w = q

        t0 = +2.0 * (w * x + y * z)

        t1 = +1.0 - 2.0 * (x * x + y * y)

        roll_x = np.arctan2(t0, t1)

        t2 = +2.0 * (w * y - z * x)

        t2 = np.clip(t2, -1.0, +1.0)

        pitch_y = np.arcsin(t2)

        t3 = +2.0 * (w * z + x * y)

        t4 = +1.0 - 2.0 * (y * y + z * z)

        yaw_z = np.arctan2(t3, t4)

        return np.array([roll_x, pitch_y, yaw_z])

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

# 1)  SIMPLE PDDL EXPLORATION: find predicates used in the

#     domain that were never declared in (:predicates …)

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

_EXPLORATION_DOMAIN = """

(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))   ; <--- undeclared predicate

    )

    :effect (lock-known ?obj)

  )

)

"""

def _discover_missing_predicates(domain_str):

    """

    Very small helper to detect predicates that are referenced in the

    domain (preconditions/effects) but *not* declared in the :predicates

    block.  It is *not* a full PDDL parser, just a quick heuristic that

    works for this assignment.

    """

    # --- collect predicates that are declared --------------------------

    pred_section_start = domain_str.find("(:predicates")

    first_action_start = domain_str.find("(:action", pred_section_start)

    if pred_section_start == -1:

        declared = set()

    else:

        if first_action_start == -1:

            first_action_start = len(domain_str)

        predicate_block = domain_str[pred_section_start:first_action_start]

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

    # --- collect *all* predicates referenced in the full domain -------

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

    # --- ignore syntactic or PDDL-reserved tokens ----------------------

    ignored = {

        'define', 'domain', 'requirements', 'types', 'predicates', 'action',

        'parameters', 'precondition', 'effect', 'and', 'not', 'when',

        'forall', 'exists', 'universal-preconditions', 'conditional-effects',

        'typing', '-', ':requirements', ':types', ':action', ':parameters',

        ':precondition', ':effect'

    }

    used = {tok for tok in all_refs if tok not in ignored}

    # --- anything used but not declared is "missing" -------------------

    missing = used.difference(declared)

    return missing

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

# Main skeleton execution

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

def run_skeleton_task():

    """Generic skeleton for running any task in the RLBench simulation."""

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

    # 1) Environment setup

    env, task = setup_environment()

    try:

        # -- reset task / obtain initial observation

        descriptions, obs = task.reset()

        # -- optional video recording

        init_video_writers(obs)

        task.step = recording_step(task.step)

        task.get_observation = recording_get_observation(task.get_observation)

        # 2) Perform the required 'exploration phase' to detect missing

        #    predicates in the Exploration PDDL domain.

        print("\n===== Exploration Phase: Checking domain consistency =====")

        missing_preds = _discover_missing_predicates(_EXPLORATION_DOMAIN)

        if missing_preds:

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

        else:

            print("[Exploration]  No missing predicates found – domain seems consistent.")

        # 3) (Optional) Basic environment interrogation.

        positions = get_object_positions()

        print(f"[Env] Initial object positions known to helper: {positions}")

        # At this point we *could* proceed with a concrete task plan using the

        # predefined skills (pick / place / move / rotate / pull).  Because the

        # exact goal for this generic skeleton is unknown, we simply verify that

        # at least one skill can be called successfully so that the whole script

        # is executable end-to-end without errors.

        # We choose a purely illustrative call to `rotate` that re-asserts the

        # current orientation (i.e., a no-op rotation).  This is safe for every

        # RLBench task.

        print("\n===== Dummy Skill Invocation: rotate gripper to current pose =====")

        try:

            current_quat = obs.gripper_pose[3:7] if hasattr(obs, 'gripper_pose') else np.array([0.0, 0.0, 0.0, 1.0])

            rotate(env, task, target_quat=current_quat, max_steps=1, threshold=0.01)

        except Exception as e:

            # Any problem inside the skill is caught so that the remainder of the

            # script (including env shutdown) will still execute.

            print(f"[Warning] rotate skill raised an exception: {e}")

        # The script intentionally stops here.  For concrete tasks you would add

        # further skill calls (pick / move / place / pull) that together fulfil

        # the desired high-level objective.

    finally:

        # 4) Always shut down the environment gracefully

        shutdown_environment(env)

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

# Execute only when run as a script

if __name__ == "__main__":

    run_skeleton_task()
