[Frozen Code Start]

[Frozen Code End]

import numpy as np
from env import setup_environment, shutdown_environment
from skill_code import pick, place, move, rotate, pull
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions

def run_skeleton_task():
    print("===== Starting Skeleton Task =====")
    
    # === Environment Setup ===
    env, task = setup_environment()
    try:
        # Reset the task to its initial state
        descriptions, obs = task.reset()

        # Initialize video writers for recording
        init_video_writers(obs)

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

        # Retrieve object and pose positions
        positions = get_object_positions()
        # Quaternions for 0° and 90° rotations (fallback values if not provided)
        zero_quat = positions.get('zero_deg', np.array([0.0, 0.0, 0.0, 1.0]))
        ninety_quat = positions.get('ninety_deg', np.array([0.0, 0.7071, 0.0, 0.7071]))
        # Drawer poses
        side_bottom = positions.get('bottom_side_pos')
        anchor_bottom = positions.get('bottom_anchor_pos')
        # Tomato and plate poses
        tomato1_pos = positions.get('tomato1')
        tomato2_pos = positions.get('tomato2')
        plate_pos = positions.get('plate')

        # === Oracle Plan Execution ===

        # Step 1: rotate gripper from zero_deg to ninety_deg
        print("[Step 1] Rotating gripper to ninety degrees")
        obs, reward, done = rotate(env, task, target_quat=ninety_quat)
        if done:
            print("[Step 1] Task ended prematurely")
            return

        # Step 2: move to side position of bottom drawer
        print("[Step 2] Moving to side position of bottom drawer")
        obs, reward, done = move(env, task, target_pos=side_bottom)
        if done:
            print("[Step 2] Task ended prematurely")
            return

        # Step 3: move to anchor position of bottom drawer
        print("[Step 3] Moving to anchor position of bottom drawer")
        obs, reward, done = move(env, task, target_pos=anchor_bottom)
        if done:
            print("[Step 3] Task ended prematurely")
            return

        # Step 4: pick the drawer handle (using pick on the anchor pose)
        print("[Step 4] Picking drawer handle")
        obs, reward, done = pick(env, task, target_pos=anchor_bottom)
        if done:
            print("[Step 4] Task ended prematurely")
            return

        # Step 5: pull to open the drawer
        print("[Step 5] Pulling open the drawer")
        # pull_distance and axis are chosen to approximate drawer opening
        obs, reward, done = pull(env, task, pull_distance=0.05, pull_axis='x')
        if done:
            print("[Step 5] Task ended prematurely")
            return

        # Step 6: pick tomato1 from table
        print("[Step 6] Picking tomato1 from table")
        obs, reward, done = pick(env, task, target_pos=tomato1_pos)
        if done:
            print("[Step 6] Task ended prematurely")
            return

        # Step 7: place tomato1 onto plate
        print("[Step 7] Placing tomato1 onto plate")
        obs, reward, done = place(env, task, target_pos=plate_pos)
        if done:
            print("[Step 7] Task ended prematurely")
            return

        # Step 8: pick tomato2 from table
        print("[Step 8] Picking tomato2 from table")
        obs, reward, done = pick(env, task, target_pos=tomato2_pos)
        if done:
            print("[Step 8] Task ended prematurely")
            return

        # Step 9: place tomato2 onto plate
        print("[Step 9] Placing tomato2 onto plate")
        obs, reward, done = place(env, task, target_pos=plate_pos)
        if done:
            print("[Step 9] Task ended prematurely")
            return

        print("[Task] Oracle plan executed successfully. Reward:", reward)

    finally:
        # Always ensure the environment is properly shutdown
        shutdown_environment(env)
        print("===== End of Skeleton Task =====")

if __name__ == "__main__":
    run_skeleton_task()