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 *  # Use provided skills: move, pick, place, rotate, pull

from video import init_video_writers, recording_step, recording_get_observation

from object_positions import get_object_positions


def run_skeleton_task():
    '''Task: Open a drawer fully, pick up the rubbish, and place it in the bin.'''
    print("===== Starting Skeleton Task =====")
    
    # === Environment Setup ===
    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 ===
        positions = get_object_positions()
        # Object keys (from object list and likely mapping):
        # 'bottom_anchor_pos', 'bottom_joint_pos', 'bottom_side_pos',
        # 'middle_anchor_pos', 'middle_joint_pos', 'middle_side_pos',
        # 'top_anchor_pos', 'top_joint_pos', 'top_side_pos',
        # 'waypoint1', 'item1', 'item2', 'item3', 'bin'
        # We'll use 'bottom' drawer for this plan.

        # Defensive: check all required keys exist
        required_keys = [
            'bottom_side_pos', 'bottom_anchor_pos', 'bottom_joint_pos',
            'item3', 'bin'
        ]
        for k in required_keys:
            if k not in positions:
                raise RuntimeError(f"Required object position '{k}' not found in get_object_positions().")

        # Plan:
        # 1. Move gripper to the side position of the bottom drawer.
        # 2. Rotate gripper to 90 degrees (for grasping drawer handle).
        # 3. Move gripper to anchor position (handle) of the bottom drawer.
        # 4. Pull the drawer open (along x axis, e.g. +0.18m).
        # 5. Move gripper above the rubbish (item3).
        # 6. Pick up the rubbish.
        # 7. Move gripper above the bin.
        # 8. Place the rubbish in the bin.

        # Step 1: Move to side position of bottom drawer
        print("[Plan] Step 1: Move to bottom_side_pos")
        bottom_side_pos = np.array(positions['bottom_side_pos'])
        obs, reward, done = move(env, task, target_pos=bottom_side_pos)
        if done:
            print("[Task] Task ended after move to bottom_side_pos!")
            return

        # Step 2: Rotate gripper to 90 degrees (for drawer handle)
        print("[Plan] Step 2: Rotate gripper to 90 degrees")
        # We'll use quaternion for 90 deg rotation about z axis
        from scipy.spatial.transform import Rotation as R
        quat_90deg = R.from_euler('z', 90, degrees=True).as_quat()  # xyzw
        obs, reward, done = rotate(env, task, target_quat=quat_90deg)
        if done:
            print("[Task] Task ended after rotate!")
            return

        # Step 3: Move to anchor position (handle) of bottom drawer
        print("[Plan] Step 3: Move to bottom_anchor_pos")
        bottom_anchor_pos = np.array(positions['bottom_anchor_pos'])
        obs, reward, done = move(env, task, target_pos=bottom_anchor_pos)
        if done:
            print("[Task] Task ended after move to bottom_anchor_pos!")
            return

        # Step 4: Pull the drawer open (along x axis, e.g. +0.18m)
        print("[Plan] Step 4: Pull the drawer open")
        # Pull distance: positive x direction, typical drawer open distance
        pull_distance = 0.18
        obs, reward, done = pull(env, task, pull_distance=pull_distance, pull_axis='x')
        if done:
            print("[Task] Task ended after pull!")
            return

        # Step 5: Move above the rubbish (item3)
        print("[Plan] Step 5: Move above item3 (rubbish)")
        item3_pos = np.array(positions['item3'])
        # Approach from above: raise z by 0.10m for safety
        approach_offset = np.array([0, 0, 0.10])
        above_item3_pos = item3_pos + approach_offset
        obs, reward, done = move(env, task, target_pos=above_item3_pos)
        if done:
            print("[Task] Task ended after move above item3!")
            return

        # Step 6: Pick up the rubbish (item3)
        print("[Plan] Step 6: Pick up item3")
        obs, reward, done = pick(env, task, target_pos=item3_pos, approach_distance=0.10, approach_axis='-z')
        if done:
            print("[Task] Task ended after pick!")
            return

        # Step 7: Move above the bin
        print("[Plan] Step 7: Move above bin")
        bin_pos = np.array(positions['bin'])
        above_bin_pos = bin_pos + approach_offset
        obs, reward, done = move(env, task, target_pos=above_bin_pos)
        if done:
            print("[Task] Task ended after move above bin!")
            return

        # Step 8: Place the rubbish in the bin
        print("[Plan] Step 8: Place item3 in bin")
        obs, reward, done = place(env, task, target_pos=bin_pos, approach_distance=0.10, approach_axis='-z')
        if done:
            print("[Task] Task completed successfully! Reward:", reward)
        else:
            print("[Task] Task not completed yet (done=False).")

    except Exception as e:
        print(f"[Task] Exception occurred: {e}")
    finally:
        # Always ensure the environment is properly shutdown
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()