import numpy as np
from env import setup_environment, shutdown_environment
from skill_code import move, rotate, pull, normalize_quaternion
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions
from scipy.spatial.transform import Rotation as R

def run_open_drawer_task():
    print("===== Starting OpenDrawer Task =====")
    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)
        
        positions = get_object_positions()
        drawer_option = 'bottom'
        anchor_pos = positions[f'{drawer_option}_anchor_pos']
        side_pos = positions[f'{drawer_option}_side_pos']
        
        print(f"[Task] Target drawer: {drawer_option}")
        print(f"[Task] Anchor position: {anchor_pos}")
        
        initial_quat = normalize_quaternion(obs.gripper_pose[3:7])
        initial_rot = R.from_quat(initial_quat)
        relative_rot = R.from_euler('x', 90, degrees=True)
        target_rot = initial_rot * relative_rot
        target_quat = target_rot.as_quat()
        
        obs, reward, done = rotate(
            env, task,
            target_quat=target_quat,
            max_steps=200,
            threshold=0.1,
            timeout=20.0
        )
        if done:
            print("[Task] Task ended after initial rotation!")
            return
        
        print("[Task] Moving to side position far left of handle:", side_pos)
        obs, reward, done = move(
            env, task,
            target_pos=side_pos,
            max_steps=100,
            threshold=0.01,
            timeout=10.0
        )
        if done:
            print("[Task] Task ended after moving to side!")
            return
        
        print("[Task] Moving horizontally to handle at:", anchor_pos)
        obs, reward, done = move(
            env, task,
            target_pos=anchor_pos,
            max_steps=100,
            threshold=0.01,
            timeout=10.0
        )
        if done:
            print("[Task] Task ended after moving to handle!")
            return
        
        print("[Task] Closing gripper to grasp handle")
        action = np.zeros(env.action_shape)
        action[-1] = -1.0
        current_quat = normalize_quaternion(obs.gripper_pose[3:7])
        for _ in range(10):
            action[:3] = anchor_pos
            action[3:7] = current_quat
            obs, reward, done = task.step(action)
            if done:
                print("[Task] Task ended after grasping!")
                return
        
        print("[Task] Pulling the drawer open")
        obs, reward, done = pull(
            env, task,
            pull_distance=0.2,
            pull_axis='y',
            max_steps=100,
            threshold=0.01,
            timeout=10.0
        )
        if done:
            print("[Task] Task completed successfully! Reward:", reward)
        else:
            print("[Task] Task not completed yet (done=False).")

    finally:
        shutdown_environment(env)

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

if __name__ == "__main__":
    run_open_drawer_task()