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/position name mapping (from object list and environment) ---
        # Drawers: bottom, middle, top
        # We'll open the "bottom" drawer for this plan.
        # Positions for bottom drawer:
        #   - bottom_side_pos: side position for approach
        #   - bottom_anchor_pos: anchor position for grasping
        #   - bottom_joint_pos: likely the center of the drawer (for pulling)
        #   - waypoint1: generic waypoint (could be used for intermediate moves)
        # Objects: rubbish, tomato1, tomato2, bin

        # Defensive: check all required keys exist
        required_keys = [
            'bottom_side_pos', 'bottom_anchor_pos', 'bottom_joint_pos',
            'rubbish', 'bin'
        ]
        for k in required_keys:
            if k not in positions:
                raise RuntimeError(f"Required object/position '{k}' not found in get_object_positions().")

        # Plan:
        # 1. rotate: rotate gripper to 90 deg (vertical) for drawer grasp
        # 2. move: move to bottom_side_pos (side of drawer)
        # 3. move: move to bottom_anchor_pos (grasp point on drawer)
        # 4. pick: grasp the drawer handle at anchor_pos
        # 5. pull: pull the drawer open (along x axis, positive direction)
        # 6. pick: pick up the rubbish from the table
        # 7. place: place the rubbish in the bin

        # Step 1: rotate gripper to 90 deg (vertical)
        # We'll use a quaternion for 90 deg rotation about x axis.
        # RLBench uses [x, y, z, w] (xyzw) format.
        # 0 deg: [0, 0, 0, 1]
        # 90 deg about x: [0.7071, 0, 0, 0.7071]
        ninety_deg_quat = np.array([0.7071, 0, 0, 0.7071])
        try:
            obs, reward, done = rotate(env, task, target_quat=ninety_deg_quat)
            if done:
                print("[Task] Task ended after rotate!")
                return
        except Exception as e:
            print(f"[Task] Exception during rotate: {e}")
            return

        # Step 2: move to bottom_side_pos
        try:
            target_pos = positions['bottom_side_pos']
            obs, reward, done = move(env, task, target_pos=target_pos)
            if done:
                print("[Task] Task ended after move to bottom_side_pos!")
                return
        except Exception as e:
            print(f"[Task] Exception during move to bottom_side_pos: {e}")
            return

        # Step 3: move to bottom_anchor_pos
        try:
            target_pos = positions['bottom_anchor_pos']
            obs, reward, done = move(env, task, target_pos=target_pos)
            if done:
                print("[Task] Task ended after move to bottom_anchor_pos!")
                return
        except Exception as e:
            print(f"[Task] Exception during move to bottom_anchor_pos: {e}")
            return

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

        # Step 5: pull the drawer open
        try:
            # Pull along x axis, positive direction. Pull distance: 0.18m (tunable)
            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"[Task] Exception during pull (drawer): {e}")
            return

        # Step 6: pick up the rubbish from the table
        try:
            # The rubbish is on the table; get its position
            obs, reward, done = pick(env, task, target_pos=positions['rubbish'], approach_distance=0.15, approach_axis='z')
            if done:
                print("[Task] Task ended after pick (rubbish)!")
                return
        except Exception as e:
            print(f"[Task] Exception during pick (rubbish): {e}")
            return

        # Step 7: place the rubbish in the bin
        try:
            obs, reward, done = place(env, task, target_pos=positions['bin'], 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 (rubbish in bin): {e}")
            return

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

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


if __name__ == "__main__":
    run_skeleton_task()