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():
    print("===== Starting Skeleton Task =====")
    env, task = setup_environment()
    try:
        # Reset the task
        descriptions, obs = task.reset()

        # Initialize video recording
        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)

        # Retrieve object positions
        positions = get_object_positions()

        # Select the bottom drawer handle positions
        drawer_side_key = 'bottom_side_pos'
        drawer_anchor_key = 'bottom_anchor_pos'
        if drawer_side_key not in positions or drawer_anchor_key not in positions:
            print(f"[Task] Error: Drawer positions '{drawer_side_key}' or '{drawer_anchor_key}' not found.")
            return

        drawer_side_pos = np.array(positions[drawer_side_key])
        drawer_anchor_pos = np.array(positions[drawer_anchor_key])

        # Step 1: rotate gripper to current orientation (no-op rotate)
        print("[Task] Step 1: rotate gripper to target orientation (no-op)")
        target_quat = obs.gripper_pose[3:7]
        obs, reward, done = rotate(env, task, target_quat=target_quat)
obs, reward, done = move(env, task, target_pos=drawer_side_pos)

        if done:
            print("[Task] Ended after move to side.")
            return

        # Step 3: move to the anchor position of the drawer
        print(f"[Task] Step 3: move to drawer anchor position: {drawer_anchor_pos}")
        obs, reward, done = move(env, task, target_pos=drawer_anchor_pos)
        if done:
            print("[Task] Ended after move to anchor.")
            return

        # Step 4: pull the drawer to open it
        pull_distance = 0.1
        pull_axis = 'x'
        print(f"[Task] Step 4: pull drawer with distance {pull_distance} along '{pull_axis}' axis")
        obs, reward, done = pull(env, task, pull_distance=pull_distance, pull_axis=pull_axis)
        if done:
            print("[Task] Ended after pull.")
            return

        # Step 5: pick up the rubbish (use item3 as the rubbish)
        rubbish_key = 'item3'
        if rubbish_key not in positions:
            print(f"[Task] Error: Rubbish object '{rubbish_key}' not found.")
            return
        rubbish_pos = np.array(positions[rubbish_key])
        print(f"[Task] Step 5: pick up rubbish at: {rubbish_pos}")
        obs, reward, done = pick(env, task, target_pos=rubbish_pos)
        if done:
            print("[Task] Ended after pick.")
            return

        # Step 6: place the rubbish into the bin
        bin_key = 'bin'
        if bin_key not in positions:
            print(f"[Task] Error: Bin position '{bin_key}' not found.")
            return
        bin_pos = np.array(positions[bin_key])
        print(f"[Task] Step 6: place rubbish in bin at: {bin_pos}")
        obs, reward, done = place(env, task, target_pos=bin_pos)
        if done:
            print(f"[Task] Task completed successfully! Reward: {reward}")
        else:
            print("[Task] Task not completed (done=False).")

    finally:
        shutdown_environment(env)

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

if __name__ == "__main__":
    run_skeleton_task()