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 only the provided skills

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: see object list
        # [
        # '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',
        # 'item1', 'item2', 'item3', 'bin'
        # ]
        # We'll use the 'bottom' drawer for this plan.

        # Drawer positions
        bottom_side_pos = positions.get('bottom_side_pos', None)
        bottom_anchor_pos = positions.get('bottom_anchor_pos', None)
        bottom_joint_pos = positions.get('bottom_joint_pos', None)
        # For pulling, we assume pulling along the x-axis from anchor to joint
        # Rubbish object (item3) and bin
        rubbish_pos = positions.get('item3', None)
        bin_pos = positions.get('bin', None)

        # Defensive: check all required positions
        if None in [bottom_side_pos, bottom_anchor_pos, bottom_joint_pos, rubbish_pos, bin_pos]:
            print("[Error] One or more required object positions are missing. Aborting task.")
            return

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

        # Step 2: Rotate gripper to 90 degrees (for proper grasping orientation)
        print("[Step 2] Rotate gripper to 90 degrees.")
        # We'll assume the target orientation is provided as a quaternion in positions dict
        # If not, we can construct it: 90 deg about z axis
        from scipy.spatial.transform import Rotation as R
        target_rot = R.from_euler('z', 90, degrees=True)
        target_quat = target_rot.as_quat()  # xyzw
        obs, reward, done = rotate(
            env, task, target_quat=target_quat
        )
        if done:
            print("[Task] Task ended after rotate!")
            return

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

        # Step 4: Pick (grasp) the drawer handle at anchor position
        print("[Step 4] Pick (grasp) the drawer handle at anchor position.")
        # For pick, we use the anchor position, approach from above (z)
        obs, reward, done = pick(
            env, task, target_pos=np.array(bottom_anchor_pos),
            approach_distance=0.10, approach_axis='z'
        )
        if done:
            print("[Task] Task ended after pick drawer handle!")
            return

        # Step 5: Pull the drawer open (fully)
        print("[Step 5] Pull the drawer open.")
        # Pull along x-axis from anchor to joint; distance = |joint - anchor|
        pull_vec = np.array(bottom_joint_pos) - np.array(bottom_anchor_pos)
        pull_distance = np.linalg.norm(pull_vec)
        pull_axis = 'x' if pull_vec[0] >= 0 else '-x'
        obs, reward, done = pull(
            env, task, pull_distance=pull_distance, pull_axis=pull_axis
        )
        if done:
            print("[Task] Task ended after pull!")
            return

        # Step 6: Move to the rubbish (item3) position
        print("[Step 6] Move to rubbish position.")
        obs, reward, done = move(
            env, task, target_pos=np.array(rubbish_pos)
        )
        if done:
            print("[Task] Task ended after move to rubbish!")
            return

        # Step 7: Rotate gripper to 0 degrees (upright) for picking up rubbish
        print("[Step 7] Rotate gripper to 0 degrees.")
        target_rot = R.from_euler('z', 0, degrees=True)
        target_quat = target_rot.as_quat()
        obs, reward, done = rotate(
            env, task, target_quat=target_quat
        )
        if done:
            print("[Task] Task ended after rotate to 0 deg!")
            return

        # Step 8: Move slightly above the rubbish for approach
        print("[Step 8] Move above rubbish for approach.")
        above_rubbish = np.array(rubbish_pos) + np.array([0, 0, 0.10])
        obs, reward, done = move(
            env, task, target_pos=above_rubbish
        )
        if done:
            print("[Task] Task ended after move above rubbish!")
            return

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

        # Step 10: Move to the bin
        print("[Step 10] Move to bin position.")
        obs, reward, done = move(
            env, task, target_pos=np.array(bin_pos)
        )
        if done:
            print("[Task] Task ended after move to bin!")
            return

        # Step 11: Place the rubbish in the bin
        print("[Step 11] Place the rubbish in the bin.")
        # Approach from above
        obs, reward, done = place(
            env, task, target_pos=np.array(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("[Exception] An error occurred during task execution:", str(e))
    finally:
        # Always ensure the environment is properly shutdown
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()