# run_skeleton_task.py (Filled-in version)

import re

import numpy as np

from pyrep.objects.shape import Shape

from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment

# All primitives (pick, place, move, rotate, pull, …) are imported here.

# We do NOT redefine them – we only call them if/when they are required.

from skill_code import *      # noqa: F401,F403

from video import init_video_writers, recording_step, recording_get_observation

from object_positions import get_object_positions

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

# A tiny helper that searches a PDDL domain description and reports

# every predicate that is USED in any :precondition / :effect but

# NOT declared inside the (:predicates …) list.

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

def find_missing_predicates(domain_text: str):

    # Strip comments

    clean = re.sub(r";.*", "", domain_text)

    # 1) Collect declared predicates

    declared_section = re.search(r"\(:predicates([^)]*)\)", clean, flags=re.DOTALL)

    declared_predicates = set()

    if declared_section:

        raw = declared_section.group(1)

        # remove parentheses around each declaration and split

        for token in re.findall(r"\(\s*([^\s)]+)", raw):

            declared_predicates.add(token.strip())

    # 2) Collect predicates that appear in actions

    used_predicates = set()

    # We simply regard anything that looks like "(name" as a predicate usage

    for token in re.findall(r"\(\s*([^\s()]+)", clean):

        # filter out PDDL keywords / connectives

        if token.lower() in {

            "and", "or", "not", "forall", "exists",

            "when", "imply", "=", "<", ">", "<=", ">=",

            ":action", ":parameters", ":precondition", ":effect",

            ":types", ":requirements"

        }:

            continue

        used_predicates.add(token.strip())

    # 3) Missing = used – declared

    missing = sorted(list(used_predicates - declared_predicates))

    return missing, sorted(list(declared_predicates)), sorted(list(used_predicates))

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 the task to its initial state

        descriptions, obs = task.reset()

        # (Optional) Initialize video writers for capturing your simulation

        init_video_writers(obs)

        # Wrap the task steps for recording (if needed)

        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)

        # === Retrieve Object Positions ===

        positions = get_object_positions()

        print(f"[Info] Detected {len(positions)} objects from get_object_positions().")

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

        #  Exploration Phase – Identify missing predicate(s)

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

        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)

          )

          (: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)

          )

        )

        """

        missing, declared, used = find_missing_predicates(exploration_domain_pddl)

        print("\n===== Exploration Report =====")

        print(f"Declared predicates : {declared}")

        print(f"Predicates used     : {used}")

        print(f"---\nMissing predicate(s): {missing}")

        print("================================\n")

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

        # Optional: Demonstrate one primitive (if an object exists)

        #           – We first verify that the chosen object is still present

        #             before executing the pick skill.

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

        if positions:

            # Select the first object in the dictionary

            first_key = list(positions.keys())[0]

            obj_info = positions[first_key]

            # The helper returns various formats; try to be permissive

            if isinstance(obj_info, dict) and "position" in obj_info:

                target_pos = np.array(obj_info["position"])

            elif isinstance(obj_info, (list, tuple, np.ndarray)):

                target_pos = np.array(obj_info)

            else:

                print(f"[Warn] Unable to parse position for {first_key}; skipping demo pick.")

                target_pos = None

            if target_pos is not None and len(target_pos) == 3:

                try:

                    print(f"[Demo] Attempting a pick on '{first_key}' at {target_pos} ...")

                    obs, reward, done = pick(

                        env,

                        task,

                        target_pos=target_pos,

                        approach_distance=0.15,

                        max_steps=120,

                        threshold=0.01,

                        approach_axis='z',

                        timeout=10.0

                    )

                    if done:

                        print("[Demo] Task reported completion during demo pick.")

                except Exception as e:

                    # We catch broad exceptions here because the demo pick is optional

                    print(f"[Error] Demo pick failed with exception: {e}")

            else:

                print("[Info] No valid position for demo pick – skipping.")

        # Nothing else to do – this script’s main purpose was to locate

        # the missing predicate and optionally exercise an existing skill.

    finally:

        # Always ensure the environment is properly shutdown

        shutdown_environment(env)

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

if __name__ == "__main__":

    run_skeleton_task()
