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: Pull open any unlocked drawer, then drop the rubbish into 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 name mapping (from object list and environment) ---
        # Drawers: bottom, middle, top
        # Gripper: gripper
        # Items: item1, item2, item3 (rubbish is likely item3)
        # Bin: bin
        # Drawer anchor/side positions: e.g., bottom_anchor_pos, bottom_side_pos, etc.

        # For this plan, we will:
        # 1. Move gripper to the side position of an unlocked drawer (choose 'bottom' for simplicity)
        # 2. Rotate gripper to 90 degrees (if not already)
        # 3. Move to anchor position of the drawer
        # 4. Pick the drawer handle (simulate grasp)
        # 5. Pull the drawer open
        # 6. Move to the rubbish (item3) on the table
        # 7. Pick the rubbish
        # 8. Move to the bin
        # 9. Place the rubbish in the bin

        # --- Step 1: Move gripper to side position of bottom drawer ---
        try:
            bottom_side_pos = positions['bottom_side_pos']
        except KeyError:
            print("[Error] Could not find 'bottom_side_pos' in object positions.")
            shutdown_environment(env)
            return

        obs, reward, done = move(
            env,
            task,
            target_pos=bottom_side_pos,
            max_steps=100,
            threshold=0.01,
            timeout=10.0
        )
        if done:
            print("[Task] Task ended after move to bottom_side_pos!")
            return

        # --- Step 2: Rotate gripper to 90 degrees (for drawer grasp) ---
        # Assume 90 deg = [0, 0, sin(45°), cos(45°)] in quaternion (z axis)
        # But get the actual quaternion from the environment if possible
        # For RLBench, let's use a helper if available, else use a standard z-rotation
        from scipy.spatial.transform import Rotation as R
        ninety_deg_quat = R.from_euler('z', 90, degrees=True).as_quat()  # xyzw

        obs, reward, done = rotate(
            env,
            task,
            target_quat=ninety_deg_quat,
            max_steps=100,
            threshold=0.05,
            timeout=10.0
        )
        if done:
            print("[Task] Task ended after rotate to 90 degrees!")
            return

        # --- Step 3: Move to anchor position of bottom drawer ---
        try:
            bottom_anchor_pos = positions['bottom_anchor_pos']
        except KeyError:
            print("[Error] Could not find 'bottom_anchor_pos' in object positions.")
            shutdown_environment(env)
            return

        obs, reward, done = move(
            env,
            task,
            target_pos=bottom_anchor_pos,
            max_steps=100,
            threshold=0.01,
            timeout=10.0
        )
        if done:
            print("[Task] Task ended after move to bottom_anchor_pos!")
            return

        # --- Step 4: Pick the drawer handle (simulate grasp) ---
        # Use pick skill to close gripper at anchor position
        obs, reward, done = pick(
            env,
            task,
            target_pos=bottom_anchor_pos,
            approach_distance=0.10,
            max_steps=100,
            threshold=0.01,
            approach_axis='z',
            timeout=10.0
        )
        if done:
            print("[Task] Task ended after picking drawer handle!")
            return

        # --- Step 5: Pull the drawer open ---
        # Pull along x axis (assume positive x is outwards)
        pull_distance = 0.15  # Adjust as needed for drawer
        obs, reward, done = pull(
            env,
            task,
            pull_distance=pull_distance,
            pull_axis='x',
            max_steps=100,
            threshold=0.01,
            timeout=10.0
        )
        if done:
            print("[Task] Task ended after pulling drawer!")
            return

        # --- Step 6: Move to rubbish (item3) on table ---
        try:
            item3_pos = positions['item3']
        except KeyError:
            print("[Error] Could not find 'item3' in object positions.")
            shutdown_environment(env)
            return

        obs, reward, done = move(
            env,
            task,
            target_pos=item3_pos,
            max_steps=100,
            threshold=0.01,
            timeout=10.0
        )
        if done:
            print("[Task] Task ended after move to item3!")
            return

        # --- Step 7: Pick the rubbish (item3) ---
        obs, reward, done = pick(
            env,
            task,
            target_pos=item3_pos,
            approach_distance=0.15,
            max_steps=100,
            threshold=0.01,
            approach_axis='z',
            timeout=10.0
        )
        if done:
            print("[Task] Task ended after picking item3!")
            return

        # --- Step 8: Move to bin ---
        try:
            bin_pos = positions['bin']
        except KeyError:
            print("[Error] Could not find 'bin' in object positions.")
            shutdown_environment(env)
            return

        obs, reward, done = move(
            env,
            task,
            target_pos=bin_pos,
            max_steps=100,
            threshold=0.01,
            timeout=10.0
        )
        if done:
            print("[Task] Task ended after move to bin!")
            return

        # --- Step 9: Place the rubbish in the bin ---
        obs, reward, done = place(
            env,
            task,
            target_pos=bin_pos,
            approach_distance=0.15,
            max_steps=100,
            threshold=0.01,
            approach_axis='z',
            timeout=10.0
        )
        if done:
            print("[Task] Task completed successfully! Reward:", reward)
        else:
            print("[Task] Task not completed yet (done=False).")

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

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


if __name__ == "__main__":
    run_skeleton_task()