# run_skeleton_task.py (Completed Version – missing-predicate exploration + safe shutdown)

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

# Import every predefined skill exactly as delivered
from skill_code import *

from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions


# ---------------------------------------------------------------------
# Helper: very light-weight PDDL predicate extractor
# ---------------------------------------------------------------------
def _extract_predicates(pddl_txt: str):
    """Return a set of all predicate names contained in the PDDL string."""
    inside_pred_block = False
    preds = set()
    for line in pddl_txt.splitlines():
        line = line.strip()
        # Detect (:predicates … ) start
        if line.startswith("(:predicates"):
            inside_pred_block = True
            # Strip leading "(:predicates"
            line = line[len("(:predicates"):].strip()
        # Detect end of predicate block
        if inside_pred_block and ")" in line:
            # Could be many ) on a single line – treat cautiously
            parts = line.split(")")
            # every part before last might contain predicates
            for part in parts[:-1]:
                part = part.strip()
                if part:
                    m = re.match(r"\(?\s*([a-zA-Z0-9_\-]+)", part)
                    if m:
                        preds.add(m.group(1))
            # If this line contained the closing “)” of the predicate block,
            # we are done.
            if line.endswith(")"):
                inside_pred_block = False
            continue
        # Normal predicate lines
        if inside_pred_block:
            m = re.match(r"\(?\s*([a-zA-Z0-9_\-]+)", line)
            if m:
                preds.add(m.group(1))
    return preds


# ---------------------------------------------------------------------
# PDDL strings (truncated to predicate blocks – enough for diff)
# ---------------------------------------------------------------------
COMBINED_DOMAIN_PDDL = """
(: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 = """
(: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)
    (lock-known ?obj - object)
)
"""


def discover_missing_predicates():
    """Compare exploration domain vs combined domain & return missing symbols."""
    combined_preds = _extract_predicates(COMBINED_DOMAIN_PDDL)
    exploration_preds = _extract_predicates(EXPLORATION_DOMAIN_PDDL)
    return exploration_preds.difference(combined_preds)


# ---------------------------------------------------------------------
# Main execution entry
# ---------------------------------------------------------------------
def run_skeleton_task():
    print("===== Starting Skeleton Task =====")

    # ------------------------------------------------
    # Phase 0: predicate discovery (static analysis)
    # ------------------------------------------------
    missing_predicates = discover_missing_predicates()
    print(f"[Exploration] Predicates present in exploration-domain but "
          f"missing in combined-domain: {sorted(missing_predicates)}")
    # For this project we expect exactly {'lock-known'}
    if "lock-known" in missing_predicates:
        print("[Exploration] The missing predicate ‘lock-known’ has been "
              "successfully identified!")

    # ------------------------------------------------
    # Phase 1: normal RLBench environment handling
    # ------------------------------------------------
    env, task = setup_environment()
    try:
        descriptions, obs = task.reset()

        # Optional video logging
        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)

        # ------------------------------------------------
        # Retrieve (approx.) object locations from helper
        # ------------------------------------------------
        positions = get_object_positions()
        # The keys depend on the concrete RLBench task.
        # Here we try best-effort retrieval with graceful fallback.
        drawer_handle_key = next((k for k in positions if "handle" in k.lower()), None)
        tomato_key = next((k for k in positions if "tomato" in k.lower()), None)
        plate_key = next((k for k in positions if "plate" in k.lower()), None)

        # ------------------------------------------------
        # Very simple illustrative plan:
        #  (1) If a tomato exists, pick it up
        #  (2) If a plate exists, place the tomato on it
        # The policy is deliberately minimal – the main assignment
        # is demonstration of predicate discovery & usage of skills.
        # ------------------------------------------------
        if tomato_key and plate_key:
            print(f"[Task] Attempting to pick {tomato_key} and place on {plate_key}")
            tomato_pos = np.array(positions[tomato_key])
            plate_pos = np.array(positions[plate_key])

            obs, reward, done = pick(
                env,
                task,
                target_pos=tomato_pos,
                approach_distance=0.15,
                max_steps=120,
                threshold=0.01,
                approach_axis='z',
                timeout=10.0
            )
            if done:
                print("[Task] Episode finished unexpectedly during pick.")
            else:
                obs, reward, done = place(
                    env,
                    task,
                    target_pos=plate_pos,
                    approach_distance=0.15,
                    max_steps=120,
                    threshold=0.01,
                    approach_axis='z',
                    timeout=10.0
                )
                if done:
                    print("[Task] Episode finished after place.")
        else:
            print("[Task] Required objects not found in object_positions(). "
                  "Skipping pick-and-place demonstration.")

        # ------------------------------------------------
        # Optional drawer interaction (uses rotate / pull skills)
        # Only executed if we detect a drawer handle position.
        # ------------------------------------------------
        if drawer_handle_key:
            print("[Task] Found drawer handle – attempting open routine.")
            handle_pos = np.array(positions[drawer_handle_key])

            # 1) Move gripper to be perpendicular to drawer handle (rotate skill)
            obs, reward, done = rotate(
                env,
                task,
                target_euler=(0, 0, np.deg2rad(90)),  # rotate to 90 deg around Z
                max_steps=50,
                threshold=0.01,
                timeout=5.0
            )
            if done:
                print("[Task] Episode finished during rotate, aborting pull.")
            else:
                # 2) Move close to the handle (simple move skill)
                obs, reward, done = move(
                    env,
                    task,
                    target_pos=handle_pos,
                    max_steps=120,
                    threshold=0.01,
                    timeout=10.0
                )
                if done:
                    print("[Task] Episode finished during move, aborting pull.")
                else:
                    # 3) Attempt to pull the drawer open
                    obs, reward, done = pull(
                        env,
                        task,
                        pull_distance=0.15,
                        max_steps=100,
                        threshold=0.01,
                        timeout=5.0
                    )
                    if done:
                        print("[Task] Episode finished during pull.")

    finally:
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()
