[Frozen Code Start]

[Frozen Code End]

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():
    '''Concrete implementation for opening a drawer and tossing two tomatoes onto a plate.'''
    print("===== Starting Skeleton Task =====")
    
    # === Environment Setup ===
    env, task = setup_environment()
    try:
        # Reset the task and get initial observation
        descriptions, obs = task.reset()

        # Initialize video recording if desired
        init_video_writers(obs)

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

        # === Retrieve object positions ===
        positions = get_object_positions()
        # Expect keys: 'bottom_side_pos', 'bottom_anchor_pos', 'item1', 'item2', 'plate'
        if 'bottom_side_pos' not in positions or 'item1' not in positions or 'item2' not in positions or 'plate' not in positions:
            print("[Error] Missing required object positions in environment.")
            return

        # === Plan Execution ===

        # Step 1: Rotate gripper from zero to ninety degrees around Z axis
        print("[Step 1] rotate to 90 degrees around Z axis")
        # Construct quaternion for 90° about Z: [x,y,z,w] = [0,0,sin(45°),cos(45°)]
        qz = np.sin(np.deg2rad(90.0/2.0))
        qw = np.cos(np.deg2rad(90.0/2.0))
        target_quat = np.array([0.0, 0.0, qz, qw])
        obs, reward, done = rotate(env, task, target_quat)
        if done:
            print("[Task] Ended after rotate.")
            return

        # Step 2: Move to the side position of the bottom drawer
        print("[Step 2] move to bottom drawer side position")
        side_pos = np.array(positions['bottom_side_pos'])
        obs, reward, done = move(env, task, side_pos)
        if done:
            print("[Task] Ended after move.")
            return

        # Step 3: Pull the drawer open by 0.1 meters along its local x-axis
        print("[Step 3] pull drawer open")
        obs, reward, done = pull(env, task, pull_distance=0.10, pull_axis='x')
        if done:
            print("[Task] Ended after pull.")
            return

        # Step 4: Pick the first tomato (item1) from the table
        print("[Step 4] pick item1 (first tomato)")
        target1 = np.array(positions['item1'])
        obs, reward, done = pick(env, task, target1, approach_distance=0.15, max_steps=100, threshold=0.01, approach_axis='z', timeout=10.0)
        if done:
            print("[Task] Ended after pick(item1).")
            return

        # Step 5: Place the first tomato onto the plate
        print("[Step 5] place item1 onto plate")
        plate_pos = np.array(positions['plate'])
        obs, reward, done = place(env, task, plate_pos, approach_distance=0.15, max_steps=100, threshold=0.01, approach_axis='z', timeout=10.0)
        if done:
            print("[Task] Ended after place(item1).")
            return

        # Step 6: Pick the second tomato (item2) from the table
        print("[Step 6] pick item2 (second tomato)")
        target2 = np.array(positions['item2'])
        obs, reward, done = pick(env, task, target2, approach_distance=0.15, max_steps=100, threshold=0.01, approach_axis='z', timeout=10.0)
        if done:
            print("[Task] Ended after pick(item2).")
            return

        # Step 7: Place the second tomato onto the plate
        print("[Step 7] place item2 onto plate")
        obs, reward, done = place(env, task, plate_pos, approach_distance=0.15, max_steps=100, threshold=0.01, approach_axis='z', timeout=10.0)
        if done:
            print("[Task] Ended after place(item2).")
            return

        # Final check
        if reward is not None:
            print(f"[Task] Completed plan with reward: {reward}")
        else:
            print("[Task] Completed plan, no final reward returned.")

    except Exception as e:
        print(f"[Error] Exception during task execution: {e}")
    finally:
        # Ensure we always shut down the environment
        shutdown_environment(env)

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

if __name__ == "__main__":
    run_skeleton_task()