# 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 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():
    '''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 keys: 'rubbish', 'tomato1', 'tomato2', 'bin', 'bottom_anchor_pos', 'bottom_side_pos', etc.

        # === PLAN EXECUTION ===
        # Goal: Open a drawer fully, pick up the rubbish, and place it in the bin.

        # --- Step 1: Rotate gripper to 90 degrees (vertical for drawer handle) ---
        # We'll use the 'bottom' drawer for this plan.
        # Assume 'ninety_deg' means gripper z-axis is vertical (RLBench: [0, 0, 0.7071, 0.7071])
        # If object_positions provides a quaternion for 'ninety_deg', use it; else, use a standard one.
        try:
            # Try to get a canonical vertical quaternion from the environment, else use default
            if 'ninety_deg' in positions:
                target_quat = np.array(positions['ninety_deg'])
            else:
                # Standard RLBench vertical quaternion (z-up): [0, 0, 0.7071, 0.7071]
                target_quat = np.array([0, 0, 0.7071, 0.7071])
        except Exception as e:
            print("[Task] Could not get 'ninety_deg' quaternion, using default vertical.")
            target_quat = np.array([0, 0, 0.7071, 0.7071])

        print("[Task] Rotating gripper to vertical (90 deg)...")
        obs, reward, done = rotate(env, task, target_quat)
        if done:
            print("[Task] Task ended during rotate!")
            return

        # --- Step 2: Move to side position of bottom drawer ---
        # Move gripper to the side position of the bottom drawer (for anchoring)
        try:
            side_pos = positions['bottom_side_pos']
        except Exception as e:
            print("[Task] Could not find 'bottom_side_pos' in positions.")
            shutdown_environment(env)
            return

        print("[Task] Moving to bottom drawer side position...")
        obs, reward, done = move(env, task, np.array(side_pos))
        if done:
            print("[Task] Task ended during move to side position!")
            return

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

        print("[Task] Moving to bottom drawer anchor position...")
        obs, reward, done = move(env, task, np.array(anchor_pos))
        if done:
            print("[Task] Task ended during move to anchor position!")
            return

        # --- Step 4: Pick (grasp) the drawer handle (anchor position) ---
        # Use pick at the anchor position (simulate grasping the handle)
        print("[Task] Grasping the bottom drawer handle (anchor)...")
        obs, reward, done = pick(env, task, np.array(anchor_pos), approach_distance=0.10, approach_axis='z')
        if done:
            print("[Task] Task ended during pick (drawer handle)!")
            return

        # --- Step 5: Pull the drawer open (along x axis, e.g. +x) ---
        # Pull distance: estimate 0.20m (adjust as needed for your environment)
        pull_distance = 0.20
        print("[Task] Pulling the bottom drawer open...")
        obs, reward, done = pull(env, task, pull_distance, pull_axis='x')
        if done:
            print("[Task] Task ended during pull (drawer)!")
            return

        # --- Step 6: Pick up the rubbish from the table ---
        try:
            rubbish_pos = positions['rubbish']
        except Exception as e:
            print("[Task] Could not find 'rubbish' in positions.")
            shutdown_environment(env)
            return

        print("[Task] Picking up the rubbish from the table...")
        obs, reward, done = pick(env, task, np.array(rubbish_pos), approach_distance=0.15, approach_axis='z')
        if done:
            print("[Task] Task ended during pick (rubbish)!")
            return

        # --- Step 7: Place the rubbish in the bin ---
        try:
            bin_pos = positions['bin']
        except Exception as e:
            print("[Task] Could not find 'bin' in positions.")
            shutdown_environment(env)
            return

        print("[Task] Placing the rubbish in the bin...")
        obs, reward, done = place(env, task, np.array(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).")

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

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


if __name__ == "__main__":
    run_skeleton_task()
