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 Task =====")
    
    # === Environment Setup ===
    env, task = setup_environment()
    try:
        descriptions, obs = task.reset()
        init_video_writers(obs)
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)
        
        # === Retrieve Object Positions ===
        positions = get_object_positions()
        
        # Helper to map plan names to position keys
        def lookup_pos(name):
            if name in positions:
                return np.array(positions[name])
            parts = name.split('-')
            if len(parts) == 3 and parts[0] in ('side', 'anchor'):
                _, _, suffix = parts
                key = f"{suffix}_{parts[0]}_pos"
                if key in positions:
                    return np.array(positions[key])
            raise KeyError(f"Unknown position key for '{name}'")
        
        # Precompute quaternions
        zero_deg_quat = np.array([0.0, 0.0, 0.0, 1.0])
        ninety_deg_quat = R.from_euler('z', 90, degrees=True).as_quat()
        
        # === Oracle Plan Execution ===
        # Step 1: rotate to 90°
        print("[Plan] Step 1: rotate to 90°")
        obs, reward, done = rotate(env, task, ninety_deg_quat)
        if done:
            print("Task ended early after rotate.")
            return
        
        # Step 2: move to side-pos-bottom
        print("[Plan] Step 2: move to side-pos-bottom")
        target = lookup_pos("side-pos-bottom")
        obs, reward, done = move(env, task, target)
        if done:
            print("Task ended early after move to side-pos-bottom.")
            return
        
        # Step 3: move to anchor-pos-bottom
        print("[Plan] Step 3: move to anchor-pos-bottom")
        target = lookup_pos("anchor-pos-bottom")
        obs, reward, done = move(env, task, target)
        if done:
            print("Task ended early after move to anchor-pos-bottom.")
            return
        
        # Step 4: pick drawer handle
        print("[Plan] Step 4: pick drawer handle")
        obs, reward, done = pick(env, task, lookup_pos("anchor-pos-bottom"), approach_axis='x')
        if done:
            print("Task ended early after pick drawer handle.")
            return
        
        # Step 5: pull drawer open
        print("[Plan] Step 5: pull drawer")
        obs, reward, done = pull(env, task, pull_distance=0.1, pull_axis='x')
        if done:
            print("Task ended early after pull.")
            return
        
        # Step 6: move to table
        print("[Plan] Step 6: move to table")
        target = lookup_pos("table")
        obs, reward, done = move(env, task, target)
        if done:
            print("Task ended early after move to table.")
            return
        
        # Step 7: pick tomato1
        print("[Plan] Step 7: pick tomato1")
        obs, reward, done = pick(env, task, lookup_pos("tomato1"))
        if done:
            print("Task ended early after pick tomato1.")
            return
        
        # Step 8: move to plate
        print("[Plan] Step 8: move to plate")
        target = lookup_pos("plate")
        obs, reward, done = move(env, task, target)
        if done:
            print("Task ended early after move to plate (tomato1).")
            return
        
        # Step 9: place tomato1
        print("[Plan] Step 9: place tomato1 on plate")
        obs, reward, done = place(env, task, lookup_pos("plate"))
        if done:
            print("Task ended early after place tomato1.")
            return
        
        # Step 10: move to table
        print("[Plan] Step 10: move to table for tomato2")
        target = lookup_pos("table")
        obs, reward, done = move(env, task, target)
        if done:
            print("Task ended early after move to table (tomato2).")
            return
        
        # Step 11: pick tomato2
        print("[Plan] Step 11: pick tomato2")
        obs, reward, done = pick(env, task, lookup_pos("tomato2"))
        if done:
            print("Task ended early after pick tomato2.")
            return
        
        # Step 12: move to plate
        print("[Plan] Step 12: move to plate for tomato2")
        target = lookup_pos("plate")
        obs, reward, done = move(env, task, target)
        if done:
            print("Task ended early after move to plate (tomato2).")
            return
        
        # Step 13: place tomato2
        print("[Plan] Step 13: place tomato2 on plate")
        obs, reward, done = place(env, task, lookup_pos("plate"))
        if done:
            print("Task ended early after place tomato2.")
            return
        
        print("===== Task Completed Successfully! Reward:", reward)
        
    except Exception as e:
        print("Error during task execution:", e)
    finally:
        shutdown_environment(env)
        print("===== Environment Shutdown =====")

if __name__ == "__main__":
    run_skeleton_task()