# run_skeleton_task.py (Completed)

import inspect

import time

from typing import Dict, Any, Set, List

import numpy as np

from pyrep.objects.shape import Shape

from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment

from skill_code import *                                # noqa: F403 – we really do need the star-import here

from video import init_video_writers, recording_step, recording_get_observation

from object_positions import get_object_positions

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

# ---------------  Utility helpers (domain / exploration)  ------------------ #

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

_COMBINED_DOMAIN_PDDL = r"""

(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 = 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 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 _extract_predicates(pddl_text: str) -> Set[str]:

    """Return the set of predicate names explicitly declared in (:predicates …)."""

    decl_predicates: Set[str] = set()

    inside_block = False

    for line in pddl_text.splitlines():

        line = line.strip()

        if line.startswith("(:predicates"):

            inside_block = True

            line = line[len("(:predicates"):].strip()

        if inside_block:

            # collect predicate names on this line

            tokens = line.replace("(", " ").replace(")", " ").split()

            for tok in tokens:

                if tok and tok[0].isalpha() and not tok.startswith(':'):

                    decl_predicates.add(tok)

        if inside_block and line.endswith(")"):

            # end of predicates section (assume closed properly)

            if line.count("(") == line.count(")"):

                inside_block = False

    return decl_predicates

def _extract_all_symbols(pddl_text: str) -> Set[str]:

    """Return every symbol appearing immediately after a '(' character."""

    symbols: Set[str] = set()

    token = ''

    in_token = False

    for ch in pddl_text:

        if ch == '(':

            in_token = True

            token = ''

        elif ch in ' \t\n)':

            if in_token and token:

                symbols.add(token)

            in_token = False

        elif in_token:

            token += ch

    return symbols

def detect_missing_predicates(*pddl_strings: str) -> Set[str]:

    """Identify predicates used inside action definitions but missing from declarations."""

    declared: Set[str] = set()

    used: Set[str] = set()

    for text in pddl_strings:

        declared.update(_extract_predicates(text))

        used.update(_extract_all_symbols(text))

    # Remove :requirements, :action, :parameters, etc.

    non_predicate_keywords = {

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

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

        'when', 'forall', 'typing', 'strips', 'negative-preconditions',

        'conditional-effects', 'universal-preconditions', 'equality',

        'disjunctive-preconditions'

    }

    used -= non_predicate_keywords

    # Predicates start with lower-case letter; filter anything else

    used = {u for u in used if u.islower()}

    return used - declared

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

# -----------------------------  Skill helper  ------------------------------ #

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

def _call_skill(skill_name: str, env, task, **kwargs):

    """

    Generic skill dispatcher that introspects the signature of the skill

    function and feeds it only the arguments it expects.

    """

    try:

        skill_func = globals()[skill_name]  # noqa: F405  (imported via star)

    except KeyError as e:

        raise RuntimeError(f"Skill '{skill_name}' not found in skill_code") from e

    sig = inspect.signature(skill_func)

    bound_kwargs: Dict[str, Any] = {}

    for param in sig.parameters.values():

        if param.name == 'env':

            bound_kwargs['env'] = env

        elif param.name == 'task':

            bound_kwargs['task'] = task

        elif param.name in kwargs:

            bound_kwargs[param.name] = kwargs[param.name]

        elif param.default is inspect.Parameter.empty and param.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD:

            # Required parameter not provided

            raise ValueError(f"Missing required parameter '{param.name}' for skill '{skill_name}'")

    return skill_func(**bound_kwargs)

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

# -------------------------  Main task execution  --------------------------- #

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

def _select_handle_object(positions: Dict[str, np.ndarray]) -> str:

    """

    Heuristically choose an object that looks like a drawer handle.

    Simply prioritises names that contain 'handle', otherwise returns

    the first key (deterministic).

    """

    for name in positions.keys():

        if 'handle' in name.lower():

            return name

    # Fallback: first available object

    return next(iter(positions.keys()))

def run_skeleton_task():

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

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

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

    # 1)  Environment Setup

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

    env, task = setup_environment()

    try:

        descriptions, obs = task.reset()

        # Optionally start video recording

        init_video_writers(obs)

        task.step = recording_step(task.step)

        task.get_observation = recording_get_observation(task.get_observation)

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

        # 2)  PDDL Missing-Predicate Exploration

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

        missing_predicates = detect_missing_predicates(

            _COMBINED_DOMAIN_PDDL, _EXPLORATION_DOMAIN_PDDL

        )

        if missing_predicates:

            print("[Exploration] Detected predicates referenced but not declared:")

            for mp in sorted(missing_predicates):

                print(f"  • {mp}")

        else:

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

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

        # 3)  Retrieve Object Positions from the scene

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

        positions = get_object_positions()  # Expected: Dict[str, np.ndarray]

        # Basic consistency check for feedback (#1 incorrect declaration)

        if 'bottom_anchor_pos' not in positions:

            print("[Feedback] 'bottom_anchor_pos' not found in object list. "

                  "This aligns with feedback – continuing without it.")

        if not positions:

            print("[Warning] No object positions returned; nothing to do.")

            return

        target_handle_name = _select_handle_object(positions)

        target_pos = positions[target_handle_name]

        print(f"[Task] Targeting handle object: {target_handle_name} at {target_pos}")

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

        # 4)  High-level Plan Execution

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

        done = False

        reward = 0.0

        # 4-1  Move robot gripper close to handle

        try:

            obs, reward, done = _call_skill(

                'move', env, task,

                target_pos=target_pos,

                approach_distance=0.15,

                threshold=0.01,

                timeout=10.0

            )

        except Exception as e:

            print(f"[Error] move skill failed: {e}")

            return

        if done:

            print("[Task] Task ended unexpectedly after move.")

            return

        # 4-2  Rotate gripper 90° around its approach axis (satisfy domain 'rotated ?g ninety_deg')

        try:

            ninety_deg_quat = np.array([0.0, 0.0, np.sin(np.pi / 4), np.cos(np.pi / 4)])  # z-axis 90°

            obs, reward, done = _call_skill(

                'rotate', env, task,

                target_quat=ninety_deg_quat

            )

        except Exception as e:

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

            return

        # 4-3  Pick the drawer handle

        try:

            obs, reward, done = _call_skill(

                'pick', env, task,

                target_pos=target_pos,

                approach_distance=0.02,

                threshold=0.005,

                timeout=10.0

            )

        except Exception as e:

            print(f"[Error] pick skill failed: {e}")

            return

        if done:

            print("[Task] Task ended unexpectedly after pick.")

            return

        # 4-4  Pull the handle (open drawer)

        try:

            obs, reward, done = _call_skill('pull', env, task)

        except Exception as e:

            print(f"[Error] pull skill failed: {e}")

            return

        if done:

            print("[Task] Completed drawer-pulling task!")

        else:

            print("[Task] Pull executed; task may continue depending on scenario.")

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

        # 5)  Optional: Place the handle back / cleanup

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

        try:

            obs, reward, done = _call_skill(

                'place', env, task,

                target_pos=target_pos  # put it back where it was

            )

        except Exception:

            # Place may not be strictly necessary – ignore failures here

            pass

    finally:

        # Ensure environment shuts down cleanly

        shutdown_environment(env)

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

if __name__ == "__main__":

    run_skeleton_task()
