from env import setup_environment, shutdown_environment
import numpy as np
from pyrep.objects.shape import Shape
from pyrep.objects.proximity_sensor import ProximitySensor

from skill_code import move, pick, place, rotate, pull, normalize_quaternion
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions

def run_skeleton_task():
    print("===== Starting Task: Slide open drawer and discard rubbish =====")

    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)
        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 from the environment
        positions = get_object_positions()

        # === Step 1: Move to side position of the bottom drawer ===
        print("[Task] 1) Moving to drawer side position")
        side_pos = np.array(positions['bottom_side_pos'])

        # === Step 2: Rotate gripper (no-op or maintain current orientation) ===
        print("[Task] 2) Rotating gripper")
        obs_curr = task.get_observation()
        current_quat = normalize_quaternion(obs_curr.gripper_pose[3:7])

        # === Step 3: Move to anchor position on the bottom drawer ===
        print("[Task] 3) Moving to drawer anchor position")
        anchor_pos = np.array(positions['bottom_anchor_pos'])

        # === Steps 1-4 (Frozen Code) ===
        obs, reward, done = move(env, task, target_pos=side_pos)
        obs, reward, done = rotate(env, task, target_quat=current_quat)
        obs, reward, done = move(env, task, target_pos=anchor_pos)
        obs, reward, done = pull(env, task, pull_distance=0.05, pull_axis='x')

        if done:
            print("[Task] Early termination after opening drawer.")
            return

        # === Step 5: Move to the rubbish on the table ===
        print("[Task] 5) Moving to rubbish on the table")
        rubbish_pos = np.array(positions['rubbish'])
        obs, reward, done = move(env, task, target_pos=rubbish_pos)
        if done:
            print("[Task] Early termination after moving to rubbish.")
            return

        # === Step 6: Pick up the rubbish ===
        print("[Task] 6) Picking up the rubbish")
        obs, reward, done = pick(env, task, target_pos=rubbish_pos)
        if done:
            print("[Task] Early termination during pick.")
            return

        # === Step 7: Move to bin ===
        print("[Task] 7) Moving to bin")
        bin_pos = np.array(positions['bin'])
        obs, reward, done = move(env, task, target_pos=bin_pos)
        if done:
            print("[Task] Early termination after moving to bin.")
            return

        # === Step 8: Place the rubbish into the bin ===
        print("[Task] 8) Placing rubbish into bin")
        obs, reward, done = place(env, task, target_pos=bin_pos)
        if done:
            print("[Task] Early termination during place.")
            return

        print("[Task] Completed all steps successfully! Final reward:", reward)

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

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

if __name__ == "__main__":
    run_skeleton_task()