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():
    '''Generic skeleton for running any task in your simulation.'''
    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: dict mapping object names to their 3D positions
        positions = get_object_positions()

        # === PLAN EXECUTION ===
        # Goal: Open one of the drawers in the cabinet, then put the rubbish in the bin.

        # For this plan, we will:
        # 1. Move gripper to the side position of the bottom drawer
        # 2. Rotate gripper to 90 degrees (to align for grasping drawer handle)
        # 3. Move gripper to the anchor position of the bottom drawer
        # 4. Pick (grasp) the drawer handle at anchor position
        # 5. Pull the drawer open
        # 6. Move gripper above the rubbish (item3) on the table
        # 7. Pick the rubbish (item3)
        # 8. Move gripper above the bin
        # 9. Place the rubbish in the bin

        # --- Step 1: Move to side position of bottom drawer ---
        try:
            bottom_side_pos = positions.get('bottom_side_pos')
            if bottom_side_pos is None:
                raise KeyError("bottom_side_pos not found in object positions.")
            print("[Task] Moving to bottom_side_pos:", 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
        except Exception as e:
            print(f"[Error] Step 1 (move to bottom_side_pos) failed: {e}")
            return

        # --- Step 2: Rotate gripper to 90 degrees (align for drawer handle) ---
        try:
            # 90 degrees about z axis: quaternion [x, y, z, w]
            # Euler (0, 0, 90 deg) in radians: (0, 0, pi/2)
            from scipy.spatial.transform import Rotation as R
            quat_ninety_deg = R.from_euler('z', 90, degrees=True).as_quat()  # xyzw
            print("[Task] Rotating gripper to 90 degrees (z axis):", quat_ninety_deg)
            obs, reward, done = rotate(env, task, target_quat=quat_ninety_deg)
            if done:
                print("[Task] Task ended after rotate!")
                return
        except Exception as e:
            print(f"[Error] Step 2 (rotate to 90 deg) failed: {e}")
            return

        # --- Step 3: Move to anchor position of bottom drawer ---
        try:
            bottom_anchor_pos = positions.get('bottom_anchor_pos')
            if bottom_anchor_pos is None:
                raise KeyError("bottom_anchor_pos not found in object positions.")
            print("[Task] Moving to bottom_anchor_pos:", 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
        except Exception as e:
            print(f"[Error] Step 3 (move to bottom_anchor_pos) failed: {e}")
            return

        # --- Step 4: Pick (grasp) the drawer handle at anchor position ---
        try:
            # We assume the drawer handle is at bottom_anchor_pos
            print("[Task] Picking (grasping) bottom drawer handle at anchor position.")
            obs, reward, done = pick(env, task, target_pos=bottom_anchor_pos, approach_distance=0.10, approach_axis='y')
            if done:
                print("[Task] Task ended after pick (drawer handle)!")
                return
        except Exception as e:
            print(f"[Error] Step 4 (pick drawer handle) failed: {e}")
            return

        # --- Step 5: Pull the drawer open ---
        try:
            # Pull along x axis (assume drawer opens along +x, adjust if needed)
            print("[Task] Pulling the bottom drawer open.")
            obs, reward, done = pull(env, task, pull_distance=0.18, pull_axis='x')
            if done:
                print("[Task] Task ended after pull (drawer)!")
                return
        except Exception as e:
            print(f"[Error] Step 5 (pull drawer) failed: {e}")
            return

        # --- Step 6: Move above the rubbish (item3) on the table ---
        try:
            item3_pos = positions.get('item3')
            if item3_pos is None:
                raise KeyError("item3 not found in object positions.")
            # Raise z to approach from above
            item3_above = np.array(item3_pos) + np.array([0, 0, 0.10])
            print("[Task] Moving above item3 (rubbish):", item3_above)
            obs, reward, done = move(env, task, target_pos=item3_above)
            if done:
                print("[Task] Task ended after move above item3!")
                return
        except Exception as e:
            print(f"[Error] Step 6 (move above item3) failed: {e}")
            return

        # --- Step 7: Pick the rubbish (item3) ---
        try:
            print("[Task] Picking item3 (rubbish) from table.")
            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 (item3)!")
                return
        except Exception as e:
            print(f"[Error] Step 7 (pick item3) failed: {e}")
            return

        # --- Step 8: Move above the bin ---
        try:
            bin_pos = positions.get('bin')
            if bin_pos is None:
                raise KeyError("bin not found in object positions.")
            bin_above = np.array(bin_pos) + np.array([0, 0, 0.15])
            print("[Task] Moving above bin:", bin_above)
            obs, reward, done = move(env, task, target_pos=bin_above)
            if done:
                print("[Task] Task ended after move above bin!")
                return
        except Exception as e:
            print(f"[Error] Step 8 (move above bin) failed: {e}")
            return

        # --- Step 9: Place the rubbish in the bin ---
        try:
            print("[Task] Placing item3 (rubbish) 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"[Error] Step 9 (place item3 in bin) failed: {e}")
            return

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

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


if __name__ == "__main__":
    run_skeleton_task()