# run_skeleton_task.py  (Completed Version)

import re

import time

import numpy as np

# --- RL-Bench / simulation helpers (provided externally) --------------------

from env import setup_environment, shutdown_environment

from video import init_video_writers, recording_step, recording_get_observation

from object_positions import get_object_positions

# --- Low-level skill primitives (already implemented in skill_code) ---------

# NOTE:  we do *not* redeclare or modify these – they are simply imported.

from skill_code import pick, place, move, rotate, pull          # noqa: F401

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

#  Utility:  very small PDDL “parser”

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

KEYWORDS = {

    'and', 'not', 'when', 'forall', 'exists',

    '=', '>', '<', '<=', '>=', '+', '-', '*', '/', 'imply'

}

def _extract_predicate_block(pddl_txt: str) -> str:

    """Return the string that sits inside the top-level (:predicates …)."""

    pattern = r'\(:predicates([^)]*)\)'

    m = re.search(pattern, pddl_txt, re.S | re.I)

    return m.group(1) if m else ''

def _tokenise_predicate_block(block: str):

    """Return the set of predicate *names* declared in the block."""

    return {

        token for token in re.findall(r'\(\s*([a-zA-Z0-9\-]+)', block)

        if token not in KEYWORDS

    }

def _tokenise_all_predicates(pddl_txt: str):

    """Return every predicate symbol that appears anywhere in the domain."""

    return {

        token for token in re.findall(r'\(\s*([a-zA-Z0-9\-]+)', pddl_txt)

        if token not in KEYWORDS

    }

def find_missing_predicates(domain_pddl: str) -> set:

    """Compare declared predicates with actually referenced ones."""

    declared = _tokenise_predicate_block(_extract_predicate_block(domain_pddl))

    referenced = _tokenise_all_predicates(domain_pddl)

    return referenced - declared

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

#  A *tiny* exploration domain string – identical to the one in the prompt

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

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

  )

)

"""

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

#  Main task runner

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

def run_skeleton_task():

    """

    Generic entry-point for *any* task.

    Adds:

      1)   An “exploration” phase that scans the given PDDL domain

           to work out what predicate is missing.

      2)   A minimal demonstration of calling the available

           low-level skills – guarded by safe existence checks.

    """

    print("==========  RLBench Skeleton Task – START ==========")

    # 1)  First, perform the “exploration” of the PDDL text

    print("\n--- Exploration phase: looking for undeclared predicates ---")

    missing = find_missing_predicates(EXPLORATION_PDDL)

    if missing:

        print(f"[Exploration] Missing predicate(s) detected  ➜  {sorted(missing)}")

    else:

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

    # 2)  Set-up simulation

    env, task = setup_environment()

    try:

        # Reset the task to a known initial state

        descriptions, obs = task.reset()

        print(f"[Env] Task description: {' / '.join(descriptions)}")

        # (Optional) initialise video capture

        init_video_writers(obs)

        task.step = recording_step(task.step)

        task.get_observation = recording_get_observation(task.get_observation)

        # 3)  Retrieve object positions (provided externally)

        positions = get_object_positions()          # Dict[str, Tuple[float, float, float]]

        print(f"[Env] Objects detected from helper module: {list(positions.keys())}")

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

        # 4)  Mini-Demo plan

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

        #

        #     The exact task for the combined domain is not specified

        #     in the prompt.  Instead, we provide a *safe* illustrative

        #     sequence that uses only the officially allowed skills

        #     ('move', 'pick', 'place', 'rotate', 'pull').

        #

        #     Each call is wrapped in try/except so that the script

        #     continues gracefully even if the concrete task does not

        #     contain the referenced objects or the skill raises.

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

        # Helper: choose arbitrary objects (if they exist)

        candidate_objs = [k for k in positions.keys() if 'drawer' in k or 'handle' in k]

        if candidate_objs:

            target_obj = candidate_objs[0]

            target_pos = positions[target_obj]

            print(f"[Plan] Selected drawer-like object: {target_obj} @ {target_pos}")

        else:

            # fallback to anything

            if positions:

                target_obj, target_pos = next(iter(positions.items()))

                print(f"[Plan] Fallback object: {target_obj} @ {target_pos}")

            else:

                target_obj, target_pos = None, None

                print("[Plan] WARNING: No objects found – skipping manipulation demo.")

        # 4-A)  MOVE (if supported in your skill code – signature may vary)

        #       We assume skill_code.move accepts: (env, task, target_pos, …)

        if target_pos is not None:

            try:

                print("[Plan] move(): approaching the object")

                # Generic placeholder arguments; adapt to the real signature if needed

                move(

                    env,

                    task,

                    target_pos=target_pos,

                    approach_distance=0.20,

                    max_steps=120,

                    threshold=0.01,

                    approach_axis='z',

                    timeout=10.0

                )

            except Exception as e:

                print(f"[Plan] move() failed or is not applicable: {e}")

        # 4-B)  ROTATE  – demonstrate the rotation helper on the gripper itself

        try:

            print("[Plan] rotate(): setting gripper orientation to identity")

            identity_quat_xyzw = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32)

            rotate(

                env,

                task,

                target_quat=identity_quat_xyzw,

                max_steps=60,

                threshold=0.03,

                timeout=5.0

            )

        except Exception as e:

            print(f"[Plan] rotate() failed or is not applicable: {e}")

        # 4-C)  PULL  – if we managed to get hold of a drawer handle earlier

        try:

            if target_obj is not None:

                print("[Plan] pull(): attempting to pull / open object")

                pull(

                    env,

                    task,

                    obj_name=target_obj,

                    max_distance=0.10,

                    max_steps=80,

                    timeout=8.0

                )

        except Exception as e:

            print(f"[Plan] pull() failed or is not applicable: {e}")

        # 5)  Final step count / reward check  (optional)

        final_obs = task.get_observation()

        print(f"[Env] Final gripper pose  : {final_obs.gripper_pose}")

        # (If the task delivers a reward signal, one could print it here)

    finally:

        shutdown_environment(env)

        print("==========  RLBench Skeleton Task – END ==========")

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

#  Entry-point

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

if __name__ == "__main__":

    run_skeleton_task()
