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

# PyRep helpers (fallback Shape look-up when name is absent in get_object_positions)
from pyrep.objects.shape import Shape

# Environment life-cycle utilities
from env import setup_environment, shutdown_environment

# Low-level skill primitives – DO NOT REDEFINE
from skill_code import move, pick, place, rotate, pull

# Optional video utilities
from video import init_video_writers, recording_step, recording_get_observation

# Convenience routine that returns {name : [x, y, z]}
from object_positions import get_object_positions



# ════════════════════════════════════════════════════════════════════════
#                 Helper utilities & robust name mappings
# ════════════════════════════════════════════════════════════════════════
_ALT_NAME_MAP = {
    # ─────────── Drawer positions ───────────
    'side-pos-bottom':   ['bottom_side_pos',   'side_pos_bottom'],
    'anchor-pos-bottom': ['bottom_anchor_pos', 'anchor_pos_bottom'],
    'bottom':            ['bottom_joint_pos',  'bottom_anchor_pos'],
    # ─────────── Tomatoes & plate ───────────
    'tomato1': ['item1'],
    'tomato2': ['item2'],
    'plate':   ['plate', 'plate_pos']         # either name may appear
}


def quat_from_euler(roll=0.0, pitch=0.0, yaw=0.0, degrees=True):
    """Return an xyzw-ordered quaternion from Euler angles."""
    order = 'xyz'
    if degrees:
        return R.from_euler(order, [roll, pitch, yaw], degrees=True).as_quat()
    return R.from_euler(order, [roll, pitch, yaw], degrees=False).as_quat()


def _possible_keys(target_name):
    """Generate every plausible key that may refer to the logical object."""
    return [target_name] + _ALT_NAME_MAP.get(target_name, [])


def safe_get_pos(logical_name, positions_dict):
    """
    Robustly obtain a world-frame XYZ position for a named object.

    1) look up every plausible key present in `positions_dict`
    2) fall back on direct PyRep Shape look-up
    """
    for key in _possible_keys(logical_name):
        if key in positions_dict:
            return np.asarray(positions_dict[key], dtype=float)

    # Fallback: query the simulator directly
    for key in _possible_keys(logical_name):
        try:
            shape = Shape(key)
            return np.asarray(shape.get_position(), dtype=float)
        except Exception:
            continue

    raise RuntimeError(f"[ERROR] Cannot locate position for: '{logical_name}'")



# ════════════════════════════════════════════════════════════════════════
#                   Main routine that executes the plan
# ════════════════════════════════════════════════════════════════════════
def run_skeleton_task():
    print("===== Starting Skeleton Task (Tomato-Drawer) =====")

    # -------------------------------------------------------------------
    # 1) Launch simulator & task
    # -------------------------------------------------------------------
    env, task = setup_environment()
    try:
        _, obs = task.reset()

        # Optional: enable video recording
        init_video_writers(obs)
        task.step           = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # Convenience lambda
        refresh_positions = lambda: get_object_positions()

        # ────────────────────────────────────────────────────────────────
        # Drawer opening sequence (steps 1-5 from specification)
        # ────────────────────────────────────────────────────────────────
        positions = refresh_positions()

        target_quat        = quat_from_euler(0, 0, 90)                # ninety_deg
        side_pos           = safe_get_pos('side-pos-bottom',   positions)
        anchor_pos         = safe_get_pos('anchor-pos-bottom', positions)
        bottom_handle_pos  = safe_get_pos('bottom',            positions)

        # The following block is FROZEN – DO NOT EDIT
        # ---------------------------------------------------------------
        # [Frozen Code Start]
        obs, reward, done = rotate(env, task, target_quat)
        obs, reward, done = move(env, task, side_pos)
        obs, reward, done = move(env, task, anchor_pos)
        obs, reward, done = pick(
                            env, task,
                            target_pos=bottom_handle_pos,
                            approach_axis='z'      # top-down grasp
                        )
        obs, reward, done = pull(
                            env, task,
                            pull_distance=0.20,
                            pull_axis='x'
                        )
        # [Frozen Code End]
        # ---------------------------------------------------------------
        if done:
            print("[Task] Terminated during drawer-opening sequence.")
            return

        # ────────────────────────────────────────────────────────────────
        # Tomato #1  (steps 6-7)
        # ────────────────────────────────────────────────────────────────
        positions    = refresh_positions()
        tomato1_pos  = safe_get_pos('tomato1', positions)
        plate_pos    = safe_get_pos('plate',   positions)

        # Move near tomato1
        obs, reward, done = move(env, task, tomato1_pos)
        if done:
            print("[Task] Terminated approaching tomato1."); return

        # Pick tomato1
        obs, reward, done = pick(
            env, task,
            target_pos=tomato1_pos,
            approach_axis='z'
        )
        if done:
            print("[Task] Terminated after picking tomato1."); return

        # Move to plate
        positions   = refresh_positions()      # plate may have moved
        plate_pos   = safe_get_pos('plate', positions)
        obs, reward, done = move(env, task, plate_pos)
        if done:
            print("[Task] Terminated en-route to plate (tomato1)."); return

        # Place tomato1
        obs, reward, done = place(
            env, task,
            target_pos=plate_pos,
            approach_axis='z'
        )
        if done:
            print("[Task] Terminated after placing tomato1."); return

        # ────────────────────────────────────────────────────────────────
        # Tomato #2  (steps 8-9)
        # ────────────────────────────────────────────────────────────────
        positions    = refresh_positions()
        tomato2_pos  = safe_get_pos('tomato2', positions)

        # Move near tomato2
        obs, reward, done = move(env, task, tomato2_pos)
        if done:
            print("[Task] Terminated approaching tomato2."); return

        # Pick tomato2
        obs, reward, done = pick(
            env, task,
            target_pos=tomato2_pos,
            approach_axis='z'
        )
        if done:
            print("[Task] Terminated after picking tomato2."); return

        # Move to plate
        positions   = refresh_positions()
        plate_pos   = safe_get_pos('plate', positions)
        obs, reward, done = move(env, task, plate_pos)
        if done:
            print("[Task] Terminated en-route to plate (tomato2)."); return

        # Place tomato2
        obs, reward, done = place(
            env, task,
            target_pos=plate_pos,
            approach_axis='z'
        )

        # ----------------------------------------------------------------
        #                           Final status
        # ----------------------------------------------------------------
        if done:
            print("[Task] Goal achieved – both tomatoes placed! Reward:", reward)
        else:
            print("[Task] Plan executed (done=False). Verify success conditions.")

    except Exception as e:
        # Catch any unexpected issues so that we can gracefully shut down
        print("[Task] Exception occurred:", str(e))

    finally:
        shutdown_environment(env)
        print("===== End of Skeleton Task =====")



if __name__ == "__main__":
    run_skeleton_task()