from pyrep.objects.shape import Shape
from pyrep.objects.proximity_sensor import ProximitySensor
import numpy as np
from scipy.spatial.transform import Rotation as R

from env import setup_environment, shutdown_environment
from skill_code import rotate, move, pull, pick, place
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:
        descriptions, obs = task.reset()
        init_video_writers(obs)
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # Retrieve object positions
        positions = get_object_positions()
        quat_zero = R.from_euler('z', 0.0, degrees=True).as_quat()
        quat_ninety = R.from_euler('z', 90.0, degrees=True).as_quat()

        # Map PDDL names to environment positions
        drawer_side_pos = positions.get('bottom_side_pos')
        if drawer_side_pos is None:
            raise KeyError("Position for 'bottom_side_pos' not found in object positions.")
        rubbish_pos = positions.get('item3')
        if rubbish_pos is None:
            raise KeyError("Position for 'item3' (rubbish) not found.")
        bin_pos = positions.get('bin')
        if bin_pos is None:
            raise KeyError("Position for 'bin' not found.")

        # Execute oracle plan up to moving to bin
        print("[Plan] Step 1: rotate to ninety_deg orientation")
        obs, reward, done = rotate(env, task, target_quat=quat_ninety)
        if done:
            print("Task ended early during rotate. Reward:", reward)
            return

        print("[Plan] Step 2: move to drawer side position")
        obs, reward, done = move(env, task, target_pos=np.array(drawer_side_pos))
        if done:
            print("Task ended early during move-to-side. Reward:", reward)
            return

        print("[Plan] Step 3: pull drawer open")
        obs, reward, done = pull(env, task, pull_distance=0.05, pull_axis='x')
        if done:
            print("Task ended early during pull. Reward:", reward)
            return

        print("[Plan] Step 4: move to rubbish at item3 position")
        obs, reward, done = move(env, task, target_pos=np.array(rubbish_pos))
        if done:
            print("Task ended early during move-to-rubbish. Reward:", reward)
            return

        print("[Plan] Step 5: pick up rubbish")
        obs, reward, done = pick(env, task, target_pos=np.array(rubbish_pos))
        if done:
            print("Task ended early during pick. Reward:", reward)
            return

        print("[Plan] Step 6: move to bin")
        obs, reward, done = move(env, task, target_pos=np.array(bin_pos))
        if done:
            print("Task ended early during move-to-bin. Reward:", reward)
            return

        [Frozen Code Start]
        obs, reward, done = rotate(env, task, target_quat=quat_ninety)
        obs, reward, done = move(env, task, target_pos=np.array(drawer_side_pos)
        obs, reward, done = pull(env, task, pull_distance=0.05, pull_axis='x')
        obs, reward, done = move(env, task, target_pos=np.array(rubbish_pos)
        obs, reward, done = pick(env, task, target_pos=np.array(rubbish_pos)
        obs, reward, done = move(env, task, target_pos=np.array(bin_pos)
        [Frozen Code End]

        # Step 7: place the rubbish in the bin
        print("[Plan] Step 7: place rubbish in bin")
        obs, reward, done = place(env, task, target_pos=np.array(bin_pos))
        if done:
            print("Task ended after place. Final reward:", reward)
            return

        print("===== Task completed. Final reward:", reward)
    finally:
        shutdown_environment(env)
    print("===== End of Skeleton Task =====")

if __name__ == "__main__":
    run_skeleton_task()