# run_skeleton_task.py  (Completed Version)

import re
import time
import numpy as np
from pyrep.objects.shape import Shape
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment

# Import * AFTER env is imported so that skill modules can safely grab env-specific
# info on import if they need it.
from skill_code import *                     # noqa: F403

from video import (
    init_video_writers,
    recording_step,
    recording_get_observation,
)

from object_positions import get_object_positions


# ----------------------------------------------------------------------------- #
#                            Helper / Exploration Logic                         #
# ----------------------------------------------------------------------------- #
_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)
  )
)
"""


def _extract_predicates(pddl_src: str):
    """Return a set of predicate symbols that appear in a PDDL :predicates block."""
    preds = set()
    # Grab lines between (:predicates ... )
    pattern = r"\(:predicates([\s\S]*?)\)"
    match = re.search(pattern, pddl_src, re.MULTILINE)
    if not match:
        return preds
    block = match.group(1)
    for line in block.splitlines():
        line = line.strip()
        if not line or line.startswith(";"):
            continue
        # Each predicate line starts with '(' and ends with ')'  -> take first token
        if line[0] == "(":
            token = line[1:].split()[0]  # predicate name is first after '('
            preds.add(token)
    return preds


def run_skeleton_task():
    """Generic skeleton for running any task in the simulation with exploration."""

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

    # -------------------------------------------------------------------------- #
    #                (1)  Exploration – find missing predicate                  #
    # -------------------------------------------------------------------------- #
    print("[Exploration] Analysing domain predicates ...")
    predicates_combined = _extract_predicates(_COMBINED_DOMAIN_PDDL)
    predicates_exploration = _extract_predicates(_EXPLORATION_DOMAIN_PDDL)

    # Anything needed by the exploration domain but not in the task domain?
    missing_predicates = sorted(predicates_exploration - predicates_combined)
    if missing_predicates:
        print(f"[Exploration] Missing predicates discovered: {missing_predicates}")
    else:
        print("[Exploration] No missing predicates found when compared with the "
              "exploration domain.")

    # We received explicit feedback that `handempty` had caused issues previously.
    # Double-check its presence and warn the user if it is absent.
    if "handempty" not in predicates_combined:
        print("[Exploration] WARNING: predicate `handempty` is NOT present in the "
              "combined domain. This will break `pick` actions!")
    else:
        print("[Exploration] Predicate `handempty` confirmed present in the domain.")

    # -------------------------------------------------------------------------- #
    #                (2)  Environment setup and optional recording               #
    # -------------------------------------------------------------------------- #
    env, task = setup_environment()
    try:
        descriptions, obs = task.reset()

        # If video helpers are available we start recording automatically
        try:
            init_video_writers(obs)
            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)
        except Exception as err:
            # Recording is optional; we do not abort if it fails.
            print(f"[Video] Initialising recording failed (ignored): {err}")

        # ---------------------------------------------------------------------- #
        #            (3)  Retrieve object positions from helper module           #
        # ---------------------------------------------------------------------- #
        try:
            positions = get_object_positions()
            if not positions:
                print("[Info] No object positions returned from get_object_positions().")
        except Exception as err:
            print(f"[Warning] Could not retrieve object positions: {err}")
            positions = {}

        # ---------------------------------------------------------------------- #
        #                   (4)  Very simple demo plan / policy                  #
        # ---------------------------------------------------------------------- #
        #
        # We do NOT know the true evaluation task. Instead we show how to use
        # pre-existing skills safely.  The policy simply:
        #
        #   • selects the first object (if any) we can find via helper
        #   • moves the gripper 10 cm above that object (skill: move)
        #   • returns to the home position afterwards
        #
        # This demonstrates correct skill invocation and leaves plenty of room
        # for downstream task-specific logic to be inserted.

        # Choose an arbitrary “home” pose – here we take the current one at reset.
        home_pos = obs.gripper_pose[:3]

        # Execute the demo only if at least one object was reported.
        if positions:
            # Grab the first (obj_name, position) pair
            obj_name, obj_pos = next(iter(positions.items()))
            obj_pos_np = np.array(obj_pos, dtype=np.float32)

            print(f"[Demo] Moving above first object '{obj_name}' at {obj_pos_np}.")

            # Target 10 cm above the object's reported position
            approach_offset = np.array([0.0, 0.0, 0.10], dtype=np.float32)
            target_above = obj_pos_np + approach_offset

            # ------------------------------------------------------------------ #
            #                             Skill: move                             #
            # ------------------------------------------------------------------ #
            try:
                print("[Demo] Executing `move` skill to target position ...")
                obs, reward, done = move(env, task, target_above)        # noqa: F405
                if done:
                    print("[Demo] Task reported completion whilst moving – ending early.")
                    return
            except Exception as err:
                print(f"[Error] move skill failed: {err}")

            # Potential area to add more skills such as pick / place here.
            # We consciously do NOT perform pick/place – many benchmark tasks
            # demand pulling handles or rotating drawers etc. with higher-level
            # supervision.  The demo therefore ends by returning to the home pose.

            print("[Demo] Returning to home position.")
            try:
                obs, reward, done = move(env, task, home_pos)            # noqa: F405
            except Exception as err:
                print(f"[Error] move-to-home failed: {err}")

        else:
            print("[Demo] No objects available to demonstrate movement.")

        print("[Task] Finished demo run successfully.")

    finally:
        # Ensure simulator shuts down even when exceptions occur.
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()
