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

[Frozen Code Start]

[Frozen Code End]

def run_skeleton_task():
    print("===== Starting Skeleton Task =====")
    env, task = setup_environment()
    try:
        descriptions, obs = task.reset()
        init_video_writers(obs)
        # wrap step and get_observation 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 positions of objects
        positions = get_object_positions()

        try:
            # Choose an unlocked drawer: bottom is known unlocked
            drawer_name = 'bottom'
            # compute key for its side position
            side_pos_key = f'{drawer_name}_side_pos'
            if side_pos_key not in positions:
                raise KeyError(f"Side position for '{drawer_name}' not found in positions")
            # Step 1: rotate to current orientation (no-op but satisfies spec)
            obs = task.get_observation()
            current_quat = obs.gripper_pose[3:7]
            print(f"[Task] Step 1: Rotating to current orientation: {current_quat}")
            obs, reward, done = rotate(env, task, current_quat)
            if done:
                print("[Task] Task ended during rotation.")
                return

            # Step 2: move to the side position of the drawer
            target_side = np.array(positions[side_pos_key])
            print(f"[Task] Step 2: Moving to side position of '{drawer_name}' drawer at {target_side}")
            obs, reward, done = move(env, task, target_side)
            if done:
                print("[Task] Task ended during move.")
                return

            # Step 3: pull the drawer open
            pull_distance = 0.1
            print(f"[Task] Step 3: Pulling the drawer by {pull_distance} meters")
            obs, reward, done = pull(env, task, pull_distance=pull_distance, pull_axis='x')
            if done:
                print("[Task] Task ended during pull.")
                return

            # Step 4: pick up the rubbish (fallback to 'item3' if no 'rubbish')
            rubbish_key = 'rubbish' if 'rubbish' in positions else 'item3'
            if rubbish_key not in positions:
                raise KeyError(f"Neither 'rubbish' nor 'item3' found in positions")
            target_rubbish = np.array(positions[rubbish_key])
            print(f"[Task] Step 4: Picking up '{rubbish_key}' at {target_rubbish}")
            obs, reward, done = pick(env, task, target_rubbish)
            if done:
                print("[Task] Task ended during pick.")
                return

            # Step 5: place the rubbish into the bin
            bin_key = 'bin'
            if bin_key not in positions:
                raise KeyError(f"Bin position not found in positions")
            target_bin = np.array(positions[bin_key])
            print(f"[Task] Step 5: Placing '{rubbish_key}' into bin at {target_bin}")
            obs, reward, done = place(env, task, target_bin)
            if done:
                print("[Task] Completed successfully! Reward:", reward)
            else:
                print("[Task] Task not completed yet (done=False).")

        except Exception as e:
            print("[Task] Exception occurred during execution:", e)

    finally:
        shutdown_environment(env)

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

if __name__ == "__main__":
    run_skeleton_task()