# run_skeleton_task.py (Completed Version – generic, exploration-ready, skill-only)

import re

import time

import traceback

from typing import 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 *          # only pre-defined skills will be called

from video import init_video_writers, recording_step, recording_get_observation

from object_positions import get_object_positions

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

# 1) Utilities – domain parsing to discover missing predicates

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

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_declared_predicates(domain: str) -> Set[str]:

    declared: Set[str] = set()

    # everything inside (:predicates ... )

    pred_section = re.findall(r"\(:predicates(.*?)\)\s*\)", domain, re.S)

    if pred_section:

        # flatten all lines

        for raw_line in re.findall(r"\([^)]+\)", pred_section[0]):

            token = raw_line.strip()[1:].split()[0]

            declared.add(token)

    return declared

def _extract_referenced_predicates(domain: str) -> Set[str]:

    referenced: Set[str] = set()

    # every token starting with '(' that is not ':', 'and', 'not', 'forall', 'when'

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

        token = raw.strip()[1:].split()[0]

        if token.startswith(":"):

            continue

        if token in {"and", "not", "forall", "when", "exists"}:

            continue

        referenced.add(token)

    return referenced

def find_missing_predicates(base_domain: str, exploration_domain: str) -> List[str]:

    declared = _extract_declared_predicates(base_domain) | _extract_declared_predicates(exploration_domain)

    referenced = _extract_referenced_predicates(exploration_domain)

    return sorted([p for p in referenced if p not in declared])

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

# 2) Main execution routine

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

def run_skeleton_task():

    """Generic task runner with simple exploration and predicate-checking."""

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

    # === 2-A. Predicate Exploration Phase ==================================

    missing_preds = find_missing_predicates(COMBINED_DOMAIN_PDDL, EXPLORATION_DOMAIN_PDDL)

    if missing_preds:

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

    else:

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

    # =======================================================================

    # === 2-B. RLBench Environment Set-up ===================================

    env, task = setup_environment()

    try:

        descriptions, obs = task.reset()

        # (optional) initialise video writer

        init_video_writers(obs)

        # wrap task.step / get_observation for recording

        task.step = recording_step(task.step)

        task.get_observation = recording_get_observation(task.get_observation)

        # === 2-C. Obtain object positions from helper module ===============

        try:

            positions = get_object_positions()  # dict: name -> (x,y,z)

            print(f"[Info] object_positions keys: {list(positions.keys())}")

        except Exception as e:

            positions = {}

            print("[Warning] get_object_positions() failed:", e)

        # === 2-D. VERY Simple Demonstration Plan (if any object exists) ====

        if positions:

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

            first_obj_pos = positions[first_obj_name]

            print(f"[Plan] Attempting to pick first object '{first_obj_name}' at {first_obj_pos}")

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

            # 1) pick

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

            try:

                obs, reward, done = pick(

                    env,

                    task,

                    target_pos=np.asarray(first_obj_pos, dtype=float),

                    approach_distance=0.15,

                    max_steps=120,

                    threshold=0.01,

                    approach_axis='z',

                    timeout=10.0

                )

                if done:

                    print("[Plan] Task finished after pick.")

                    return

            except Exception as e:

                print("[Error] pick skill failed:", e)

                traceback.print_exc()

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

            # 2) rotate gripper to neutral orientation (identity quaternion)

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

            try:

                identity_quat = np.array([0.0, 0.0, 0.0, 1.0])  # xyzw

                obs, reward, done = rotate(

                    env,

                    task,

                    target_quat=identity_quat,

                    max_steps=80,

                    threshold=0.05,

                    timeout=6.0

                )

                if done:

                    print("[Plan] Task finished during rotate.")

                    return

            except Exception as e:

                print("[Warning] rotate skill failed:", e)

                traceback.print_exc()

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

            # 3) place the object back (simple heuristic: drop slightly above original)

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

            try:

                place_pos = np.asarray(first_obj_pos, dtype=float) + np.array([0.0, 0.0, 0.05])

                obs, reward, done = place(

                    env,

                    task,

                    target_pos=place_pos,

                    approach_distance=0.15,

                    max_steps=120,

                    threshold=0.01,

                    approach_axis='z',

                    timeout=10.0

                )

                if done:

                    print("[Plan] Task finished after place.")

                    return

            except Exception as e:

                print("[Error] place skill failed:", e)

                traceback.print_exc()

        else:

            print("[Plan] No objects discovered – idle run.")

        # === 2-E. Pulling drawers if available (optional) ===================

        try:

            drawer_handle = Shape("drawer_handle")

            sensor = ProximitySensor("drawer_open_sensor")

            handle_pos = drawer_handle.get_position()

            print(f"[Plan] Drawer handle detected at {handle_pos}, attempting pull.")

            # We assume move() skill exists and uses names rather than coords

            # If not, this will be caught.

            obs, reward, done = move(

                env,

                task,

                target_pos=np.asarray(handle_pos),

                max_steps=120,

                threshold=0.01,

                timeout=10.0

            )

            if done:

                print("[Plan] Task finished during move to drawer.")

                return

            # pull (no positional parameters)

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

            if done:

                print("[Plan] Task finished after pull.")

                return

        except Exception as e:

            print("[Info] No drawer handle or pull failed:", e)

    finally:

        shutdown_environment(env)

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

if __name__ == "__main__":

    run_skeleton_task()
