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: pick, place, move, 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: Slide open any available drawer to access its contents, then chuck away any rubbish remaining on the table.'''
    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 names from object list and environment:
        # Drawers: 'bottom', 'middle', 'top'
        # Drawer positions: e.g., 'bottom_anchor_pos', 'bottom_side_pos', etc.
        # Items: 'item1', 'item2', 'item3'
        # Bin: 'bin'
        # Gripper: assumed to be controlled via skills

        # --- Step 1: Rotate gripper to 90 degrees (for side approach to drawer) ---
        # We'll use the bottom drawer as the "available" drawer (not locked, or at least not top, which is locked)
        # From the initial state, bottom is not locked, so we use it.
        # The rotate skill expects a quaternion. Let's get the quaternion for 90 deg rotation about z.
        # We'll use scipy's Rotation to get the quaternion.
        from scipy.spatial.transform import Rotation as R
        try:
            # Get current gripper pose
            obs = task.get_observation()
            start_quat = obs.gripper_pose[3:7]
            # Target: rotate 90 deg about z axis
            target_rot = R.from_euler('z', 90, degrees=True)
            target_quat = target_rot.as_quat()  # xyzw
            obs, reward, done = rotate(env, task, target_quat)
            if done:
                print("[Task] Task ended after rotate!")
                return
        except Exception as e:
            print(f"[Task] Exception during rotate: {e}")
            shutdown_environment(env)
            return

        # --- Step 2: Move to side position of bottom drawer ---
        try:
            side_pos = positions.get('bottom_side_pos', None)
            if side_pos is None:
                raise KeyError("bottom_side_pos not found in positions.")
            obs, reward, done = move(env, task, side_pos)
            if done:
                print("[Task] Task ended after move to side position!")
                return
        except Exception as e:
            print(f"[Task] Exception during move to side position: {e}")
            shutdown_environment(env)
            return

        # --- Step 3: Move to anchor position of bottom drawer ---
        try:
            anchor_pos = positions.get('bottom_anchor_pos', None)
            if anchor_pos is None:
                raise KeyError("bottom_anchor_pos not found in positions.")
            obs, reward, done = move(env, task, anchor_pos)
            if done:
                print("[Task] Task ended after move to anchor position!")
                return
        except Exception as e:
            print(f"[Task] Exception during move to anchor position: {e}")
            shutdown_environment(env)
            return

        # --- Step 4: Pick the drawer handle (simulate grasp) ---
        try:
            # For picking the drawer, we use the anchor position as the target
            obs, reward, done = pick(env, task, anchor_pos, approach_distance=0.10, approach_axis='z')
            if done:
                print("[Task] Task ended after pick drawer!")
                return
        except Exception as e:
            print(f"[Task] Exception during pick drawer: {e}")
            shutdown_environment(env)
            return

        # --- Step 5: Pull the drawer open ---
        try:
            # Pull along x axis (assume drawer opens along +x, adjust if needed)
            pull_distance = 0.15  # Adjust as needed for your environment
            obs, reward, done = pull(env, task, pull_distance, pull_axis='x')
            if done:
                print("[Task] Task ended after pull!")
                return
        except Exception as e:
            print(f"[Task] Exception during pull: {e}")
            shutdown_environment(env)
            return

        # --- Step 6: Pick up rubbish (item) from table ---
        # We'll pick the first available item on the table (item1)
        try:
            item_name = None
            for candidate in ['item1', 'item2', 'item3']:
                if candidate in positions:
                    item_name = candidate
                    break
            if item_name is None:
                raise KeyError("No item found on table to pick.")
            item_pos = positions[item_name]
            obs, reward, done = pick(env, task, item_pos, approach_distance=0.15, approach_axis='z')
            if done:
                print("[Task] Task ended after pick item!")
                return
        except Exception as e:
            print(f"[Task] Exception during pick item: {e}")
            shutdown_environment(env)
            return

        # --- Step 7: Place the item in the bin ---
        try:
            bin_pos = positions.get('bin', None)
            if bin_pos is None:
                raise KeyError("bin not found in positions.")
            obs, reward, done = place(env, task, bin_pos, approach_distance=0.15, 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 during place in bin: {e}")
            shutdown_environment(env)
            return

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

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


if __name__ == "__main__":
    run_skeleton_task()