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 *

from video import init_video_writers, recording_step, recording_get_observation

from object_positions import get_object_positions

def run_skeleton_task():
    """Execute the oracle plan: throw away the trash, leaving any other objects alone."""
    print("===== Starting Skeleton Task =====")

    # === Environment Setup ===
    env, task = setup_environment()
    try:
        # Reset environment and initialize recording
        descriptions, obs = task.reset()
        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 Object Positions ===
        positions = get_object_positions()
        try:
            # These keys must exist in the environment
            item1_pos   = positions['item1']
            item2_pos   = positions['item2']
            rubbish_pos = positions['rubbish']
            bin_pos     = positions['bin']
        except KeyError as e:
            print(f"[Error] Missing position for object: {e}")
            return

        # === Oracle Plan Execution ===
        # Frozen Code Region: do not modify
        obs, reward, done = pick(env, task, target_pos=item1_pos)
        obs, reward, done = place(env, task, target_pos=bin_pos)
        obs, reward, done = pick(env, task, target_pos=item2_pos)
        # End of Frozen Code Region

        # Place item2 back to its original location
        print(f"[Task] Placing item2 back at original position: {item2_pos}")
        obs, reward, done = place(env, task, target_pos=item2_pos)
        if done:
            print("[Task] Task ended prematurely during place(item2 back).")
            return

        # Place item1 back to its original location
        print(f"[Task] Placing item1 back at original position: {item1_pos}")
        obs, reward, done = place(env, task, target_pos=item1_pos)
        if done:
            print("[Task] Task ended prematurely during place(item1 back).")
            return

        # Now perform the true goal: throw away the rubbish
        print(f"[Task] Picking up rubbish at: {rubbish_pos}")
        obs, reward, done = pick(env, task, target_pos=rubbish_pos)
        if done:
            print("[Task] Task ended prematurely during pick(rubbish).")
            return

        print(f"[Task] Placing rubbish into bin at: {bin_pos}")
        obs, reward, done = place(env, task, target_pos=bin_pos)
        if done:
            print(f"[Task] Task ended during place(rubbish). Reward: {reward}")
            return

        # Final status
        if not done:
            print(f"[Task] Trash disposed. Final reward: {reward}")
        else:
            print("[Task] Task signaled done=True before completing all steps.")

    finally:
        shutdown_environment(env)

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

if __name__ == "__main__":
    run_skeleton_task()