import numpy as np
from scipy.spatial.transform import Rotation as R

# === KEEP ORIGINAL IMPORTS INTACT ===
from pyrep.objects.shape import Shape           # (needed by RLBench internals)
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment

# Import skills AFTER the env-related imports
from skill_code import rotate, move, pick, pull, place

from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions


# ---------------------------------------------------------------------
#  Helper utilities
# ---------------------------------------------------------------------
def quat_from_euler_xyz(x=0.0, y=0.0, z=0.0, rad=False):
    """
    Convert XYZ Euler angles to an xyzw quaternion.

    Args:
        x, y, z (float): Euler angles.
        rad (bool): If True, the angles are already in radians.

    Returns:
        np.ndarray(4,): Quaternion (xyzw).
    """
    if not rad:
        x, y, z = np.deg2rad([x, y, z])
    return R.from_euler('xyz', [x, y, z]).as_quat()


def lookup_position(positions, keys, fallback_msg):
    """
    Return the first available object position amongst a list of possible names.

    Args:
        positions (dict): name → np.ndarray(3,)
        keys (list[str]): Ordered list of keys to try.
        fallback_msg (str): Error message if none are found.

    Returns:
        np.ndarray(3,)
    """
    for k in keys:
        if k in positions:
            return np.asarray(positions[k])
    raise KeyError(fallback_msg)


# ---------------------------------------------------------------------
#  MAIN EXECUTION LOGIC
# ---------------------------------------------------------------------
def run_task_plan():
    """
    Execute the fixed seven-step oracle plan:

        1) rotate  – orient gripper to 90° about world-Z
        2) move    – go to the bottom drawer’s side position
        3) move    – go to the drawer handle (anchor) position
        4) pick    – grasp the drawer handle
        5) pull    – slide the drawer open
        6) pick    – grasp the rubbish on the table
        7) place   – drop the rubbish into the bin
    """
    print("\n=====  START TASK EXECUTION  =====")

    # ------------------------------------------------------------
    # 1. Environment / Task initialisation
    # ------------------------------------------------------------
    env, task = setup_environment()
    try:
        # Reset task and obtain first observation
        descriptions, obs = task.reset()

        # (Optional) start video capture
        init_video_writers(obs)
        task.step = recording_step(task.step)                  # wrap step() for recording
        task.get_observation = recording_get_observation(task.get_observation)

        # --------------------------------------------------------
        # 2. Retrieve object positions
        # --------------------------------------------------------
        positions = get_object_positions()

        # Drawer way-points
        bottom_side_pos = lookup_position(
            positions,
            ['bottom_side_pos', 'bottom_side', 'side_pos_bottom'],
            "[Plan] Could not find bottom drawer side position in object_positions."
        )
        bottom_anchor_pos = lookup_position(
            positions,
            ['bottom_anchor_pos', 'bottom_anchor', 'anchor_pos_bottom'],
            "[Plan] Could not find bottom drawer anchor position in object_positions."
        )

        # Bin
        bin_pos = lookup_position(
            positions,
            ['bin', 'trash_bin', 'rubbish_bin'],
            "[Plan] Could not find bin position in object_positions."
        )

        # Rubbish on the table – try multiple possible identifiers
        rubbish_pos = lookup_position(
            positions,
            ['rubbish', 'item3', 'item1', 'trash', 'garbage'],
            "[Plan] Could not find rubbish position on table in object_positions."
        )

        # --------------------------------------------------------
        # 3. STEP-BY-STEP PLAN
        # --------------------------------------------------------
        done = False
        reward_total = 0.0

        # STEP-1 ROTATE gripper 90 deg about world-Z
        target_quat = quat_from_euler_xyz(0, 0, 90)            # xyzw quaternion
        obs, reward, done = rotate(env, task, target_quat)
        reward_total += reward
        if done:
            print("[Plan] Finished after rotate.")
            return

        # STEP-2 MOVE to the drawer side waypoint
        obs, reward, done = move(env, task, bottom_side_pos)
        reward_total += reward
        if done:
            print("[Plan] Finished after move-to-side.")
            return

        # STEP-3 MOVE to the drawer handle (anchor)
        obs, reward, done = move(env, task, bottom_anchor_pos)
        reward_total += reward
        if done:
            print("[Plan] Finished after move-to-anchor.")
            return

        # STEP-4 PICK the drawer handle (approach from above)
        obs, reward, done = pick(
            env,
            task,
            target_pos=bottom_anchor_pos,
            approach_distance=0.10,
            approach_axis='-z'
        )
        reward_total += reward
        if done:
            print("[Plan] Finished after pick-drawer.")
            return

        # STEP-5 PULL the drawer outward (+X, 15 cm)
        obs, reward, done = pull(
            env,
            task,
            pull_distance=0.15,
            pull_axis='x'
        )
        reward_total += reward
        if done:
            print("[Plan] Finished after pull.")
            return

        # STEP-6 MOVE directly above the rubbish (add 10 cm Z offset)
        above_rubbish = rubbish_pos + np.array([0.0, 0.0, 0.10])
        obs, reward, done = move(env, task, above_rubbish)
        reward_total += reward
        if done:
            print("[Plan] Finished after approach-rubbish.")
            return

        # STEP-7 PICK the rubbish
        obs, reward, done = pick(
            env,
            task,
            target_pos=rubbish_pos,
            approach_distance=0.08,
            approach_axis='-z'
        )
        reward_total += reward
        if done:
            print("[Plan] Finished after pick-rubbish.")
            return

        # STEP-8 MOVE above the bin (15 cm safety)
        above_bin = bin_pos + np.array([0.0, 0.0, 0.15])
        obs, reward, done = move(env, task, above_bin)
        reward_total += reward
        if done:
            print("[Plan] Finished after approach-bin.")
            return

        # STEP-9 PLACE the rubbish into the bin
        obs, reward, done = place(
            env,
            task,
            target_pos=bin_pos,
            approach_distance=0.10,
            approach_axis='-z'
        )
        reward_total += reward

        # --------------------------------------------------------
        # 4. Final status
        # --------------------------------------------------------
        if done:
            print("[Plan] Task completed successfully! Total reward:", reward_total)
        else:
            print("[Plan] Plan executed, but task not yet signalled done. Reward so far:", reward_total)

    finally:
        # Ensure simulator shuts down cleanly
        shutdown_environment(env)
        print("=====  END TASK EXECUTION  =====")


if __name__ == "__main__":
    run_task_plan()