[Frozen Code Start]
    
[Frozen Code End]

import numpy as np
from env import setup_environment, shutdown_environment
from skill_code import move, pick, place, rotate, pull, normalize_quaternion
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 =====")
    
    env, task = setup_environment()
    try:
        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)

        positions = get_object_positions()
        
        # Oracle plan to open a drawer, pick up rubbish, and dispose it
        
        # STEP 1: Rotate gripper (no‐op rotation to satisfy plan)
        obs = task.get_observation()
        current_quat = normalize_quaternion(obs.gripper_pose[3:7])
        print("[Task] STEP 1: Rotating gripper to current orientation")
        try:
            obs, reward, done = rotate(env, task, current_quat)
            if done:
                print("[Task] Ended after rotate. Reward:", reward)
                return
        except Exception as e:
            print(f"[Task] Error during rotate: {e}")
            raise

        # STEP 2: Move to drawer side position
        bottom_side = positions['bottom_side_pos']
        print("[Task] STEP 2: Moving to bottom side position:", bottom_side)
        obs, reward, done = move(env, task, bottom_side)
        if done:
            print("[Task] Ended after move to side. Reward:", reward)
            return

        # STEP 3: Move to drawer anchor position
        bottom_anchor = positions['bottom_anchor_pos']
        print("[Task] STEP 3: Moving to bottom anchor position:", bottom_anchor)
        obs, reward, done = move(env, task, bottom_anchor)
        if done:
            print("[Task] Ended after move to anchor. Reward:", reward)
            return

        # STEP 4: Pull drawer to open
        print("[Task] STEP 4: Pulling drawer to open")
        obs, reward, done = pull(env, task, pull_distance=0.1, pull_axis='y')
        if done:
            print("[Task] Ended after pull. Reward:", reward)
            return

        # STEP 5: Retreat to side position
        print("[Task] STEP 5: Moving back to bottom side position")
        obs, reward, done = move(env, task, bottom_side)
        if done:
            print("[Task] Ended after retreat. Reward:", reward)
            return

        # STEP 6: Pick up rubbish (assume item3 is rubbish)
        rubbish_pos = positions['item3']
        print("[Task] STEP 6: Picking up rubbish at:", rubbish_pos)
        obs, reward, done = pick(env, task, target_pos=rubbish_pos)
        if done:
            print("[Task] Ended after pick. Reward:", reward)
            return

        # STEP 7: Move to trash bin
        bin_pos = positions['bin']
        print("[Task] STEP 7: Moving to trash bin at:", bin_pos)
        obs, reward, done = move(env, task, bin_pos)
        if done:
            print("[Task] Ended after move to bin. Reward:", reward)
            return

        # STEP 8: Place rubbish in the bin
        print("[Task] STEP 8: Placing rubbish in bin")
        obs, reward, done = place(env, task, target_pos=bin_pos)
        if done:
            print("[Task] Completed successfully! Reward:", reward)
        else:
            print("[Task] Plan executed but done=False. Reward:", reward)

    finally:
        shutdown_environment(env)

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

if __name__ == "__main__":
    run_skeleton_task()