import numpy as np
from scipy.spatial.transform import Rotation as R
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():
    '''Task: Open one drawer and put the rubbish in the bin'''
    print("===== Starting Task: Open drawer and dispose rubbish =====")
    env, task = setup_environment()
    try:
        # Reset task and initialize recording
        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 from the scene
        positions = get_object_positions()
        # Choose the bottom drawer for simplicity
        bottom_side_pos   = positions['bottom_side_pos']
        bottom_anchor_pos = positions['bottom_anchor_pos']
        # Rubbish and bin positions
        rubbish_pos = positions.get('rubbish', positions.get('item3'))
        bin_pos     = positions['bin']

        # Step 1: Rotate gripper to 90 degrees about X to face drawer handle
        print("[Plan] Step 1: rotate gripper to 90° about X")
        target_quat = R.from_euler('x', 90, degrees=True).as_quat()  # [x,y,z,w]
        obs, reward, done = rotate(env, task, target_quat)
        if done:
            print("[Task] Ended unexpectedly after rotate.")
            return

        # Step 2: Move gripper to the side position of the bottom drawer
        print("[Plan] Step 2: move to drawer side position")
        obs, reward, done = move(env, task, bottom_side_pos)
        if done:
            print("[Task] Ended unexpectedly after move to side.")
            return

        # Step 3: Move gripper to the anchor (handle) position of the bottom drawer
        print("[Plan] Step 3: move to drawer anchor position")
        obs, reward, done = move(env, task, bottom_anchor_pos)
        if done:
            print("[Task] Ended unexpectedly after move to anchor.")
            return

        # Step 4: Pull the drawer open
        print("[Plan] Step 4: pull drawer open")
        # pull_distance: how far to open; pull_axis: axis along which to pull
        obs, reward, done = pull(env, task, pull_distance=0.2, pull_axis='x')
        if done:
            print("[Task] Ended unexpectedly after pull.")
            return

        # Step 5: Pick up the rubbish
        print("[Plan] Step 5: pick up rubbish")
        if rubbish_pos is None:
            raise RuntimeError("Rubbish position not found.")
        obs, reward, done = pick(env, task, rubbish_pos)
        if done:
            print("[Task] Ended unexpectedly after pick.")
            return

        # Step 6: Place the rubbish in the bin
        print("[Plan] Step 6: place rubbish in bin")
        obs, reward, done = place(env, task, bin_pos)
        if done:
            print("[Task] Completed: rubbish placed in bin. Reward:", reward)
        else:
            print("[Task] Finished steps but 'done' flag is False. Reward:", reward)

    except Exception as e:
        print(f"[Error] Exception during task execution: {e}")
    finally:
        shutdown_environment(env)
        print("===== Task End =====")

if __name__ == "__main__":
    run_skeleton_task()