# run_skeleton_task.py

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():
    '''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 = get_object_positions()
        # Object names from 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'
        # ]

        # For this task:
        # 1. Open an accessible drawer (cabinet unlock = open a drawer)
        # 2. Throw away the trash (one of item1, item2, item3) into the bin
        # 3. Leave other objects alone

        # We'll use the bottom drawer (most accessible) and item3 as trash (arbitrary choice)
        # If you want to randomize, you can, but here we pick item3 for determinism

        # Drawer positions
        bottom_side_pos = positions.get('bottom_side_pos')
        bottom_anchor_pos = positions.get('bottom_anchor_pos')
        bottom_joint_pos = positions.get('bottom_joint_pos')
        # Gripper will approach side, then anchor, then pull

        # Trash and bin positions
        item3_pos = positions.get('item3')
        bin_pos = positions.get('bin')

        # Plan steps (from specification):
        # 1. move
        # 2. rotate
        # 3. move
        # 4. pick
        # 5. pull
        # 6. move
        # 7. rotate
        # 8. move
        # 9. pick
        # 10. move
        # 11. place

        # Step 1: Move gripper to bottom_side_pos (side of bottom drawer)
        print("[Step 1] Move gripper to bottom_side_pos")
        try:
            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] Move to bottom_side_pos failed: {e}")
            return

        # Step 2: Rotate gripper to 90 degrees (for drawer grasp)
        print("[Step 2] Rotate gripper to 90 degrees")
        try:
            # Assume 90 deg = [x, y, z, w] quaternion for 90 deg about z
            # We'll use scipy to get quaternion for 90 deg about z
            from scipy.spatial.transform import Rotation as R
            quat_90deg = R.from_euler('z', 90, degrees=True).as_quat()  # xyzw
            obs, reward, done = rotate(env, task, target_quat=quat_90deg)
            if done:
                print("[Task] Task ended after rotate to 90 deg!")
                return
        except Exception as e:
            print(f"[Error] Rotate to 90 deg failed: {e}")
            return

        # Step 3: Move gripper to bottom_anchor_pos (front of drawer handle)
        print("[Step 3] Move gripper to bottom_anchor_pos")
        try:
            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] Move to bottom_anchor_pos failed: {e}")
            return

        # Step 4: Pick the drawer handle (simulate grasp)
        print("[Step 4] Pick drawer handle (bottom_anchor_pos)")
        try:
            obs, reward, done = pick(env, task, target_pos=bottom_anchor_pos, approach_distance=0.08, approach_axis='y')
            if done:
                print("[Task] Task ended after pick drawer handle!")
                return
        except Exception as e:
            print(f"[Error] Pick drawer handle failed: {e}")
            return

        # Step 5: Pull the drawer open (pull along -y axis, e.g., 0.15m)
        print("[Step 5] Pull drawer open")
        try:
            obs, reward, done = pull(env, task, pull_distance=0.15, pull_axis='-y')
            if done:
                print("[Task] Task ended after pull!")
                return
        except Exception as e:
            print(f"[Error] Pull drawer failed: {e}")
            return

        # Step 6: Move gripper up and away from drawer (to avoid collision)
        print("[Step 6] Move gripper up (z+)")
        try:
            # Move to a safe position above the drawer
            safe_pos = np.array(bottom_anchor_pos) + np.array([0, 0, 0.10])
            obs, reward, done = move(env, task, target_pos=safe_pos)
            if done:
                print("[Task] Task ended after move up!")
                return
        except Exception as e:
            print(f"[Error] Move up after pull failed: {e}")
            return

        # Step 7: Rotate gripper back to 0 deg (neutral orientation)
        print("[Step 7] Rotate gripper to 0 deg")
        try:
            quat_0deg = R.from_euler('z', 0, degrees=True).as_quat()
            obs, reward, done = rotate(env, task, target_quat=quat_0deg)
            if done:
                print("[Task] Task ended after rotate to 0 deg!")
                return
        except Exception as e:
            print(f"[Error] Rotate to 0 deg failed: {e}")
            return

        # Step 8: Move gripper to above item3 (trash)
        print("[Step 8] Move gripper to above item3")
        try:
            above_item3 = np.array(item3_pos) + np.array([0, 0, 0.10])
            obs, reward, done = move(env, task, target_pos=above_item3)
            if done:
                print("[Task] Task ended after move above item3!")
                return
        except Exception as e:
            print(f"[Error] Move above item3 failed: {e}")
            return

        # Step 9: Pick item3 (trash)
        print("[Step 9] Pick item3 (trash)")
        try:
            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] Pick item3 failed: {e}")
            return

        # Step 10: Move gripper to above bin
        print("[Step 10] Move gripper to above bin")
        try:
            above_bin = np.array(bin_pos) + np.array([0, 0, 0.10])
            obs, reward, done = move(env, task, target_pos=above_bin)
            if done:
                print("[Task] Task ended after move above bin!")
                return
        except Exception as e:
            print(f"[Error] Move above bin failed: {e}")
            return

        # Step 11: Place item3 in bin
        print("[Step 11] Place item3 in bin")
        try:
            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] 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()
