# run_skeleton_task.py
#
# Description:
#  - This code sets up the RLBench environment, retrieves object positions,
#    and executes a hard‑coded oracle plan for opening a drawer and moving
#    two tomatoes onto a plate using predefined skills.

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 move, rotate, pull, pick, place

from video import init_video_writers, recording_step, recording_get_observation

from object_positions import get_object_positions


def run_skeleton_task():
    '''Runs the specific task plan to open a drawer and place tomatoes on a plate.'''
    print("===== Starting Skeleton Task =====")
    
    env, task = setup_environment()
    try:
        descriptions, obs = task.reset()
        init_video_writers(obs)
        # wrap step/get_observation for recording
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)
        
        # Retrieve 3D positions of relevant objects
        positions = get_object_positions()
        # Drawer approach positions
        bottom_side_pos = np.array(positions['bottom_side_pos'])
        bottom_anchor_pos = np.array(positions['bottom_anchor_pos'])
        # Tomatoes (renamed as item1/item2 in the scene)
        tomato1_pos = np.array(positions['item1'])
        tomato2_pos = np.array(positions['item2'])
        # Plate position
        plate_pos = np.array(positions['plate'])
        
        # === Oracle Plan Execution ===
        
        # Step 1: Move gripper to side position of bottom drawer
        print("[Task] Step 1: move to bottom side position")
        obs, reward, done = move(env, task, target_pos=bottom_side_pos)
        if done:
            print("[Task] Terminated early after step 1")
            return
        
        # Step 2: Rotate gripper from zero to ninety degrees around Z axis
        print("[Task] Step 2: rotate gripper to ninety degrees")
        # Define quaternion for 90° around Z: [x,y,z,w] (sin(45°)=.7071, cos(45°)=.7071)
        target_quat = np.array([0.0, 0.0, np.sin(np.pi/4), np.cos(np.pi/4)])
        obs, reward, done = rotate(env, task, target_quat=target_quat)
        if done:
            print("[Task] Terminated early after step 2")
            return
        
        # Step 3: Move gripper to anchor position of bottom drawer
        print("[Task] Step 3: move to bottom anchor position")
        obs, reward, done = move(env, task, target_pos=bottom_anchor_pos)
        if done:
            print("[Task] Terminated early after step 3")
            return
        
        # Step 4: Pull the drawer open
        print("[Task] Step 4: pull drawer open")
        # Pull distance chosen experimentally (e.g., 0.2m) along the local X axis
        obs, reward, done = pull(env, task, pull_distance=0.2, pull_axis='x')
        if done:
            print("[Task] Terminated early after step 4")
            return
        
        # Step 5: Pick up the first tomato (item1)
        print("[Task] Step 5: pick first tomato")
        obs, reward, done = pick(env, task, target_pos=tomato1_pos)
        if done:
            print("[Task] Terminated early after step 5")
            return
        
        # Step 6: Place the first tomato on the plate
        print("[Task] Step 6: place first tomato on plate")
        obs, reward, done = place(env, task, target_pos=plate_pos)
        if done:
            print("[Task] Terminated early after step 6")
            return
        
        # Step 7: Pick up the second tomato (item2)
        print("[Task] Step 7: pick second tomato")
        obs, reward, done = pick(env, task, target_pos=tomato2_pos)
        if done:
            print("[Task] Terminated early after step 7")
            return
        
        # Step 8: Place the second tomato on the plate
        print("[Task] Step 8: place second tomato on plate")
        obs, reward, done = place(env, task, target_pos=plate_pos)
        if done:
            print("[Task] Completed all steps with done=True, reward:", reward)
        else:
            print("[Task] All steps executed, final done flag:", done, "reward:", reward)
    
    except Exception as e:
        print("Error during task execution:", repr(e))
    
    finally:
        shutdown_environment(env)
        print("===== End of Skeleton Task =====")


if __name__ == "__main__":
    run_skeleton_task()