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, then put the rubbish 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 names from object list and environment
        # Drawers: bottom, middle, top
        # Items: item1, item2, item3 (assume item3 is rubbish as per observation)
        # Bin: bin
        # Gripper: gripper (not in positions, but gripper_pose in obs)
        # Drawer anchor/side/joint positions: e.g., bottom_anchor_pos, bottom_side_pos, etc.

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

        # --- Step 1: Move to side of bottom drawer ---
        try:
            bottom_side_pos = positions['bottom_side_pos']
        except KeyError:
            print("[Error] bottom_side_pos not found in positions. Aborting.")
            return

        print("[Task] Moving to side of bottom drawer at:", bottom_side_pos)
        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 drawer side!")
            return

        # --- Step 2: Rotate gripper to 90 degrees (for grasping handle) ---
        # Assume 90 deg is [0, 0, sin(45°), cos(45°)] in xyzw (z axis)
        # But let's try to get the quaternion for 90 deg about z
        from scipy.spatial.transform import Rotation as R
        quat_ninety_z = R.from_euler('z', 90, degrees=True).as_quat()  # xyzw
        print("[Task] Rotating gripper to 90 degrees about z axis:", quat_ninety_z)
        obs, reward, done = rotate(
            env,
            task,
            target_quat=quat_ninety_z,
            max_steps=100,
            threshold=0.05,
            timeout=10.0
        )
        if done:
            print("[Task] Task ended after rotate!")
            return

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

        print("[Task] Moving to anchor position of bottom drawer at:", bottom_anchor_pos)
        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 anchor!")
            return

        # --- Step 4: Pick the drawer handle (simulate pick at anchor) ---
        print("[Task] Picking the drawer handle 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, adjust if needed)
        print("[Task] Pulling the drawer open along x axis.")
        obs, reward, done = pull(
            env,
            task,
            pull_distance=0.15,  # 15cm pull, adjust as needed for your sim
            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] item3 (rubbish) not found in positions. Aborting.")
            return

        print("[Task] Moving to rubbish (item3) at:", item3_pos)
        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 rubbish!")
            return

        # --- Step 7: Pick the rubbish (item3) ---
        print("[Task] Picking 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 rubbish!")
            return

        # --- Step 8: Move to bin ---
        try:
            bin_pos = positions['bin']
        except KeyError:
            print("[Error] bin not found in positions. Aborting.")
            return

        print("[Task] Moving to bin at:", bin_pos)
        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 rubbish in bin ---
        print("[Task] Placing rubbish in 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()