[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():
    '''Generic skeleton for running the drawer-open and rubbish-disposal task.'''
    print("===== Starting Skeleton Task =====")

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

        # (Optional) Initialize video writers for capturing your simulation
        init_video_writers(obs)

        # Wrap the task steps for recording (if needed)
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # === Retrieve Object Positions ===
        positions = get_object_positions()
        # Select the bottom drawer for simplicity
        anchor_key = 'bottom_anchor_pos'
        side_key = 'bottom_side_pos'
        # Get the relevant positions and orientations
        drawer_anchor_pos = np.array(positions[anchor_key])
        drawer_side_pos = np.array(positions[side_key])
        # Try to retrieve a quaternion for ninety degrees, else default
        if 'ninety_deg' in positions:
            ninety_quat = np.array(positions['ninety_deg'])
        else:
            # approximate 90° rotation around z axis
            ninety_quat = np.array([0.0, 0.0, 0.7071, 0.7071])

        # === Step 1: Rotate gripper to ninety-degree orientation ===
        print("[Task] Step 1: Rotate gripper to ninety-degree orientation")
        obs, reward, done = rotate(env, task, target_quat=ninety_quat)
        if done:
            print("[Task] Terminating early after rotate. Reward:", reward)
            return

        # === Step 2: Pull the drawer to open it ===
        pull_distance = float(np.linalg.norm(drawer_side_pos - drawer_anchor_pos))
        print(f"[Task] Step 2: Pull drawer by distance {pull_distance:.3f}")
        obs, reward, done = pull(env, task, pull_distance=pull_distance, pull_axis='x')
        if done:
            print("[Task] Terminating early after pull. Reward:", reward)
            return

        # === Step 3: Pick up the rubbish from the table ===
        # We assume 'item3' corresponds to the rubbish to pick
        if 'item3' not in positions:
            raise KeyError("Could not find 'item3' in positions for rubbish")
        rubbish_pos = np.array(positions['item3'])
        print("[Task] Step 3: Pick up rubbish at position", rubbish_pos)
        obs, reward, done = pick(env, task, target_pos=rubbish_pos)
        if done:
            print("[Task] Terminating early after pick. Reward:", reward)
            return

        # === Step 4: Place the rubbish into the bin ===
        if 'bin' not in positions:
            raise KeyError("Could not find 'bin' in positions for disposal")
        bin_pos = np.array(positions['bin'])
        print("[Task] Step 4: Place rubbish into bin at position", bin_pos)
        obs, reward, done = place(env, task, target_pos=bin_pos)
        if done:
            print("[Task] Task completed successfully! Reward:", reward)
        else:
            print("[Task] Task ended without success. Reward:", reward)

    finally:
        # Always ensure the environment is properly shutdown
        shutdown_environment(env)

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

if __name__ == "__main__":
    run_skeleton_task()