[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():
    '''Generic skeleton for running any task in your simulation.'''
    print("===== Starting Skeleton Task =====")
    
    # === Environment Setup ===
    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
        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()
        # Map names from our list to the keys returned by get_object_positions
        bottom_side_pos   = positions['bottom_side_pos']
        bottom_anchor_pos = positions['bottom_anchor_pos']
        item1_pos         = positions['item1']
        item2_pos         = positions['item2']
        plate_pos         = positions['plate']

        # === Execute Oracle Plan ===

        # 1) Rotate gripper from 0° to 90° around its local z-axis
        # Quaternion for 90° rotation about z: [x, y, z, w] = [0, 0, sin(45°), cos(45°)]
        target_quat = np.array([0.0, 0.0, np.sin(np.pi/4), np.cos(np.pi/4)])
        try:
            obs, reward, done = rotate(env, task, target_quat)
            if done:
                print("[Task] Ended during rotate step.")
                return
        except Exception as e:
            print(f"[Error] rotate failed: {e}")
            return

        # 2) Move to the side position of the bottom drawer
        try:
            obs, reward, done = move(env, task, bottom_side_pos)
            if done:
                print("[Task] Ended during move-to-side step.")
                return
        except Exception as e:
            print(f"[Error] move-to-side failed: {e}")
            return

        # 3) Move to the anchor (grasp) position on the bottom drawer
        try:
            obs, reward, done = move(env, task, bottom_anchor_pos)
            if done:
                print("[Task] Ended during move-to-anchor step.")
                return
        except Exception as e:
            print(f"[Error] move-to-anchor failed: {e}")
            return

        # 4) Pull the drawer open along the x-axis by 0.1m
        try:
            obs, reward, done = pull(env, task, pull_distance=0.1, pull_axis='x')
            if done:
                print("[Task] Ended during pull step.")
                return
        except Exception as e:
            print(f"[Error] pull failed: {e}")
            return

        # 5) Pick up the first tomato (item1) from the table
        try:
            obs, reward, done = pick(env, task, item1_pos)
            if done:
                print("[Task] Ended during first pick step.")
                return
        except Exception as e:
            print(f"[Error] pick(item1) failed: {e}")
            return

        # 6) Place the first tomato onto the plate
        try:
            obs, reward, done = place(env, task, plate_pos)
            if done:
                print("[Task] Ended during first place step.")
                return
        except Exception as e:
            print(f"[Error] place(item1->plate) failed: {e}")
            return

        # 7) Pick up the second tomato (item2) from the table
        try:
            obs, reward, done = pick(env, task, item2_pos)
            if done:
                print("[Task] Ended during second pick step.")
                return
        except Exception as e:
            print(f"[Error] pick(item2) failed: {e}")
            return

        # 8) Place the second tomato onto the plate
        try:
            obs, reward, done = place(env, task, plate_pos)
            if done:
                print("[Task] Ended during second place step.")
                return
        except Exception as e:
            print(f"[Error] place(item2->plate) failed: {e}")
            return

        # Task complete
        print("[Task] Goal achieved: bottom drawer opened and two tomatoes placed on plate.")

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

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


if __name__ == "__main__":
    run_skeleton_task()
