# run_skeleton_task.py (Completed Version)

import time
import math
import numpy as np

# RL-Bench / PyRep specific
from pyrep.objects.shape import Shape
from pyrep.objects.proximity_sensor import ProximitySensor

# Environment helpers (provided externally)
from env import setup_environment, shutdown_environment

# Low-level skills that already exist – DO NOT redefine them
from skill_code import rotate, pick, place, move, pull

# Video helpers (optional but kept for completeness)
from video import (
    init_video_writers,
    recording_step,
    recording_get_observation,
)

# Utility that returns the initial locations of interest
from object_positions import get_object_positions


def explore_rotated_predicate(env, task):
    """
    Minimal exploration routine for discovering / satisfying the ‘rotated’
    predicate that was reported missing in feedback.  We simply try to rotate
    the gripper 90° about the Z-axis (quaternion = [0, 0, sin(π/4), cos(π/4)])
    because the combined‐domain PDDL uses the constant ‘ninety_deg’.
    """
    print("===== [Exploration] Looking for missing predicate: rotated =====")

    # 90° rotation about the z-axis (xyzw order)
    ninety_deg_quat = np.array([0.0, 0.0, 0.7071068, 0.7071068])
    try:
        obs, reward, done = rotate(
            env,
            task,
            target_quat=ninety_deg_quat,
            max_steps=120,
            threshold=0.03,
            timeout=10.0,
        )
        print(
            "[Exploration] Finished rotate call – "
            f"done flag = {done}, last reward = {reward:.4f}"
        )
        return done
    except Exception as e:
        print(f"[Exploration] Rotate skill failed: {e}")
        return False


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

    # ------------------------------------------------------------------
    # 1) Environment initialisation
    # ------------------------------------------------------------------
    env, task = setup_environment()
    try:
        descriptions, obs = task.reset()
        init_video_writers(obs)  # optional screen-capture

        # Wrap step / observation so that the video module can record
        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)

        # ------------------------------------------------------------------
        # 2) Fetch high-level information about objects / positions
        # ------------------------------------------------------------------
        positions = get_object_positions()
        print("[Info] Retrieved object positions:", positions)

        # We keep the identifiers general in order to remain domain-agnostic.
        # Replace keys below with real names if available in your simulator.
        drawer_handle_pos = positions.get("drawer_handle")
        disposal_area_pos = positions.get("disposal_bin")
        tomato_pos = positions.get("tomato")

        # ------------------------------------------------------------------
        # 3) Exploration phase – try to satisfy UNKNOWN predicate “rotated”
        # ------------------------------------------------------------------
        rotated_ok = explore_rotated_predicate(env, task)
        if rotated_ok:
            print("[Main] ‘rotated’ predicate appears to be satisfied.")
        else:
            print("[Main] Could not fully confirm the ‘rotated’ predicate, "
                  "continuing anyway…")

        # ------------------------------------------------------------------
        # 4) Example oracle-like plan execution
        # ------------------------------------------------------------------
        #
        # NOTE: In a real solution you would have a full symbolic plan
        # (move-to-side → move-to-anchor → pick-drawer → pull … etc.).
        # Here we show a concise sequence purely to demonstrate how to
        # call the predefined skills without redefining any primitives.
        #

        # 4-a) If we know where the drawer handle is, approach it
        if drawer_handle_pos is not None:
            print("[Plan] Moving gripper near drawer handle:", drawer_handle_pos)
            obs, reward, done = move(
                env,
                task,
                target_pos=drawer_handle_pos,
                approach_distance=0.10,
                max_steps=80,
                threshold=0.01,
                approach_axis="xy",  # do horizontal approach
                timeout=8.0,
            )
            if done:
                print("[Plan] Task ended during move; exiting.")
                return

        # 4-b) Try to pull the drawer after a short pause
        time.sleep(0.25)
        try:
            obs, reward, done = pull(env, task)
            if done:
                print("[Plan] Drawer successfully pulled; task finished.")
                return
        except Exception as e:
            print(f"[Plan] Pull skill failed or not applicable: {e}")

        # 4-c) Optional: pick up an object (e.g., tomato) and drop it
        if tomato_pos is not None and disposal_area_pos is not None:
            print("[Plan] Picking tomato at", tomato_pos)
            obs, reward, done = pick(
                env,
                task,
                target_pos=tomato_pos,
                approach_distance=0.12,
                max_steps=100,
                threshold=0.01,
                approach_axis="z",
                timeout=10.0,
            )
            if done:
                print("[Plan] Task ended on pick; exiting.")
                return

            time.sleep(0.25)

            print("[Plan] Placing tomato in disposal area at", disposal_area_pos)
            obs, reward, done = place(
                env,
                task,
                target_pos=disposal_area_pos,
                approach_distance=0.15,
                max_steps=100,
                threshold=0.01,
                approach_axis="z",
                timeout=10.0,
            )
            if done:
                print("[Plan] Task succeeded after placing object.")
                return

        # ------------------------------------------------------------------
        # 5) Fallback completion message
        # ------------------------------------------------------------------
        print("[Main] Reached end of scripted actions. "
              "If the goal is unmet, further planning is required.")

    finally:
        # ------------------------------------------------------------------
        # 6) Always shut down environment cleanly
        # ------------------------------------------------------------------
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()