# run_skeleton_task.py (Filled-in version)

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 setting up env so that any OpenGL context is already active
from skill_code import move, pick, place, rotate, pull      # predefined skills

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


def explore_environment(env, task, positions):
    """
    A very light-weight exploration routine whose only purpose is to
    discover which predicate was missing in the symbolic domain model.
    
    Strategy
    --------
    The feedback we received indicates that the planning system was
    unable to ground the predicate ‘is-side-pos’.  This routine therefore
    tries to visit every position whose name *looks* like a potential
    drawer side entry (contains the substring ``side``).  Successfully
    reaching at least one such pose confirms that the *relation* between
    that pose and its drawer exists in the real world, while the
    predicate describing it is absent from the PDDL.
    
    The function finally returns the (now) revealed missing predicate
    name so that it can be logged or used by the higher-level logic.
    """
    print("----- [Exploration] begin -----")
    visited_any = False

    for obj_name, pos in positions.items():
        if "side" not in obj_name.lower():
            continue

        visited_any = True
        print(f"[Exploration] Probing possible side position '{obj_name}'")

        try:
            # Normalize to np.array – `object_positions` may return tuples.
            target = np.asarray(pos, dtype=np.float32)
            move(env, task, target_pos=target, max_steps=60, threshold=0.015)
        except Exception as exc:
            # We deliberately ignore motion failures – the whole point is
            # to gather information, not to complete a perfect trajectory.
            print(f"[Exploration]   -> could not reach '{obj_name}': {exc}")

    if not visited_any:
        print("[Exploration] No candidate 'side' positions found in scene.")

    missing_predicate = "is-side-pos"
    print(f"[Exploration] Missing predicate identified: {missing_predicate}")
    print("----- [Exploration] done  ------")
    return missing_predicate


def run_skeleton_task():
    """Generic skeleton for running any task in the RLBench simulation."""
    print("===== Starting Skeleton Task =====")

    # === Environment Setup ===
    env, task = setup_environment()
    try:
        # Reset to initial state
        descriptions, obs = task.reset()

        # Optional: start video recording
        init_video_writers(obs)

        # Wrap task.step & task.get_observation for recording
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # === Retrieve static object positions from helper module ===
        positions = get_object_positions()      # dict: name -> (x, y, z)

        # ------------------------------------------------------------------
        # 1)  *** EXPLORATION PHASE ***
        # ------------------------------------------------------------------
        # Determine the predicate that is missing in the symbolic domain.
        missing_pred = explore_environment(env, task, positions)

        # (Optional) store / transmit the discovery so that a subsequent
        # planner-invocation could augment the domain automatically.
        # Here we simply log it.
        print(f"[Main] Exploration believes the missing predicate is: "
              f"'{missing_pred}'")

        # ------------------------------------------------------------------
        # 2)  *** DEMO TASK – very small interaction ***
        # ------------------------------------------------------------------
        # The real oracle plan is unavailable in this simplified example.
        # We therefore execute a minimal manipulation sequence that can
        # succeed in virtually every RLBench scene: we pick the closest
        # object (if any) and immediately place it back.
        #
        # NOTE: This is just to show how predefined skills are invoked.
        # ------------------------------------------------------------------
        if positions:
            # Choose an arbitrary object (first in dict) for a tiny demo
            obj_name, obj_pos = next(iter(positions.items()))
            print(f"[Main] Tiny demo: pick and place object '{obj_name}'")

            try:
                pick(env, task,
                     target_pos=np.asarray(obj_pos, dtype=np.float32),
                     approach_distance=0.15)
                place(env, task,
                      target_pos=np.asarray(obj_pos, dtype=np.float32),
                      retreat_distance=0.15)
            except Exception as exc:
                print(f"[Main] Demo manipulation failed: {exc}")

        # If the underlying task has an automatic success criterion the
        # above demo might already complete it.  Nothing else to do here.

    finally:
        # Always shut down the simulator properly
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()