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 *

from video import init_video_writers, recording_step, recording_get_observation

from object_positions import get_object_positions

def run_skeleton_task():
    print("===== Starting Skeleton Task =====")
    
    # === Environment Setup ===
    env, task = setup_environment()
    try:
        descriptions, obs = task.reset()
        init_video_writers(obs)
        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()
        bottom_side_pos = positions.get('bottom_side_pos')
        bottom_anchor_pos = positions.get('bottom_anchor_pos')
        bottom_joint_pos = positions.get('bottom_joint_pos')
        if bottom_side_pos is None or bottom_anchor_pos is None or bottom_joint_pos is None:
            raise KeyError("Missing one of bottom_side_pos, bottom_anchor_pos, or bottom_joint_pos in positions")
        
        # === Oracle Plan Execution ===
        # Step 1: rotate gripper to ninety_deg orientation around Z
        angle = np.deg2rad(90.0)
        target_quat = np.array([0.0, 0.0, np.sin(angle/2.0), np.cos(angle/2.0)])
        print("[Plan] Step 1: rotate to ninety_deg")
    obs, reward, done = rotate(env, task, target_quat)
obs, reward, done = move(env, task, target_pos=bottom_side_pos)
        if done:
            print("[Plan] Terminated after step 2. Reward:", reward)
            return
        
        # Step 3: move to anchor-pos-bottom
        print("[Plan] Step 3: move to anchor-pos-bottom at", bottom_anchor_pos)
        obs, reward, done = move(env, task, target_pos=bottom_anchor_pos)
        if done:
            print("[Plan] Terminated after step 3. Reward:", reward)
            return
        
        # Step 4: pick at anchor-pos-bottom
        print("[Plan] Step 4: pick at anchor-pos-bottom")
        obs, reward, done = pick(env, task, target_pos=bottom_anchor_pos)
        if done:
            print("[Plan] Terminated after step 4. Reward:", reward)
            return
        
        # Step 5: pull drawer bottom to open
        vec = bottom_joint_pos - bottom_anchor_pos
        dist = np.linalg.norm(vec)
        idx = np.argmax(np.abs(vec))
        axis_map = {0: 'x', 1: 'y', 2: 'z'}
        axis = axis_map[idx] if vec[idx] >= 0 else '-' + axis_map[idx]
        print(f"[Plan] Step 5: pull drawer bottom by distance {dist:.3f} along axis '{axis}'")
        obs, reward, done = pull(env, task, pull_distance=dist, pull_axis=axis)
        if done:
            print("[Plan] Drawer opened. Reward:", reward)
        else:
            print("[Plan] Pull completed (done=False). Reward:", reward)
    
    finally:
        shutdown_environment(env)
    
    print("===== End of Skeleton Task =====")

if __name__ == "__main__":
    run_skeleton_task()